A System for Morphology-Task Generalization via Unified Representation and Behavior Distillation
Authors: Hiroki Furuta
Yusuke Iwasawa
Yutaka Matsuo
Shixiang Shane Gu
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Morphology-Task Generalization
Behavior Distillation
Supervised RL
Reinforcement Learning

Abstract: The rise of generalist large-scale models in natural language and vision has made us expect that a massive data-driven approach could achieve broader generalization in other domains such as continuous control. In this work, we explore a method for learning a single policy that manipulates various forms of agents to solve various tasks by distilling a large amount of proficient behavioral data. In order to align input-output (IO) interface among multiple tasks and diverse agent morphologies while preserving essential 3D geometric relations, we introduce morphology-task graph, which treats observations, actions and goals/task in a unified graph representation. We also develop MxT-Bench for fast large-scale behavior generation, which supports procedural generation of diverse morphology-task combinations with a minimal blueprint and hardware-accelerated simulator. Through efficient representation and architecture selection on MxT-Bench, we find out that a morphology-task graph representation coupled with Transformer architecture improves the multi-task performances compared to other baselines including recent discrete tokenization, and provides better prior knowledge for zero-shot transfer or sample efficiency in downstream multi-task imitation learning. Our work suggests large diverse offline datasets, unified IO representation, and policy representation and architecture selection through supervised learning form a promising approach for studying and advancing morphology-task generalization.

Published as a conference paper at ICLR 2023
A SYSTEM FOR MORPHOLOGY-TASK GENERALIZATION VIA UNIFIED REPRESENTATION AND BEHAVIOR DIS- TILLATION
Hiroki 1:2 Yusuke Iwasawa' Yutaka Matsuo
Shixiang Shane 2u
'The University of Tokyo "Google Research, Brain Team furuta@weblab, u-tokyo..oo ac. jp
ABSTRACT
The rise of generalist snge-scale models in natural language and vision has made expect that a massive data-driven approach could achieve broader generalization in other domains such as continuous control. this work, we explore method for learning single policy that manipulates various forms of agents to solve various tasks by distilling a large amount of proficient behavioral data. In order to align input-output (IO) interface among multiple tasks and diverse agent morphologies while preserving essential 3D geometric relations, we introduce morphology- graph, which treats observations, actions and goals/task in unified graph represen- tation. We also develop MxT-Bench for fast large-scale behavior generation, which supports procedural generation of diverse morphology-task combinations with a minimal blueprint and gardware-accelerse tation and architecture selection on MxT-Bench, we find out that a morphology-tar graph representation coupled with Transformer architecture improves the multi-task performances compared other baselines including recent discrete tokenization, and provides better prior knowledge for zero-shot transfer or sample efficiency in downstream multi-task imitation learning. Our work suggests large diverse offline datasets, unified IO representation, and policy representation and architecture se- lection through supervised learning form a promising approach for studying and advancing norphology-tase generalization!
simulator. Through efficient represen-
INTRODUCTION
The impressive success of large language models (Devlin et al., 2019; Radford et al., 2019; Bom- masani et al., 2021; Brown et al., 2020; Chowdhery et al., 2022) has encouraged the other domains, such computer vision (Radford et al., 2021; Gu et al., 2021b; Alayrac al., 2022: Jaegle al., 2021) or robotics (Ahn et al., 2022; Huang et al., 2022b), to leverage the large-scale pre-trained model trained with massive data with unified input-output interface. These large-scale pre-trained models are innately multi-task learners: they surprisingly work well not only in the fine-tuning or few-shot transfer but also in the zero-shot transfer settings (Raffel al., 2020; Chen et al., 2022a). Learning "generalist" model seems to be an essential goal in the recent machine learning paradigm with the same key ingredients: curate massive diverse dataset, define unified IO representation. and perform efficient representation and architecture selection, altogether for best generalization.
In reinforcement learning (RL) for continuous control, various aspects are important for generalization. First, we care about "task" generalization. For instance, in robotic manipulation, we care the policy to generalize for different objects and target goal positions (Kalashnikov et al., 2018; Andrychowicz et al., 2017; Yu et al., 2019; Lynch al., 2019). Recent advances in vision and language models also enable task generalization through compositional natural language instructions (Jiang et al., 2019; Shridhar et al., 2022a; Ahn et al., 2022; Cui et al., 2022). However, to scale the data, equally important is "morphology" generalization, where a single policy can control agents of different embodiment (Wang al., 2018; Noguchi et al., 2021) and can thereby ingest experiences from as
Work done as Student Rescarcher at Google.
httpp:: //s / tes. google com/ http//vie///c/ee//o rowee/eetolllleaaeaaae

---

Published as a conference paper at ICLR 2023
MxT-Bench
nn-Distribbttin Generaiization
Single-Task Task RL (PPO)
Morphology- Task Graph (Unified (O) Morph,
Zero-Shot Generalization
Morphology-Tee Task Behavior Dataset
Behavior Distillation
Downstream Fine-Tuning
Rep. Selection
Figure Behavior distillation pipeline. We train a single-task policy for each environment on MxT- Bench, and then collect proficient morphology-tass task behavior dataset (Section 4.1). To enable a single policy to learn multiple tasks and morphologies simultaneously, we convert stored transitions to the morphology- -ask graphrepresentation to align with unified IO interface (Section 4.3) for multi-task distillation (Section 4.2). After behavior distillation, the learned policy can be utilized for in-distribution zero-shot generalization (Section 5) downstream fine-tuning (Section 5.2), and representation and architecture selection (Section 5.3).
many robots in different simulators (Freeman et al., 2021; Todorov et al., 2012; Coumans & Bai, 2016) as possible. Most prior works (Mendonca al., 2021; Gupta et al., 2022) only address either the task or morphology axis separately, and achieving broad generalization over task and morphology jointly remains long-standing problem'.
This paper first proposes MxT-Bench', the first multi-morpholooe and multi-task benchmarking environments, as a step toward building the massive diverse dataset for continuous control. MxT- Bench provides various combinations of different morphologies (ant, centipede, claw, worm, and unimal (Gupta al., 2022)) and different tasks (reach, touch, and twisters). MxT-Bench is easily scalable to additional morphologies and tasks, and is built top of Brax (Freeman et al., 202 for fast behavior generation.
Next, we define unified I0 representation for architecture to ingest all the multi-morphology multi-task data. Inspired by scene graph (Johnson et al., 2015) in computer vision that represents the 3D relational information of scene, and by morphology graph (Wang et al., 2018; Chen et al., 2018; Huang et al., 2020; Gupta et al., 2022) that expresses an agent's 's geometry and actions, we graph (MTG) unified interface to encode observations,
introduce the notion of morphology-ten actions, and goals (i.e. tasks) as nodes in the shared graph representation. Goals are represented as sub-nodes, and different tasks correspond to different choices: touching is controlling a torso node, while reaching is controlling an end-effector node (Figure 3). In contrast to discretizing and tokenizing every dimension as proposed in recent work (Janner et al., 2021; Reed et al., 2022), this unified IO limits data representation it can ingest, but strongly preserves 3D geometric relationships that are crucial for any physics control problem (Wang al., 2018; Ghasemipour et al., 2022), and we empirically show it outperforms naive tokenization in our control- -ocrue- dataset.
Lastly, while conventional multi-task or meta RL studies generalization through on-policy joint training (Yu et al., 2019; Cobbe et al., 2020), we perform efficient representation and architecture selection, over 11 combinations of unified IO representation and network architectures, and local node observations, for optimal generalization through behavior distillation (Figure 1), where RL is essentially treated as a (single-task, (ow-dimensiona) behavior generator (Gu et al., 202 and multi-task supervised learning (or offline RL (Fujimoto et al., 2019)) is used for imitating all the behaviors (Singh et al., 2021; Chen et al., 2021b; Reed et al., 2022). Through offline distillation, we controllably and tractably evaluate two variants of MTG representation, along with multiple network architectures (MLP, GNN (Kipf& & Welling, 2017), Transformers (Vaswani et al., 2017)), and show that MTGv2 variant with Transformer improves the multi-task goal-reaching performances compared to other possible choices by 23% and provides better prior knowledge for zero-shot generalization (by 14%18) and fine-tuning for downstream multi-task imitation learning (by 50 55 %).
As the fields of vision and language move toward broad generalization (Chollet, 2019; Bommasani et al., 2021), we hope our work could encourage RL. and continuous control communities continue
Heer, "task" means what each agent should solve. See Section for the detailed definition. When considering the models, as slightly overloaded, it may imply morphological diversity as well. Pronounced maixed"-benel" It stands for "Morphology x Task".
2

---

Published as a conference paper at ICLR 2023
growing diverse behavior datasets, designing different 10 representations, and iterating more repre- sentation and architecture selection, and eventually optimize a single policy that can be deployed on any morphology for any task. In summary, our key contributions are:
We develop MxT-Bench as a test bed for morphology-tast generalization with fast expert behavior generator. MxT-Bench supports the scalable procedural generation of both agents and tasks with minimal blueprints.
We introduce morphology- -task graph, a universal I0 for control which treats the agent's s observations, actions and goals/tasks in unified graph representation, while preserving the task structure.
We study generalization through offline supervised behavior distillation, where we can effi- ciently try out various design choices; over 11 combinations of unified IO representation and network architectures, and 8 local node observations. As result, we find that Transformer with MTGv2 achieves the best multi-task performances among other possible designs (MIP, GNN and Transformer with MTGv1, Tokenized-MTer etc.) in both in in-distribution and downstream tasks, such as zero-shot transfer and fine-tuning for multi-task imitation learning.
2 RELATED WORK
Morphology Generalization While, in RL for continuous control, the policy typically learns to control only a single morphology (Tassa et al., 2018; Todorov et al., 2012), several works succeed in generalizing the control problem for morphologically different agents to solve a locomotion task by using morphology-avee awarre Graph Neural Network (GNN) policies (Wang et al., 2018; Huang et al., 2020; Blake et al., 2021). In addition, several work (Kurin al., 2021: Gupta et al., 2022; Hong et al., 2022; Trabucco et al., 2022) have investigated the use of Transformer (Vaswani et al., 2017). Other work jointly optimize the morphology- agnostic policy and morphology itself (Pathak al., 2019; Gupta et al., 2021; Yuan et al., 2022; Hejna et al., 2021), or transfer a controller over different morphologies (Devin et al., 2017; Chen et al., 2018; Hejna et al., 2020; Liu et al., 2022).
Morphology Genernitzation
Morphology-Task
Generalization
Task Generalization
Figure 2: We tackle morphology-task gen- eralization, which requires achieving both task and morphology generalization simulta- neously. See Appendix M for the details.
While substantial efforts have been investigated to realize morphology generalization, those works mainly focus on only a single task (e.g. running), and less attention is paid to multi-task settings, where the agents attempt to control different states of their bodies to desired goals. We believe that goal-directed control is key problem for embodied single controller. Concurrently, Feng et al. (2022) propose an RL--hased single controller that is applied to different quadruped robots and target poses in the sim-to-real setting. In contrast, our work introduces the notion of morphology graph as a unified 10 that represents observations, actions, and goals in a shared graph, and can handle more diverse morphologies solve multiple tasks.
Task Generalization In the previous research, task generalization has been explored in multi-task or meta RL literature (Wang et al., 2016; Duan al., 2017: Cabi et al., 2017: Teh et al., 2017: Colas et al., 2019; Li et al., 2020; Yang et al., 2020; Kurin et al., 2022). Each task might be defined by the difference in goals, reward functions, and dynamics (Ghasemipour et al., 2021; Eysenbach et al., 2020), under shared state and action spaces. Some works leverage graph representation to embed the compositionality of manipulation tasks (Li et al., 2019; Zhou et al., 2022: Li et al., 2021; Kumar et al., 2022; Ghasemipour et al., 2022), while others use natural language representation to specify diverse tasks (Jiang et al., 2019; Shridhar et al., 2022a; Ahn et al., 2022; Huang et al., 2022a; Cui et al., 2022). Despite the notable success of acquiring various task generalization, multi-task RL often deals with only a single morphology. We aim to extend the general behavior policy into the "cartesian product" of tasks and morphologies (as shown in Figure 2) to realize a more scalable and capable controller.
al., 2019; Kalashnikov
hhttps://github.ccom
bench
3

---

Published as a conference paper at ICLR 2023
Morphology-Tan Graph
vel of COM
TUn
position
reach
position
encodo
touch
Agent
Dsy, Xy, 2)
twisters
v1
v2
Figure 3: We propose the notion of morphology-task graph, which expresses the agent's 's observations, actions, and goals/tasks in unified graph representation, while preserving the geometric structure of the task. We develop two practical implementations; morphology-tan graph (left) accepts the morphological graph. encoded from the agent's geometric information, as an input-output interface, and merges positional goal information as a part of corresponding node features. morphology-tass graph v2 (right) treats given goals as extra disjoint RI. have focused on locomotion (run
nodes of the morphological graph. While most prior morphology-agnosee task), i.c. a single static goal node controlling (maximizing) the velocity of center of mass, morphology-tase graph could naturally extend morphology-agnoste
control to other goal-oriented tasks: single static goal node for reach, multiple static goal nodes for twisters, and single dynamic goal node tracking a movable ball for an object interaction task; touch.
Transformer for RL Recently, Chen al. (2021a) and Janner al. (2021) consider offline RL as supervised sequential modeling problem and following works achieve impressive success (Reed et al., 2022; Lee al., 2022; Furuta et al., 2022; Xu al., 2022; Shafiullah et al., 2022; Zheng et al., 2022; Paster et al., 2022). In contrast, our work leverages Transformer to handle topological and geometric information of the scene, rather than a sequential nature of the agent trajectory.
Behavior Distillation Due to the massive try-and-error and large variance, training policy from scratch in online RL an inefficient process, especially in multi- setting. It more efficient use RL for generating single-tast task behaviors (often from low dimensions) (Gu et al., 2021a) and then use supervised learning to imitate all behaviors with a large single policy (Levine et al., 2016; Rusu et al., 2016; Parisotto et al., 2016; Singh et al., 2021; Ajay et al., 2021; Chen et al., 2021b). Several works tackle the large-scale behavior distillation with Transformer (Reed et al., 2022; Lee et al., 2022), or with representation that treats observations and actions in the same dision-languas et al., 2020; Shridhar et al., 2022b). Our work utilizes similar pipeline, but focuses on finding the good representation and architecture to generalize across morphology and tasks simultaneously with proposed morphology- graph. See Appendix M for the connection to policy distillation.
space (Zeng
3 PRELIMINARIES
In RL, consider a Markov Decision Process with following tuple (S, A. p. Pı, r, ?), which consists
of state space S, action space A, state transition probability function p A S S [0, oo), initial state distribution
S + |0, oo), reward function r x A + R, and discount factor 7 E [0, 1). The agent follows Markovian policy æ A [0, oo), which is often parameterized in Deep RL, and seeks optimal policy π² that maximizes the discounted cumulative rewards:
n² arg max
1 Es (s).amm()1s) [r(s, a)l,
(1)
1-y
where P/ (st) = [f f. -p--1 II. p( a. 4-114-17 4 (a)()) and (s) = -1- E1
represent time-aligned and time-aggregated state marginal distributions following policy m.
Graph Representation for Morphology-Agnostee Control Following prior continuous control literature (Todorov et al,, 2012), we assume the agents have bodies modeled as simplified skeletons of animals. An agent's morphology is characterized by the parameter for rigid body module (torso, limbs), such radius, length, mass, and inertia, and by the connection among those modules (joints). In order to handle such geometric and topological information, an agent's 's morphology can be expressed as acyclic tree graph representation G == (V, E), where V is set of nodes o V and E is a set of edges eM e E between v' and v., The node e' corresponds to i-th module of the agent, and the edge el corresponds the hinge joint between the nodes s² and v². Each joint may
4

---

Published as a conference paper at ICLR 2023
Bluoprint
Missing
Stze/Mass
Task
Procedural Generation
Benchmark
Muis Tack
Mulii Marphlyyyy
Scalability
Centipede
MuloCo Connol MossWaeldl (Yo 2019)
X
Houng al (2020)
Clow
Gupto et (3021)
MrT-leench(Onns))
Custam
Agents
Figure 4: The overview of MxT-Bench, which can procedurally generate both various morphologies and tasks with minimal blueprints. MxT-Bench can not only construct the agents with different number of limbs, but also randomize missing limbs and size/mass of bodies. We could design the tasks with parameterized goal distributions. It also supports to import custom complex agents such as unimals (Gupta et al., 2022). Compared to relevant RL benchmarks in terms of (1) multi-task (task coverage).
(2) multi- morphology (morphology coverage), and (3) scalability, MuJeCo (Todorov et al., 2012) and DM Control (Tassa al., 2018) only have single morphology for single task. While other existing works (Yu et al., 2019; Huang et al., 2020) partially cover maakkmmorpooloy-axxx with some sort scalability, they do not satisfy all criteria.
have 1-3 actuators corresponding degree of freedom. If a joint has several actuators, the graph G is considered multipath one. This graph-based factored formulation can describe the various agents' morphologies in a tractable manner (Wang et al., 2018; Huang et al., 2020; Loynd et al., 2020).
We assume that node v² observes local sensory input st at time step f, which includes the information of limb i such as position, velocity, orientation, joint angle, or morphological parameters. To process these node features and graph structure, a morphology-agnosee
policy can be modeled as node-based GNN (Kipf& Welling, 2017; Battaglia et al., 2018; Cappart et al., 2021), which takes a set of local observations {st] V as an input and emit the actions for actuators of each joint {af] The objective of corphology-agnostie RL is the average of Equation among given morphologies.
Goal-conditional RL In goal-conditional RL (Kaelbling, 1993; Schaul et al., 2015), the agent aims to find an optimal policy x" (als, 8g) conditioned on goal 8g E G, where G stands for goal space that is sub-dimension of state space S (e.g. XYZ coordinates, velocity, or quaternion). The desired goal Sy is sampled from the given goal distribution Py G + [0, oo), where ei stands for task in the task space y (e.g. reaching the agent's is leg touching the torso to the ball). The reward function can include goal-reaching term, that is often modeled as (3,, sy) = (s,, sy), where dæ (:,-) .) is a dask-dependent distance function, such as Euclidean distance, between the sub-dimension of interest in current state 8, and given goal 8.- Some task t give multiple goals to the agents. In that case, we No overload 8g to represent set of goals; {st) (=1, where No. is the number of goals that should be satisfied in the task vb.
Morphology-Tass Generalization This paper aims to achieve morphology-task generalization, where the learned policy should generalize over tasks and morphologies simultaneously. The optimal policy should generalize over morphology space M, task V. and minimize the distance to any given goal 8g E G. Mathematically, this objective can be formulated as follows:
π = arg max M. (Ea
[Egm "m (0""""(""))))())) ..)) )) I (s
.8 [gg]]]]
(2)
E M is denoted as G. (Vm. Em), and sth : [ v.l and a := {aā {a])_err |Em.l stand for a set of local observations and actions of morphology m. While we can use multi-task online RL to maximize Equation 2 in principle, it is often sample inefficient due to the complexity of task, which requires a policy that can handle the diversity of the scene among morphology M, task V, and goal space G simultaneously.
where the graph representation of morphology
4 METHOD
4.1 MXT-BENCH AS TEST BED FOR MORPHOL.OGY-TASK GENERALIZZTTIN
To overcome these shortcomings in the existing RL environments, we develop MxT-Bench, which has wide coverage over both tasks and morphologies to test morphology-task generalization, with the
5

---

Published as a conference paper at ICLR 2023
functionalities for procedural generation from minimal blueprints (Figure 4). MxT-Bench is built on top of Brax (Freeman et al., 2021) and Composer (Gu et al., 2021a), for faster iteration of behavior environments. Beyond supporting
distillation with multi-task settings, the scalability of MxT-Bench helps to test the broader range of morphology- generalization since can easily generate autoo--distrbbution tasks and morphologies, compared to manually-designed morphology or task specifications.
and
In the morphology axis, we prepare 4 types of blueprints (ant, claw, centipede, and worm) as base morphologies, since they are good at the movement on the XY-plane. Through MxT-Bench, we can easily spawn agents that have different numbers of bodies, legs, or different sizes, lengths, and weights. Moreover, we can also import the existing complex morphology used in previous work. For instance, we include 60+ morphologies that are suitable for goal-reaching, adapted from Gupta et al. (2022) designed in MuloCo. In the task axis, we design reach, touch, and twisters" as basic tasks, which could evaluate different aspects of the agents; the simplest one is the reach task, where the agents aim to put their leg on the XY goal position. In the touch task, agents aim to create and maintain contact between a specified torso and a movable ball. The touch task requires reaching behavior, while maintaining conservative momentum to avoid kicking the ball away from the agent. Twisters tasks are the multi-goal problems; for instance, the agents should satisfy XY XY-position for one leg, and Z. height for another leg. We pre-define variants of twisters with max goals (see Appendix B.2 for the details). Furthermore, we could easily specify both initial and goal position distribution with parameterized distribution. In total, we prepare 180+ environments combining the morphology and task axis for the experiments in the later section. See Appendix B for further details.
4.2 BEHAVIOR DISTILLATION
Toward broader generalization over morphologies and tasks, a single policy should handle the diversity of the scene among morphology M, task V, and goal space G simultancously. Multi- task online RL from scratch, however, is difficult to tune, slow to iterate, and hard to reproduce. Instead, we employ behavior cloning on RL--genrraeed expert behaviors to study morphology-tan generalization. To obtain rich goal-reaching behaviors, we train single-morphology single-task policy using PPO (Schulman et al., 2017) with a simple MLP policy, which significantly more efficient than morii-morphology training done in prior work (Gupta et al., 2022; Kurin et al., 2021; Huang et al., 2020). Since MxT-Bench is built on top of Brax (Freeman et al., 2021), a hardware- accelerated simulator, training PPO policies can be completed in about 5~~30 minutes per environment (on NVIDIA RTX A6000). We collect many behaviors per morphology-tan combination from the expert policy rollout. We then train single policy To with supervised learning objective:
L M., E a ,8, - DDn. [log (a (s"" [5y)]]]
where D a is an expert dataset of morphology and task sh. Importantly, offline behavior distillation protocol runs (parallelizable) single-task online RL only once, and allows us to reuse the same fixed data try out various design choices, such as model architectures or local features of morphology-tass graph, which is often intractable in multi-task online RL.
4.3 MORPHOLOGY-TASS TASK GRAPH
To learn a single policy that could solve various morphology-task problems, it is essential to unify the input-output interface among those. Inspired by the concept of scene graph (Johnson et al., 2015) and morphological graph (Wang et al., 2018). we introduce the notion of morphology-tass graph (MTG) representation that incorporates goal information while preserving the geometric structure of the task. Morphology-tan in unified graph space. Although most prior (running) with the reward calculated from the velocity of the center of mass, morphology-task graph RL to multi-task goal-oriented settings: including static
graph could express the agent's observations, actions, and goals/tasks
happhology-ggnostie RL has focused on locomotion
could naturally extend morphology-agnose single positional goals (reaching), multiple-ge -ooll problems (twister-ge -aaeee and object interaction tasks (ball-touching). In practice, we develop two different ways to inform the goals and tasks (Figure 3); morphology-tase graph v1 (MTGv1) accepts the morphological graph, encoded from the agent's geometric information, as an onput-output interface, and merges positional goal information as a
hhttps://en http:///e..wwiiiiiee.e www.weiii.aaaaooo............. /wwk/wi/ wiki http://wiiitrreeeeeeeeeeeee

---

Published as a conference paper at ICLR 2023
Random
MLP
GNN (MTG+11
Transformer (MTGG11)
Transformer (MTGG22
Token-MTGr2
In-Detributtio
1,2019 0.41
0,5150 0,01
0.4776 0.01
0.1069 1 0.02
0.3128
0,02
0.3402 1 0.01
In-Ditcibution (unimal)
0,4839 0.02
0,4178
Compusitional (Murphulogy)
1,1419
044
0.7216 0.01
0.4940 1 0.01
0,4066
Compositional (Task)
0.8932
0.0
0.649 = 0.01
0. 0.5395 + 0.04
0.4461 0.05
Outtoff-istrbattion
0.8979 0.01
0.7821 0.02
0.6144 + 0:04
0.5266 0.04
The average normalized final distance in various types of morphology -tsk generalization on MxT- Bench. We test (1) in-distribution, (2) compositional morphology/task, and (3) (2) and (3) evaluate zero-shot transfer. We compare MLP, GNN and Transformer with MTGv1, Transformer with MTGv2, and tokenized MTGv2 (Reed et al., 2022). MTGv2 improves multi-task performance other choices by 23 % in the in in-distribution evaluation, and achieves better zero-shot transfer in the compositional and out-of-distribution settings by 14 ~v 18 %.
Table I:
generalization.
part of corresponding node features. For instance, in touch task, MTGv1 includes XY position of the movable ball as an extra node feature of the body node. Moreover, morphology traphllogy-ask graph v2 (MTGv2) considers given goals additional disjoint nodes of morphological graph representation. These morphology-tan graph strategies enable the policy to handle a lot of combinations of tasks and morphologies simultancously.
Transformer for MTG While morphology-task graph could represent the agent and goals in a unified manner, because the task structure may change over time, the policy should unravel the implicit relationship between the agent's 's modules and goals dynamically. We mainly employ Transformer as a policy architecture because it can process morphological graph as a fully-connected graph and achieve notable performance by leveraging the hidden relations between the nodes beyond manually-encoded geometric structure (Kurin et al., 2021). The combination of morphology-Een is expected to model the local relations between goals and target nodes explicitly and to embed the global relations among the modules. See Appendix for further details.
graph and Transformer
EXPERIMENTS
We first evaluate the multi-task performance of morphology-task graph representation (MTGv1, MTGv2) in terms of in-distribution (known morphology and task with different initialization), compositional (known task with unseen morphology or known morphology with unseen task), and outto--distribution generalization (either morphology or task is unseen) on MxT-Bench (Section 5.1). Then, we investigate whether morphology-task graph could contribute obtaining better control prior for multi-task fine-tuning (Section 5.2). In addition, we conduct the offline node feature selection to identify what is the most suitable node feature set for morphology-tee The results are averaged among 4 random seeds. See Appendix A for the hyperparameters We also investigate the other axis of representations or architectures (Appendix E, F, G) and test the effect of dataset size, the number of norphology-tase combinations, and model size (Appendix H). Lastly, we examine why morphology-tas graph works well by visualizing attention weights (Appendix K).
generalization (Section 5.3).
Evaluation Metric Goal-reaching tasks are evaluated by the distance to the goals at the end of episode (Pong et al., 2018; 8; 2020; Ghosh et al., 2021; Choi et al., 2021; 1; Eysenbach et al., 2021) I). However, this can be problematic in our settings, because the initial distance or the degree of goal- reaching behaviors might be different among various morphologies and tasks. We measure the performance of the policy by using normalized final distance metric d(M, V; T over morphology M and task space T with pre-defined max/min value of each morphology m and task vb.
d(M V; = Manlill EM ES Exp
1
de (x77) (
m.m.
min
mpe
af
where 8ye is the last state of the episode, di is maximum, and d.... is minimum distance of i-th goal sy with morphology m and task v.. We use a distance threshold to train the expert PPO policy as gi.m.&, and the average distance from an initial position of the scene as d Equation normalized around the range of [0. 1], and the smaller, the better. See Appendix B for the details.
5.1 BEHAVIOR DISTILLATION ON MXT-BENCE
We systematically evaluate three types of morphology- generalization; in in-distribution, composi- tional, and out-of-distiibution generalization through behavior distillation. As baselines, we compare
7

---

Published as a conference paper at ICLR 2023
Compositionl Morphology)
Composiional ("aok)
Sw--ot-Ditributtr
MSP
MLP (220
THMTOV (40
TFWIII (00
TTAmTo (120
TTATTO (126)
00
0.0
06
Avarage Norma Distance
Avaraga Normatzed Distance
Avarage Normakeed Distance
Figure 5: Multi-task goal-reaching performances on fine-tuning (multi-task imitation) for compositional and out-of-distribution evaluation. These results revcal that fine-tuning outperforms random initialization in all settings, and fine-tuned MTGv2 outperforms others by 50 ~ 55 %. See Appendix for the detailed scores.
MLP. GNN (Wang et al.. 2018) with MTGv1, Transformer with MTGv1 or MTGv2, and tokenized MTGv2, similar Reed et al. (2022) (see Appendix E for the details).
In in-distribution settings, we prepare 50 environments and 60 unimal environments adapted from Gupta et al. (2022) (see Appendix B.3 for the details) for both training and evaluation. Compositional and
settings evaluate zero-shot transfer. In compositional settings, we test morphol- ogy and task generalization separately; we prepare 38 train environments and 12 test environments with hold-out morphologies for morphology evaluation, and leverage 50 train environments and prepare 9 test environments with unseen task for task evaluation. In out-of-distribution settings, also leverage 50 environments as training set, and define 27 environments with diversified morphologies and unseen task as an evaluation set. The proficient behavioral data contains 12k transitions per each environment. See Appendix D for the details of environment division.
Table 1 reveals that MTGv2 achieves the best multi task goal-reaching performances among other possible combinations in all the aspects of generalization. Comparing average normalized distance, MTGv2 improves the multi-task performance against the second best, MTGv1, by 23% in in- distribution evaluation. Following previous works (Kurin et al., 2021; Gupta et al., 2022), Transformer with MTGv1 I achieves better goal-reaching behaviors than GNN. In compositional and out-of- distribution zero-shot evaluation, MTGv2 outperforms other choices by 14 ~ 18%. Moreover, the compositional zero-shot performance of MTGv2 is comparable with the performance of MTGv1 I in in-distribution settings. These results imply MTGv2 might be the better formulation to realize the morphology-tee -ask generalization.
5.2 DOES MORPHOLOGY TASK GRAPH OBTAIN BETTER PRIOR FOR CONTROL?
To reveal whether the distilled policy obtains reusable inductive bias for unseen morphology or task, we test the fine-tuning performance for multi-task imitation learning on MxT-BBnch. We adopt the same morphology-tast division for compositional and outtof-distribution evaluation in Section 5.1. Figure 5 shows that fine-tuning outperforms random initialization in all settings, which suggests that behavior-distilled policy works as a better prior knowledge for control. The same as zero-shot transfer results in Section 5.1, MTGv2 outperforms other baselines, and is better than the second best, MTGv1 1 by 50 55%. Furthermore, MTGv2 could work even with a small amount of dataset (4k, 8k); for instance, in compositional morphology evaluation (left in Figure 5), MTGv2 trained with 8k transitions still outperforms competitive combinations with 12k transitions, which indicates better sample efficiency (see Appendix J for the detailed scores). These results suggest MTGv2 significantly captures the structure of morphology graph as prior knowledge for downstream tasks.
5.3 WHAT IS THE BEST NODE FEATURES FOR MORPHOLOGY- TASK GRAPH?
In the agent's 's system, there are a lot of observable variables per module: such as Cartesian position (p). Cartesian velocity (v), quaternion (q), angular velocity (a), joint angle joint range (jr), limb id (id). joint velocity (jv), relative position (rp), relative rotation (rr), and morphological information (m). Morphological information contains module's ' shape, mass, inertia, actuator's gear, dof-index, etc. To shed light on the importance of the local representation for generalization, we execute an extensive ablation of node feature selections and prepossessing. In prior works, Huang et al. (2020) and others (Kurin et al., 2021; Hong et al., 2022) used [p. v. q. a, ja, jr. id) and Gupta et al. (2022) used (p, v, q, a, ja. jr, jv. rp. rr. m) Considering the intersection of those, we define (p. v. 4, a, ja, jr) base, base_settt_ and test the combination to other observations (jv, id, rp, rT, m).
8

---

Published as a conference paper at ICLR 2023
In the experiments, we evaluate in-distribution generalization, done in Section 5.1, with MTGv2 representation. Table shows that, while some additional features (id. rp, m) contribute to improving the performance, the most effective feature can be morphological information (m)) which suggests base_set contains sufficient features for control, and raw morphological properties (m) serves better task specification than manually-encoded information such as limb id (id). The key observations are; (1) morphological information is critical for morphology-tat generalization, and (2) extra observation might disrupt the performance, such as relative rotation between parent and child node (rr). Throughout the paper, we treat base_ set t- as node features for morphology-tan -task graph.
Node features
+jv
+pp
+m
Average
base_set
0,4330 0.02
bese_set-id d
0.4090 0.02
base_set-rp
0.3820 0.01
bese_set-Er
0.4543 0.01
base_aet-m
0.3128 + 0.02
base_aet xp xx
0.3869 + 0.01
base_net-vv-rp-rr basn_eet--v--p--rre
0.4000 1 0.01
0,3323 + 0.01
Table 2: Offline node feature selection. We compare the combination of Cartesian position (p): Cartesian velocity (v), quaternion (q), angular velocity (a). joint angle (ja). joint range (jr). limb id joint velocity (jv). relative position (rp), relative rotation (rr). and morphological information (m). base_ set is composed of (p. v. q. a. ja. jr). These results suggest base_ set contains sufficient features for control, and the most effective feature seems morphological information (m) for better task specification.
6 DISCUSSION AND LIMITATION
While the experimental evaluation on MxT-Bench implies that morphology-tase graph is a simple and effective method to distill the diverse proficient behavioral data into generalizable single policy, there are some limitations. For instance, we focus on distillation from expert policies only, but it is still unclear whether morphology-tas graph works with moderate random quality behaviors in offline RL (Fujimoto et al., 2019; Levine et al., 2020; Fujimoto & Gu, 2021). Combining distillation with iterative data collection (Ghosh al., 2021; Matsushima et al., 2021) or online fine-tuning (Li et al., 2022) would be promising future work. In addition, we avoided tasks where expert behaviors cannot be generated easily by single-task RL without fine-scale reward engineering or human demonstrations; incorporating such datasets or bootstrapping single-task RL from the distilled policy could be critical for scaling the pipeline to more complex tasks such as open-ended and dexterous manipulation (Lynch et al., 2019; Ghasemipour et al., 2022; Chen et al., 2022b). Since morphology-tas graph only uses readily accessible features in any simulator and could be automatically defined through URDFs or MuJoCo XMLs. in future work we aim to keep training our best morphology-task graph architecture policy on additional data from more functional and realistic control behaviors from other simulators like MuJoCo (Yu et al., 2019; Tassa et al., 2018) 8), PyBullet (Shridhar et al., 2022a), IsaacGym (Chen et al., 2022b; Peng et al., 2021), and Unity (Juliani et al., 2018), and show it achieves better scaling laws than other representations (Reed et al., 2022) on broader morphology-tas families.
7 CONCLUSION
The broader range of behavior generalization is a promising paradigm for RL. To achieve morphology- task generalization, we propose morphology-tast graph, which expresses the agent's modular ob- servations, actions, and goals as a unified graph representation while preserving the geometric task structure. As test bed for morphology-tast generalization, we also develop MxT-Bench, which en- ables the scalable procedural generation of agents and tasks with minimal blueprints. Fast-generated behavior datasets of MxT-Bench with RL allow efficient representation and architecture selection through supervised learning, and MTGv2, variant of morphology- multi-task performances among other possible designs (MLP, GNN and Transformer with MTGv1, and tokenized- Bieenized-MTG22. etc), outperforming them in in-distribution evaluation (by 23 %), zero-shot transfer among compositional
graph, achieves the best
diistri-titiiiooinn evaluation (by 14 ~ 18 %) and fine-tuning for downstream multi-tase imitation (by 50 ~ 55 %). We hope our work will encourage the community to explore scalable yet incremental approaches for building universal controller.
9

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was supported by JSPS KAKENHI Grant Number JP22J21582. We thank Mitsuhiko Nakamoto and Daniel C. Freeman for converting several MuJoCo agents to Brax, So Kuroki for the support on implementations, Yujin Tang, Kamyar Ghasemipour, Yingtao Tian, and Bert Chan for helpful feedback on this work.
REFERENCES
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Daniel Ho, Jasmine Hsu, Julian Ibarz, Brian lchter, Alex Irpan, Eric Jang, Rosario Jauregui Ruano, Kyle Jeffrey, Sally Jesmonth, Nikhil J Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Kuang- Lee, Sergey Levine, Yao Lu, Linda Luu, Carolina Parada, Peter Pastor, Jornell Quiambao, Kanishka Rao, Jarek Rettinghouse, Diego Reyes, Pierre Sermanet, Nicolas Sievers, Clayton Tan, Alexander Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Mengyuan Yan, and Andy Zeng. Do as can, not say: Grounding language in robotic affordances. arXiv preprint arxiv:2204.06691, 2022.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. OPAL: Offline primitive discovery for accelerating offline reinforcement learning. In International Conference on Learning Representations, 2021.
Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lene, Arthur Mensch, Katie Millican, Malcolm Reynolds, Roman Ring, Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob Menick, Sebastian Borgeaud, Andrew Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, and Karen Simonyan. Flamingo: visual language model for few-shot learning. arXiv preprint arxiv:2204.144 2022.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. In Advances in neural information processing systems, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv: 1607.06650 2016.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinieius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 02000.0/266 01261, 2018.
Aditya Bhatt, Max Argus, Artemij Amiranashvili, and Thomas Brox. Crossnorm: Normalization for off-policy td reinforcement learning. arXiv preprint /00Xvv:9922.0603 2019.
Charlie Blake, Vitaly Kurin, Maximilian Igl, and Shimon Whiteson. Snowflake: Scaling gnns to sich-dimensional continuous control via parameter freezing. arXiv preprint arXiv:2103.01009 2021.
Rishi Bommasani. Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg. Antoine Bossclut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri Chatterji, Annie Chen, Kathleen Creel, Jared Quincy Davis, Dora Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto. Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong. Kyle Hsu, Jing Huang, Thomas Ieard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Kohd, Mark Krass, Ranjay Krishna, Rohith Kuditipudi, Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovec, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu
10

---

Published as a conference paper at ICLR 2023
Ma, Ali Malik, Christopher D. Manning, Suvir Mirchandani, Eric Mitchell, Zancle Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Ben Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, Julian Nyarko, Giray Ogut, Laurel Orr, Isabel Papadimitrion, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Rob Reich, Hongyu Ren, Frieda Rong, Yusuf Roohani, Camilo Ruiz, Jack Ryan, Christopher Ré, Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishnan Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramèr, Rose E. Wang, William Wang. Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie, Michihiro Yasunaga, Jiaxuan You, Matei Zaharia, Michael Zhang, Tianyi Zhang, Xikun Zhang, Yuhui Zhang, Lucia Zheng, Kaitlyn Zhou, and Perey Liang. On the opportunities and risks of foundation models. arXiv preprint
2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http: /github. com/ cam//oogle////oooo jax.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh. Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv preprint 1xXiv:2005.14665 2020.
Serkan Cabi, Sergio Gómez Colmenarejo, Matthew W. Hoffman, Misha Denil, Ziyu Wang, and Nando de Freitas. The intentional unintentional agent: Learning to solve many continuous control tasks simultaneously. In Conference on Robot Learning, 2017.
Quentin Cappart, Didier Chételat, Elias Khalil, Andrea Lodi, Christopher Morris, and Petar Veličković. Combinatorial optimization and reasoning with graph neural networks. arXiv preprint arXiv:2102.09544, 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, 2021a.
Tao Chen, Adithyavairavan Murali, and Abhinav Gupta. Hardware conditioned policies for multi-
robot transfer learning. In Advances in Neural Information Processing Systems, 2018.
Tao Chen, Jie Xu, and Pulkit Agrawal. A system for general in-hand object re-orientation. In Conference on Robot Learning, 2021
Xi Chen, Xiao Wang, Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Grycner, Basil Mustafa, Lucas Beyer, Alexander Kolesnikov, Joan Puigcerver, Nan Ding, Keran Rong, Hassan Akbari, Gaurav Mishra, Linting Xue, Ashish Thapliyal, James Bradbury. Weicheng Kuo, Mojtaba Seyedhosseini, Chao Jia, Burcu Karagol Ayan, Carlos Riquelme, Andreas Steiner, Anelia Angelova, Xiaohua Zhai, Neil Houlsby, and Radu Soricut. Pali: jointly- scaled multilingual language-image model. arXiv preprint arXiv:2209.0679 2022a.
Yuanpei Chen, Yaodong Yang, Tianhao Wu, Shengjie Wang. Xidong Feng, Jiechuang Jiang, Stephen Marcus McAleer, Hao Dong, Zongqing Lu, and Song-Chun Zhu. Towards human-level bimanual dexterous manipulation with reinforcement learning. arXiv preprint arXiv:2206. 200.02206.0.866 2022b.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Variational empowerment as representation learning for goal-based reinforcement learning. In International Conference Machine Learning, 2021.
François Chollet. On the measure of intelligence. arXiv preprint arXiv: 1911.0/57.. 2019.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh,
11

---

Published as a conference paper at ICLR 2023
Kensen Shi, Sasha Tsyyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Heier-Hellstem, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023// 2022.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International Conference Machine Learning, 2020.
Cédric Colas, Pierre Fournier, Olivier Sigaud, Mohamed Chetouani, and Pierre-Yves Oudeyer. CURIOUS: Intrinsically motivated modular multi-goal reinforcement learning. In International Conference on Machine Learning. 2019.
Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning, 2016.
Yuchen Cui, Scott Nickum, Abhinav Gupta, Vikash Kumar, and Aravind Rajeswaran. Can foundation models perform zero-shot task specification for robot manipulation? In Learning for Dynamics and Control Conference, 2022.
Wojeiech Marian Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant M. Jayakumar, Grzegorz Swirszez, and Max Jaderberg. Distilling policy distillation. arXiv preprint arXiv:1902.02166 2019.
Coline Devin, Abhishek Gupta, Trevor Darrell, Pieter Abbeel, and Sergey Levine. Learning modular neural network policies for multi-task and multi-robo transfer. In International Conference Robotics and Automation, 2017.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 110.v01810.04405, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RL Fast reinforcement learning via slow reinforcement learning. In International Conference on Learning Representations, 2017.
Benjamin Eysenbach, Xinyang Geng, Sergey Levine, and Ruslan Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. In Advances in Neural Information Processing Systems, 2020.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning to achieve goals via recursive classification. In International Conference Learning Representations, 2021.
Gilbert Feng, Hongbo Zhang, Zhongyu Li, Xue Bin Peng, Bhuvan Basireddy, Linzhu Yue, Zhitao Song, Lizhi Yang. Yunhui Liu, Koushil Sreenath, and Sergey Levine. Genloco: Generalized locomotion controllers for quadrupedal robots. arXiv preprint 4rXiv:2209.05009 2022.
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax differentiable physics engine for large scale rigid body simulation, 2021. URL http: //github. com/ google/aa poooiee/rxx....
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Advances in Neural Information Processing Systems, 2021.
12

---

Published as a conference paper at ICLR 2023
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, 2019.
Hiroki Furuta, Tadashi Kozuno, Tatsuya Matsushima, Yutaka Matsuo, and Shixiang Shane Gu. Co- adaptation of algorithmic and implementational innovations in inference- based deep reinforcement learning. In Advances in Neural Information Processing Systems, 2021a
Hiroki Furuta, Tatsuya Matsushima, Tadashi Kozuno, Yutaka Matsuo, Sergey Levine, Ofir Nachum, -eermation-thoortiic
and Shixiang Shane Gu. Policy information capacity: Information theoretic measure for task complexity in deep reinforcement learning. In International Conference on Machine Learning, 2021b.
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. In International Conference Learning Representations, 2022.
Seyed Kamyar Seyed Ghasemipour, Shixiang Shane Gu, and Richard Zemel. Smile: Scalable meta inverse reinforcement learning through context-conditional policies. In Advances in Neural Information Processing Systems, 2019.
Seyed Kamyar Seyed Ghasemipour, Satoshi Kataoka, Byron David, Daniel Freeman, Shixiang Shane Gu, and Igor Mordatch. Blocks assemble! Learning to assemble with large-scale structured reinforcement learning. International Conference on Machine Learning, 2022.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. In International Conference on Learning Representations, 2021.
Jonathan Godwin, Thomas Keck, Peter Battaglia, Victor Bapst, Thomas Kipf, Yujia Li, Kimberly Stachenfeld, Petar Veličković, and Alvaro Sanchez-Gonzalez, Jraph: A library for graph neural networks in jax., 2020. URL http: // thub. com/ deepmi nd/ jraph.
Shixiang Shane Gu, Manfred Diaz, Daniel C. Freeman, Hiroki Furuta, Seyed Kamyar Seyed Ghasemipour, Anton Raichuk, Byron David, Erik Frey, Erwin Coumans, and Olivier Bachem. Brax- lines: Fast and interactive toolkit for rl-driven behavior engineering beyond reward maximization. arXiv preprint arXiv:21 10.04686, 2021a.
Xiuye Gu, Tsung-Y Yi Lin, Weicheng Kuo, and Yin Cui. Open-vocabulary object detection via vision and language knowledge distillation. arXiv preprint arxiv:2104.1392. 2021b.
Agrim Gupta, Silvio Savarese, Surya Ganguli, and Li Fei-Fei. Embodied intelligence via learning and evolution. Nature Communications, 2021.
Agrim Gupta, Linxi Fan, Surya Ganguli, and Li Fei-Fei. Metamorph: Learning universal controllers
with transformers. In International Conference on Learning Representations, 2022.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv: 522.02259 2015.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL http:// //t:://iite sithub..... com/ wwwm///ooll//////ooooooooo Sooogllefffl flax.
Donald Hejna, Pieter Abbeel, and Lerrel Pinto. Hierarchically decoupled imitation for morphologi- cal transfer. In International Conference on Machine Learning, 2020.
Donald J. Hejna, Pieter Abbeel, and Lerrel Pinto. Task-agnostic morphology evolution. In Interna- tional Conference on Learning Representations, 2021.
Sunghoon Hong, Deunsol Yoon, and Kee-Fung Kim. Structure-aware transformer policy for inhomo- geneous multi-task reinforcement learning. In International Conference Learning Representa- tions, 2022.
Wenlong Huang, Igor Mordatch, and Deepak Pathak. One policy to control them all: Shared modular policies for agent-agnostic control. In International Conference on Machine Learning, 2020.
13

---

Published as a conference paper at ICLR 2023
Wenlong Huang, Pieter Abbeel, Deepak Pathak, and Igor Mordatch. Language models zero-shot planners: Extracting actionable knowledge for embodied agents. arXiv preprint arXiv:2201.0720. 2022a.
Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, Pierre Sermanet, Noah Brown, Tomas Jackson, Linda Luu, Sergey Levine, Karol Hausman, and Brian lchter. Inner monologue: Embodied reasoning through planning with language models. arXiv preprint arxiv:2207.088. 2006000 2022b.
Andrew Jaegle, Sebastian Borgeaud, Tean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman. Oriol Vinyals, and João Carreira. Perceiver io: general architec- ture for structured inputs & outputs. arXiv preprint AXiv:2107.147955 2021.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Advances in Neural Information Processing Systems, 2021.
Yiding Jiang, Shixiang Shane Gu, Kevin P Murphy, and Chelsea Finn. Language as an abstraction for hierarchical deep reinforcement learning. Advances in Neural Information Processing Systems, 2019.
Justin Johnson, Ranjay Krishna, Michael Stark, Li-Jia Li, David A. Shamma, Michael S. Bernstein, and Li Fei-Fei. Image retrieval using scene graphs. In Conference on Computer Vision and Puttern Recognition, 2015.
Arthur Juliani, Vincent-Pieree Berges, Ervin Teng, Andrew Cohen, Jonathan Harper, Chris Elion, Chris Goy, Yuan Gao, Hunter Henry, Marwan Mattar, et al. Unity: A general platform for intelligent agents. arXiv preprint 3xXiv:8809.0262) 2018.
Leslie Pack Kaelbling. Learning to achieve goals. In International Joint Conference on Artificial Intelligence, 1993.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning. 2018.
Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic rein- forcement learning at scale. arXiv preprint arXiv:2104.08212, 2021.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
Sateesh Kumar, Jonathan Zamora, Nicklas Hansen, Rishabh Jangir, and Xiaolong Wang. Graph inverse reinforcement learning from diverse videos. Conference on Robot Learning, 2022.
Vitaly Kurin, Maximilian Igl, Tim Rocktäschel, Wendelin Boehmer, and Shimon Whiteson. My body is a cage: the role of morphology in graph-based incompatible control. International Conference on Learning Representations, 2021.
Vitaly Kurin, Alessandro De Palma, Ilya Kostrikov, Shimon Whiteson, and M. Pawan Kumar. In defense of the unitary scalarization for deep multi-task learning. arXiv preprint 22X11:2201.04122, 2022.
Kuang-Huue Huei Lee, Ofir Nachum, Mengjiao Yang. Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski. and Igor Mordatch. Multi-game decision transformers. arXiv preprint arxiv:2205. 15241, 2022.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1)1334-1373, 2016.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement Iearror Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.0/6643 2020.
14

---

Published as a conference paper at ICLR 2023
Alexander C. Li, Letrel Pinto, and Pieter Abbeel. Generalized hindsight for reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi ma- nipulation using relational reinforcement learning. International Conference on Robotics and Automation, 2019.
Yunfei Li, Yilin Wu, Huazhe Xu, Xiaolong Wang, and Yi Wu. Solving compositional reinforcement learning problems via task reduction. In International Conference on Learning Representations, 2021.
Yunfei Li, Tian Gao, Jiaqi Yang. Huazhe Xu, and Yi Wu. Phasic self-imitative reduction for sparse- reward goal-conditioned reinforcement learning. In International Conference Machine Learning, 2022.
Xingyu Liu, Deepak Pathak, and Kris M. Kitani. Revolver: Continuous evolutionary models for robot-to-robot policy transfer. arXiv preprint arXiv:2202.05244 2022.
Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan, and Matthew Hausknecht. Working memory graphs. In International Conference on Machine Learning. 2020.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and
Pierre Sermanet. Learning latent plans from play. In Conference on Robot Learning, 2019.
Zhao Mandi, Pieter Abbeel, and Stephen James. On the effectiveness of fine-tuning versus meta- reinforcement learning. arXiv preprint
2022.
Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment- efficient reinforcement learning via model-based offline optimization. In International Conference on Learning Representations, 2021.
Russell Mendonca, Oleh Rybkin, Kostas Danilidis, Danijar Hafner, and Deepak Pathak. Discovering and achieving goals via world models. In Advances in Neural Information Processing Systems, 2021.
Yuki Noguchi, Tatsuya Matsushima, Yutaka Matsuo, and Shixiang Shane Gu. Tool as embodiment for recursive manipulation. arXiv preprint 1 2021.
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv: 130944999 2016.
Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. In International Conference on Learning Representations, 2016.
Emilio Parisotto. H. Francis Song. Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. arXiv preprint arXiv:1910.06764 2019.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count on luck: Why decision transformers fail in stochastic environments. arXiv preprint 220. 2022.
Deepak Pathak, Chris Lu, Trevor Darrell, Phillip Isola, and Alexei A. Efros. Learning to control self-assembling morphologies: A study of generalization via modularity. In Advances in neural information processing systems, 2019.
Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa. AMP: Adversarial motion priors for stylized physics-based character control. arXiv preprint arXiv:2104.02/88 2021.
Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal difference models: Model- free deep for model-based control. In International Conference on Learning Representations, 2018.
15

---

Published as a conference paper at ICLR 2023
Vitchyr H. Pong, Murtaza Dalai, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew-fit: State-covering sule-vupe reinforcement learning. In International Conference on Machine Learning, 2020.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. arXiv preprint 0rXiv:2/03.000200 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-61 2020.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy meta-reinforcemen learning via probabilistic context variables. In International Conference on Machine Learning, 2019.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. generalist agent. arXiv preprint arxiv:2205.06175 2022.
Andrei A. Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirk- patrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distillation. International Conference on Learning Representations, 2016.
Tom Schaul, Dan Horgan, Karol Gregor, and David Silver. Universal Value Function Approximators. In International Conference on Machine Learning, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning & modes with one stone. arXiv preprint arxiv.2206. 1151, 2022.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Cliport: What and where pathways for robotic manipulation. In Conference on Robot Learning, 2022a.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning, 2022b.
Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforcement learning. In International Conference on Learning Representations, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancy. Timothy Lillicrap, and Martin Riedmiller. Deepmind control suite. arXiv preprint arXiv:1801.00690 2018
Yee Whye Teh, Victor Bapst, Wojeiech Marian Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, 2017.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems, 2012.
Brandon Trabucco, Mariano Phielipp, and Glen Berseth. Anymorph: Learning transferable polices by inferring agent morphology. In International Conference Machine Learning, 2022.
16

---

Published as a conference paper at ICLR 2023
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, 2017.
Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint wrXiv:161..05766 2016.
Tingwu Wang. Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International Conference on Learning Representations, 2018.
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, B. Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In International Confer- on Machine Learning. 2022.
Ruihan Yang, Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. In Advances in Neural Information Processing Systems, 2020.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: benchmark and evaluation for multi-task and meta reinforcement learning. In Conference Robot Learning, 2019.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In Advances in Neural Information Processing Systems, 2020.
Ye Yuan, Yuda Song, Zhengyi Luo, Wen Sun, and Kris M. Kitani. Transform2act: Learning transform transform-and-comee and-control policy for efficient agent design. In International Conference on Learning Representations, 2022.
Andy Zeng, Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong, Ivan Krasin, Dan Duong, Vikas Sindhwani, and Johnny Lee. Transporter networks: Rearranging the visual world for robotic manipulation. In Conference on Robot Learning, 2020.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference on Machine Learning, 2022.
Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compo- sitional generalization in control. arXiv preprint arXiv:2203.05900 2022.
17

---

(1 1,..,,, L). Published as a conference paper at ICLR 2023
APPENDIX
A DETAILS OF IMPLEMENTATION
The hyperparameters we used are listed Table We implement MLP and Transformer policy with Jax (Bradbury et al., 2018) and Flax (Heek et al., 2020). For the implementation of GNN policy, we use graph neural network library, Jraph (Godwin al., 2020).
Method
Hyperparameter
Value
Shared
Learning rate
3e-4
Batch size
256
Gradient clipping
0.1
Activation function
ReI.U
Gradient steps
100k
MLP
Hidden size
1024
of layers
2
Transformer
Embedding size Attention hidden size
256
512
# of attention heads
2
# of attention layers
3
GNN
Hidden size
256
# of layers (per node)
3
Table 3:
Hyperparameters for each method.
Transformer for Morphology-Taske Graph Transformer first encodes morphology-task graph to the latent representation vector zo with shared single-layer MLP and learnable position embedding (PE) (in case of MTGv1, we omit Sy. but include it to corresponding node observation s' instead)::
Zo MMLP("".....,MI
Voll), MLP(ss)] PE,
then multi-head attention (MHA) (Vaswani et al., 2017) and layer normalization (LayerNorm) (Ba al., 2016) are recursively applied to latent representation 2ạ at 1-th layer,
z/ LayerNorm(MHA(
z1-1),
21 LayerNorm(MLP(44 21).
Before decoding the action per module from the last-layer latent representation zz, we employ the residual connection of the node features (Kurin et al., 2021),
MLP([2L, s'), I).... MLP([=] Vml /Voll)), MLP([=//mllt) Vnl++1 ,yy)]
where shared MLP has a single layer and tanh activation to clip the output within the range of [-1, 1]. We mask out the output from the goal modules or modules that have no actuators to ensure action size as E-
Details of Morphology-Tast Graph v1 and v2 Our proposed morphology-tas graph expresses the agent's s observations, actions, and goals or tasks in unified graph representation, while preserving graph v1 accepts the morphological graph,
the geometric structure of the task. Morphology- encoded from the agent's geometric information, as input-output interface, and merges positional goal information as part of corresponding node features, and morphology-tas graph treats given goals extra disjoint nodes of morphological graph. In practice, pad the incompatible parts of goal information by 0. For both MTGv1 and MTGv2, we prepare the binary vectors that have indicators the target (or goal) nodes. MTGv1 obtains the products of the goal information and binary vectors, then leverages them as extra node features. MTGGv2 just takes such binary vectors as extra node features (and also includes goals extra nodes).
18

---

Published as a conference paper at ICLR 2023
B DETAILS OF MXT-BENCH
B.1 MORPHOLOGY
We prepare base blueprints; ant, centipede, claw, and worm, for procedural scene generation (Fig- 6). Base ant has torso and legs with 2 joints per limb, and we could procedurally generate the agents with different number of limbs. Base centipede has bodies and legs with joints per limb, and could procedurally generate the agents with different number of bodies. Base claw has 1 torso and 2 legs with 4 joints per limb (each leg consists of modules), and we could procedurally generate the agents with different number of limbs. Base worm has 2 bodies and legs, and could procedurally generate the agents with different number of bodies.
Figure 6: Examples of procedurally-genenee morphology from base blueprints in MxT-Bench. From left to right, each figure shows the example of 5-leg ant, 4-body centipede, 6-1eg claw, and 5-body worm.
Furthermore, we develop the finctiionality for morphology diversification, with missing, mass, and size parameters (Figure 7). Missing randomization lacks one module at one leg. This might be an equivalent situation that one leg is broken. Mass randomization changes the default mass of each module with specified scales. The appearance does not change, but certainly, the dynamics would differ. Size randomization changes the default length and radius of each module with specified scales.
Figure 7: Examples of diversified morphology from base blueprints in MxT-Bench. From left to right, each figure shows the example of -eigg--misssng ant, S-leg-size-randdmmized ant, 4-body-1-missing centipede and 3-body--ize--rnddomis centipede.
B.2 TASK
We prepare 4 base tasks with parameterized goal distributions; reach, touch, twisters, and push, for procedural task generation (Figure 8). Reach task requires the agents to put their one leg to the given goal position (XY). The variant, reach hard task, represents that the goal distribution farther than reach task. Touch task requires the agents to contact their body or torso to the movable ball (i.c. movable ball a goal). Twisters
multi-goal problem; the agents should satisfy given goals at the same time. There are two basic constraints: reach and handsup. Handsup requires the agents to raise their one leg to the given goal Z height. Twisters has some combinations like reach teach_handsup reach hard _andssupa reach2. _aandsupp or reach 2hanssp2. For instance, in reach handsup2, the agents should put their one leg to the given goal position, and raise their two legs to the given goal heights simultaneously. The XY position goals sampled from uniform distribution with predefined donut-shape range, and Z height goals are also from uniform distribution with predefined range. Push task requires the agents move the box object the given goal position (XY) sampled
19

---

Published as a conference paper at ICLR 2023
from uniform distribution with predefined range. Since it has richer interaction with the object, push might be a more difficult task than those three. We use this task in Appendix L.
Figure 8: Examples of pre-defined task in MxT-Bench. From left to right, each figure shows the example of reach, touch, twisters, and push task.
B.3 CUSTOM MORPHOLOGY
MxT-Bench also supports custom morphology import used in previous work. For instance, Gupta al. (2022) propose unimal agents that are generated via evolutional strategy and designed for MuJoCo. Since they are not manually designed, their morphologies seem more diverse than our ant, centipede, claw, and worm. We inspect unimals whether they are suitable for goal-reaching. and include 72 morphologies from there (and use 60 morphologies for the experiments). Figure 9 shows some example of unimals.
Figure Examples of unimal agents, adapted from Gupta al. (2022)
B.4 COMPARISON TO EXISTING BENCHMARK
Since the current RL community has paid much attention embodied control far, there no suitable benchmarks to quantify the generalization over various tasks beyond single locomotion tasks and morphologies at the same time. In addition, the scalability various morphologies tasks seems to be required for the benchmark, because we should avoid "overfitting" to manually-designe tasks.
As summarized in Figure 4, MuloCo (Todorov et al., 2012) or DM Control (Tassa et al., 2018), the most popular benchmarks in the continuous control domain, could not evaluate task morphology generalization; they only have single morphology for a single task as pre-defined environment. Yu al. (2019) propose robot manipulation benchmark for meta RI, but does not care about the morphology. Furthermore, while it has quite a diverse set of tasks, the scalability of environments seems to be limited. In contrast, previous morphology-agnosti
RL works (Huang et al., 2020; Kurin et al., 2021; Hong et al., 2022; Trabucco et al., 2022) have a set of different morphologies adapted from MuJoCo agents. Gupta et al. (2022) also provide much larger set of agents that are produced via ooittiopti of morphology and task rewards by the evolutionary strategy (Gupta et al., 2021), with kinematics and dynamics randomization. However, those works only aim to solve single locomotion tasks, i.e. running forward as fast as possible.
20

---

Published as a conference paper at ICLR 2023
C DETAILS OF EXPERT DATA GENERATION
We train the single-task PPO (Schulman et al., 2017) at each environment to obtain the expert policy. For the reward function, basically adopt dense reward (except for push task) -dw 8g)) until agents satisfy a given condition, 1 (s" $y) I min'.. See Appendix for the threshold d After convergence,
min 12k
collect the proficient behaviors. Unless otherwise specified, transitions per environment. The average normalized final distances of those datasets are mostly less than 0.1 (Figure 10).
0.11
0.10
9.08
9.34
0.00
0.00
Figure 10: Average normalized final distance in each dataset (Table 4),
D ENVIRONMENTS DIVISION
Throughout the experiments, we test a lot of morphology-task combinations to investigate the in-distribution generalization, compositional generalization for morphology and task, out-of- distribution generalization. In this section, we list up the combination of environments used in the experiments.
Table 4 explains the combinations for the experiments of in-distribution generalization, and composi- tional generalization for morphology (both zero-shot transfer and fine-tuning) in Table and Figure 5. For compositional morphology evaluation, Morph-Test division as evaluation environments. For the evaluation of dataset size and the number of morphology -takk combinations (Appendix H), In- n--Disributton, 12 Env, and 25 Env division as train datasets and test environments. Following prior works (I (Huang et al., 2020; Kurin et al., 2021; have chosen Morph-Test division that would not exceed the maximum number
Morph-Train - division training dataset and
Gupta et al., 2022), of nodes in Morph-Traij Train division for convenience. We have checked our trained model could be generalized to the agents with an extra number of limbs (c.g. 7-leg ant or 8-body centipede), but have not extensively increased the number of limbs until the model could not deal with them. We note that this may limit the morphology generalization in the open-ended evaluation.
Table 6 also explains the combinations for the experiments of in-distribution generalization with unimals (Gupta et al., 2022) Table and Table 12.
Table 5 provides the combinations for the experiments of compositional generalization for task and -f-ttiiiiitiitoo generalization (both zero-shot transfer and fine-tuning) in Table 1 and Figure We use them as test environments, and for training datasets, we leverage In-Distribution division in Table
In addition, we extensively evaluate the compositional generalization for task and out-of-distribution generalization with more different unseen task, such as push (see Appendix for the details). Table also shows the environment division for both zero-shot transfer and fine-tuning.
21

---

Published as a conference paper at ICLR 2023
Sub-domain
Environment
Distribation
Morrh-IIraia Morph-Iest
snt nerch
0.1
8.75
antreach
201 reach_
0.1
8.75
201 7
0.1
8.75
ant 4
0.5
ant 1
snt t
0.5
3.5
ant nach? handup
(0.1.0.1.0.)
(1.4, 1.4. 0.655)
reach2 hundup_ 5
(0.1,0.1,00 0.11
[144 1.4, 0.625)
antrracc22 handupp
[0.1. o...o.1....
[1.4, 1.4, 0.625
reach, handsp2
0..10.11,00 0.11
0.55
ant_twisters
aetreach,
[0.1. 0.1.0.1.00 OLL)
(1.5,0.55, 0.55, 0.55]
reach, handup2_5
(0.1..11,00 0.11
0.550
mntreach, handsahnddppp
10.1 OL11
(155,..5,,00 0.55)
hundup_4
mitreach, handspp
[..1.0.1]
[4.5, 0.55
handspp6
claw reach 2
0.1
8.7
reach
clew reuch
claw reach
0.1
875
reach.
claw reach
011
875
3
clow touch
cls 1much
0,5
3.5
Juuch_5
cls 6msch
0.5
3.5
claw reach handsup.
[0.1.0.1)
(14.0.625)
cluw necc hindeup
[0.1,0.13
[1450.625]
clow, claw_twistees twisters
claw reach handsup.
(0.1.0.1)
(1.40.625)
claw peach_handoup.
114,
claw reach hard handspp55
(0.1.0.1)
[4.5, 0.55)
claw Jandspp
[0.1.0.1
centppdde mach cenppode 3
0.1
3.0
centipale rench
cnntppode mach 4
011
3.0
cenppede
antipodee rach. 6
011
3.0
centipede
jentipodt_touee 3ouch__3
0.5
10.5
centipede
centipede_uree touch
centipode_toucs souch_5
0.5
10.5
centipede_
centipode_soocs souch_7
05
10.5
worm touch
touch.4
Wome
touch
6 worm touch
Table 4: The combinations environments used in the experiments in-distribution generalization, and compositional generalization for morphology.
22

---

Published as a conference paper at ICLR 2023
Sub-domain
Environment
de
de
Tusk-Test
OOD-Test
reach hard ant hot_reach__aae hard 3
0.1
11.0
ant rerrrruuhhharreee 4
0.1
11.0
ant ant_reach___ hard_5
0.1
11.0
ant reache hard_ faa_rechhhaad__f6
0.1
11.0
centipede reach hard haad__
0.1
centipede_ree reacc_eaaarreeeeeeeeeeeeeeeeaaae hard
0.1
centipede_ reach hard
centipede_reane rewcc_achhhaa seach_hard_$
0.1
5.5
centipede reach rea_ch_aand_66666 hard_6
0.1
5.5
centipede wer_rreahhhhhdeeeeeeeeeee hard__
ant fnl_reuch_hardd hard_4_b
0.1
11.0
ant_reuch__ss ant_reach_hart___2 h_nd___bb
ant ant_reuch_hane hard 4__0___m_000_0 0.5_110_0.51103..0
0.1
11.0
ant se__recchhaadd 1_0_4___________ _..__..0___111000
0.1
11.0
ant reach diverse
ant reach ant ant_reach_haal hard size 0.9_1.0__1._
mass, 1.0. 3.0...33.00
11.0
0.1
11.0
ant reach hard 5 mass, _0.__0_1.000 0.5_1.0_33.
0.1
11.0
ant hat__rach___as hard mass, 0.5_1.0_1..0
0.1
11.0
ant ant_reach_ hard mass_ 1.0_3.003.0
ant ant_reuch_haat hard 5 size 0..___09_.011111
0.1
11.0
centipede_raat reach hard 3___r____00____
centipede_reat reaudde__eee_____eeeeeee reach_hard_____ hard_a_dd__bbb_ 33_____0
0.1
centipede_eaach_ @racc_hard_3_b_sr.
0.1
5.5
centipede_reats rracc. hard_3_b_l_____
0.1
centipede_reet
rwwcchhharddaaeeeeeeeeeaaaaaaaae hard 4_b_r_0
0.1
centipede_reat rewcc_eaaaaaeeeeeeeeeeeeeeeeeaae 1_a_____b_r___
0.1
5.5
centipede reach hard_4 hand_4_b_L_
@each_hard_____
0.1
5.5
centipede_reat reach hard diverse
centipede reach hard 4_b_r_2
centipede_reaat hhteereaahhhaaaaaaaaeeeeeeeeeeee _______r_____2__2-
0.1
5.5
centipede_reade saa__raahhharrre Saard_3_iiee___ 0___sie___9_110
09111.0_111
0.1
5.5
centipede reach. hard ____mas_____ 0.5_1.0._33
0.1
centipede_reee hard ___maa__00 0..__1........
0.1
centipede_roe reach hard 3 mass 1.0_3.033.0
0.1
5.5
centipede_reate Beaeddd__raa__heeeeee Barch_hard__4_ 4 size 0.9__.00_111
0.1
5.5
centipede_reat hard
0.5_ 0.5_1.03..0
0.1
centipederreat www_rrrahhharrreeeeeeeeeeeeeeeee mass 1.5_1.0_1.0
0.1
5.5
centipede_reate Bectll__eec____eee Baicchharrd_4mmee mass 1.0__.0__.000
0.1
5.5
Table The combinations environments used experiments compositional generalization out-of-distribution generalization.
23

---

Published as a conference paper at ICLR 2023
unioall
reach
touch
twislers
5506-0-13-7-12-26-11
5906-12-12:0-11130000
875
5506 5906-15-16.0-14--7..1
0.35
3.5
5505-6.5-01-14-2442
(0.1,0.1
(1.5, 0.55)
(0.10.11
(15,0.55)
5996-12-1--1-15-33--01
875
5506-15-66-0--2221-00
8.75
5996-6---1-499-9--06
5506-0---0-12--45.66
[0.1,0.1
0.550.55
130-022-1400
5906-2-0-0-11-22--1
[0.10.31
(1.5.055)
1-06-7-6-17-12-20-
0.35
3.5
5506-0-7-01-15.34-1
3.5
1VX6-12--17--08--66
5506-2-16-0111--5--25
5506-8-11-01-15228
0.35
3.5
5506-1-
1906-1-12-17-11-10--
0.1
8.75
5506-12-6-17112-2--66
875
5906-2-17-12-0-11--6
875
5996-8-1-41-11--2---6
0.35
3.5
5500-1-13-7-11-33--66
0.1
8.75
(0.1.0.1
(15,0.55)
5906-2-9-1--11-1044
8.75
5306
5996-1-55-7-77-55-477
(0.1.031
(1..0055)
1W0-113-0-17-11-25-44
0.1
8:75
5906---1041-14-19906
875
5906-8-16-00 190.8-11--0211-----2
(0.1,0.1
(155,0555
5506 1506-12-00-22-22-111
(0.10..11
(1.5,055)
5506-13-1--1-11-099BBB
[0.1.03]
[15,0555)
50063-15--1-11-36--0
5506-8-77-17499888.299
3506-155-0-19-22.33
(0.1,0.1
[.5,0555
5506
13----33-0222
(0.1.0.1)
(15,0055)
5505.3-15-1712-188.68
8.75
5506-65-02-21-39220
(0.1.0.)
(L5,0.55)
5906-10-0-1-15-44-55
0.35
3.5
5906-11--02-2144.007
8775
5906-4112.1-15510--92
5905..6--1-15.22.66
(0.1.0.1)
[1.5.
5506-10-12-022
100-10-12-02-1233--1
0.1
8.75
5506-13-5-22-21-35541
875
5306-1-11-1-14-33-477
5996-9-12-01-10-2222
0.1
8:75
5506.
100--00-13-0-11-03--4
0.35
3.5
(0.1.0.]
(1550.55)
4505-4116.1-05-464--
(0.1.0.1]
(1550.55)
5996-10-14417110-88.44
0.35
3.5
5906-14-1-01-12-42-42
[0.1.0.1
(1.50055)
5906-4-3-01-09225.18
3.5
5906-9-3-01-14-23.90
(0.1.0.11
(1.5,
5506-
300--00---01-15-2--1
0.35
3.5
5906-14-15-1-11-20--5
3.5
5906-9-7-01-13-44-02
875
5506-10.3-11-41-449.22
5906-14-242-15-144466
875
5506-5-16-0221-1555
S506 1006-11--2400 1112--01-14--1-400
0.35
3.5
-1----1-55-5552
5506 5-0----11------55555
[0.1...1
(1.5,055)
5506-11-4-1-11255--0
[0.1,0.]
[1.,,055)
5906-15-11001-11-04--44
3.5
5906-6-1--01-14-16.099
0.35
3.5
5906-11.6-1-11-43.666
190.-15-11.0111-55..35
[0.1,0.1]
(1.5,0555
5506 00066---1-00-16-44
0.35
3.5
5506- 196-12-11-17455-56--6
0.35
3.5
5506-55-11-1442-11-22
5506-6-3-01-55-9--30
3.5
Table 6: Unimal IDs we adapted from Gupta et al. (2022). We inspect 100 morphologies and select 72 morphologies that work healthily, In the experiment of Table and Table 12, we select 20 morphologies each for tasks (reach, touch, twisters) listed above.
24

---

Published as a conference paper at ICLR 2023
Sub-domain
Environment
d
d
Task- Test
OOD- Test
ant pash_p ant push_4 4
1.0
4.0
unt_push
ant ani_push_5
1.0
4.0
ant pal_pus__6
1.0
4.0
centipede_ppng push_
1,5
3,75
centipede_ 4
1.5
3,75
centipede push
centipede_pushe
1.5
3.75
centipede_push
1.5
3.75
centipede__psh.
1.5
3.75
ant ant_push_3 porm__push_4
1.5
3.25
worm_push
1.5
3.25
worm_push_55
1.5
3.25
ant pot__ushh_4.. b_b
1.0
4.0
ant p_t_p_us___bb
1.0
4.0
ant ant_push_4 4_ mass 0.5_1.03..0
ant ant_push mass 0.5_1.0_1.0
1.0
4.0
unt_push_divers
ant ant_push__4__ 4_mass 1.0_3.033.0
1.0
4.0
ant ant_push_44__e Beuush__size__00 009 __.10_1.11
1.0
4.0
ant ant_push_5. 1aaih_5_mas____ mass 0.5_1.033.0
1.0
ant pot__push_5, 5 mass 0.5_110_ 1.0_1110
1.0
4.0
ant pat__push_5.. 1achh___mass_ mass____________________________ 1.0_3.0.3.0
1.0
4.0
ant_push__5 size 1.9_1.0_1..
1.0
4.0
centipede push__ ______
1.5
3.75
centipede_puv _uchdd____h______ 3_b___
1.5
3.75
centipede_pung Bacpsddd__ps______ 3_ 3_b_r_1
1.5
3.75
sentipede_pustee p 3_b___1
1.5
3.75
1.5
3.75
cennipede_push_11 __b____ssh______________________
1.5
3.75
centipede_puse
Btttieed__pssh___b___1 1_______b__11
1.5
3.75
centipede_ push diverse
1.5
3.75
centipede_push___bb..
1.5
3.75
Becripede_pushh___ 3_siee__________ 0.9 1.911.0_111
1.5
3.75
centipede_pushh_.1 mass 0.5__.003.0
1.5
3.75
centipede_push__3 0.5 1.0
1.5
3.75
centipede_pushe
Bouub____as______3____0_ 1.0_3.0_..0
1.5
3.75
Secipede__psh__444 siize____iie_____ 0...__09.1....11
1.5
3.75
centipede_puse
pwwpppusshhaaaaee mass. _.5__110_3.0
1.5
3.75
teciipede_pushh.e
0.5 1.0
1.5
3.75
centipede_pute push mass _a__aa___________ 3.0
1.5
3.75
Table 7: The extra combinations of environments used in the experiments of compositional generalization for task and out-of-distrihuttin generalization with push task.
25

---

Published as a conference paper at ICLR 2023
E ARCHITECTURE SELECTION: TOKENIZED MORPHOLOGY-TASS GRAPH
In the recent literature, offline RL is considered as supervised sequential modeling problem (Chen al., 2021a), and some works (Janner al., 2021; Reed al., 2022) tokenize the continuous observations and actions, like an analogy of vision transformer (Dosovitskiy et al., 2020); treating input modality like language.
As part of offline architecture selection, we examine the effectiveness of tokenization. We mainly follow the protocol in Reed al. (2022); we first apply mu-law encoding (Oord et al., 2016) to morphology-task graph representation:
log(1a μ + 1)
mu mu_law<x)== = sgn(z)
log( M + 1)²
with μ = 100 and M = 256. This pre-processing could normalize the input to the range of [o, 1]. Then, discretize pre-processed observations and actions with 1024 bins. To examine the broader range of design choice, we prepare the following 6 variants of tokenized morphology-task graph: whether layer normalization is added after embedding function (LN) (Bhatt al., 2019; Parisotto et al., 2019; Furuta et al., 202 Chen et al., 2021a), or outputting discretized action (D), smoothed action by taking the average of bins (DA), and continuous values directly (C).
As shown in Table 8, predicting continuous value reveals the better performance than discretized action maybe due to some approximation errors. However, the performance of Token-MTGv2 (C) or Token-MTGG (C, LN) is still lower than MTGv2 itself. This might be because tokenization looses some morphological invariance among nodes. In Table 1, we adopt Token-MTee (C) for comparison.
Sab-doanalo
Teles-MTGr220
Token-MTCs2LD
Tekou-MTGx2
Bokcn-HGGv2LC
(MT(22)
ant reach
0.905 2 D.0
0:152 6.001
0.335 = 0.01
0.3100 2: 0.02
0.5215 2 0.00
0.3846 0/13
0.3200 2-0
stLacechh
1.0073 2 D.01
118511 3.011
0.4655 2 I.00
183338 0.00
18800 2 8.11
0.3455 2.00
9.2888 3.001
2
ckorJosch
1.000
10354422 6.001
0.523 0.07
10292 0.03
LUB85 0,01
0.3129:001
022133
81672
0.01
0AT26 0.000
0.001 : 0.00
06341 00.
06607 8:0
0.3110 9003
0.3112 101
comtipde nouch
01967
0.516 1.111
0.00
04771
0:15
comppde
Inach
1.0996
1/331
09357 0.14
10519
0.01
0.035
0.4952
Assrage liss.
08824 52 000
03232 012
0,3972 0.02
03348 0.0
0.225 0.02
0.342 000
0.1128 012
Table 8: The average normalized final distance in in-distribution evaluation. We extensively evaluate the tokenized morphology-task graph variants, similar to Reed et al. (2022).
26

---

Published as a conference paper at ICLR 2023
F ARCHITECTURE SELECTION: POSITION EMBEDDING
As a part of architecture selection, we investigate whether position embedding (PE) contributes to the generalization. Our empirical results in Table 9 suggest that, multi-task goal reaching performance seems comparable between those, However, in more diverse morphology domains, PE plays an important role. Therefore, we include PE into a default design.
Sub-domain
Transformer (MTGv2) (w/ PE) Transformer (MTGv2) (w/o PE)
ant_reach ant touch dla_reachh
0,3206 = 0.06
0.3966 + 0.08
0.2668 + 0.08
0.4573 0.14
0.1039 0.05
0.1569 0.02
0,3581 = 0.04
0.3508 + 0.03
claw_touch
0.2573 = 0.08
0.3278 0.07
claw twisters
0.3442 0.04
0.3071 0.04
centipede reach
0,1057 + 0.114
0.0610 + 0.02
centipede touch
0.3869 + 0.04
0.1687 0.05
worm_touch Norm_touch
0,8952 0.05
0.8427 0.03
Average Dist.
0.3128 = 0.02
0.3142 0.03
unimal reach
0,4532 = 0.01
0.5856 0.02
unimal_touch touch
0,4461 + 0.05
0.3799 + 0.06
unimal twisters
0.3540 +: 0.02
0.3236 0.04
Average Dist.
0,4178 + 0.01
0.4297 + 0.03
Table 9: The average normalized final distance in-distribution evaluation. We compare the effect of position embedding.
G REPRESENTATION SELECTION: MORPHOLOGY-TASK GRAPH WITH HISTORY
Some meta RL multi-task RL algorithms take historical observations as inputs, or leverage recurrent neural networks to encode the temporal information the tasks (Wang al., 2016; Teh et al., 2017; Rakelly et al., 2019). To examine whether morphology-tase graph could be augmented with historical observations, we test the morphology-tast graph with history. which concatenates recent H frames morphology-task graph as inputs while predicting one-step actions, in various types of morphology-tase generalization with ant agents. In all the setting, we choose ant (re_ccech, touch, twisters) (in Table 4) for training. In morphology generalization setting, tasks for evaluation. In task generalization setting, we use ant_reach_______ hard (in Table 5) as a test set. In out-of-distribution setting, use ant_reach_ history length as H = 3.
hold out ant agent
(in Table 5) for evaluation. We also set the
Table 10 implies that when the policy faces unseen tasks (i.e. in Compositional (Task) or Out-of- Distribution settings), morphology-tan -task graph with history may help to improve the performance, which seems promising results to extend our framework more complex tasks.
MI.P
Transfrmmrr (MTGv1)
Transformer (MTGv2)
Transformer (NTGrv-history)
In-Distribution
0.3035
0.03
0.1364
0.04
0.0961 1: 0.02
0.1015 0.04
Compositional (Morphology)
0.7480
0.02
0.1477
0.04
0.1010 1 0.01
0,1070 +: 0.03
Compositional (Task)
0.4468
0.11
0.2243
0.03
0.2201 4 0.04
0.2170 0.01
Dutto--Dist
0.6788 0.01
0.3663
0.05
0.3370 + 0.03
0.3316 0.05
Table 10: The average normalized final distance in various types of morphology-task generalization on MxT- Bench (cspecially, ant). We concatenate recent three frames morphology-task graph as inputs to Transformer. The results show that when the policy faces unseen tasks, MTGv2-history may help to improve the performance.
27

---

Published as a conference paper at ICLR 2023
H DOES TRANSFORMER WITH MORPHOLOGY-TASK GRAPH SCALE UP WITH DATAGETTAORFFOOLOOYYYSAAAAAAAOAD SIZE?
The important aspect of the success in large language models is the scalability the size of training data, the number of tasks for joint-training, and the number of parameters (Raffel et al., 2020; Brown al., 2020). One natural question is whether similar trend holds even in RL.
Figure 11 suggests that the performance can get better if we increase the number of datasets and the number of parameters. The performance of Transformer with 0.4M parameters is equivalent that of MLP with 3.IM. In contrast, when we increase the number of environments, the performance degrades while Transformer with MTGv1 and MTGv2 surpasses the degree of degradation, which seems inevitable trends in multi-task RL (Yu et al., 2020; Kurin et al., 2022) and an important future direction towards generalist controllers.
Avg. Normelized Distance
0.7
0.50
0.40
0.4
0.30
0.3
Number Transitions Environment)
Number of Envtronments
Number Model Parameters
Transformer
(MTGv2)
Figure 11: The average normalized final distance with different size of datasets (left), morphology-tabe com- binations (middle), and model size (right). The smaller value means the better multi-task performance (see Appendix D for environment division). These results suggest that the performance increase the number of datasets, and morphology-tas graph can surpasses the degradation the performance when we increase the number environments. Transformer parameter-efficient than MLP, and improves the performance as many parameters.
get better when
I PERCENTAGE OF IMPROVEMENT
For clarification, we compute the percentage of improvement between two average normalized final distances (defined Equation 3) di and d2 as follows (d) dz):
d2 d
100
d2
28

---

Published as a conference paper at ICLR 2023
J ADDITIONAL RESULTS
In this section, we provide the detailed performance of in-distribution generalization (Table 11 and Table 12), compositional morphology and task generalization (Table 13 and Table 14), and out-of- distribution generalization (Table 15). For fine-tuning experiments, we summarized the detailed scores of Figure in Table 16.
Sub-domain
Random
MI.P
GNN (MTGv1)
Transformer (MTGv1)
Transformer (MTGv2)
unt ant_reuch
0.9637 + 0.02
0.6426 + 0.03
0.6240 0.03
0.3657 = 0.04
0.3206 0.06
ant touch
1,0817 0.02
0.3689 0.02
0,4434 0.03
0.1140 + 0.03
0,2668 0.08
0.9113 0.84
0.3708 0.01
0,2513
0.2517 0.02
0.1039 0.05
claw_reach
1.0760 + 0.20
0.6617
- 0.00
0.6214 0.02
0,7158 = 0.03
0.3581 0.04
claw claw_toudh touch
1.5240 0.88
0.6824 0.06
0,3135 0.06
0,6121 = 0.04
0,2573 0.08
claw jtw_ttssters
1.3786 1.42
0.4907 0.03
0.6063 0.07
0,4614 0.07
0,3442 0.04
centipede_ peacpeeeeeeachhhhheeee reach
0.5843 + 0.35
0.0803 + 0.02
0.1088 0.02
0.0981 + 0.03
0.1057 = 0.04
centiedde_ touch
1.0077 1: 0.01
0,4743 0.03
0,5089 0.01
0.4609 = 0.07
0,3869 0.04
worm_ touch
2.7087 1.73
1.104
1.0599 0.04
0.7708 0.03
0.8952 0.05
Average Dist.
0.4419:
0.5150
0.4776 0.01
0,4069 0.02
0.3128 0.02
Table 11:
The average normalized final distance for in-distribution evaluation on MxT-Bench (as shown in
Table 1).
Sub-domain
Random
ML.P
Transformer (MTGv1)
Transformer (MTGv2)
unimal_reac reach
0,9662 + 0.01
0.7448 + 0.02
0.5692 0.01
0.4532 + 0.01
unimal, touch
1.1302
0.7634
0.5290 0.04
0.4461 0.05
unimal-twiste
0.6305 + 0.02
0.5027 + 0.02
0.3534 0.03
0.3540 + 0.02
Average Dist.
0.9090 + 0.03 0.6703 + 0.01
0.4839 0.02
0.4178 + 0.01
Table average normalized final distance in-distribation evaluation on MxT-Bench with challenging morphologies from Gupta et al. (2022) (as shown in Table 1).
Sub-domain
Random
MLP
Transformer (MTGvI)
Transformer (MTGv2)
ant reach
0.9697 = 0.03
0.8541 =
0.5511 +
0.4170 0.04
ant Jouch
1.0821 + 0.00
0.7742 : 0.07
0.4464 : 0.07
0.3752 +: 0.12
ant_twisters
0,9215 = 0.95
0,5356 = 0.02
0.2569 : 0.04
0,2608 + 0.02
claw_ Lewch
0.9915 0.04
0.9370 =
0.7332
0.06
claw touch
1.0695 + 0.02
1.0283 + 0.02
0.7074 0.10
0.1375 0.01
claw_twin _ww_ttwsters
0.7001 = 0.17
0.5456 = 0.02
0,5486 : 0.02
0.5589 + 0.05
centipede, reach
centipede_ touch
1.0121 = 0.01
0,7632 = 0.02
0.4971 = 0.06
0.3474 = 0.09
worm worn_touch
2,7376 + 2.23
1.1417 = 0.10
0.8608 - 0.07
1.0110 + 0.03
Average Dist.
1.1419 = 0.41 0,7216 = 0.01
0.4940 1 0.01
0.4066 + 0.01
Table 13: The average normalized final distance for compositional morphology evaluation MxT-Bench (as shown Table l).
29

---

Published as a conference paper at ICLR 2023
Sub-domain
Random
MLP
Transformer (MTGy1)
Transformer (MTG+2)
ant setcr hard
0.9542 + 0.01
0.8299 0.03
0.6176 0.07
0.4522 0.06
centipede_rent reach hard
0.8443 0.02
0.5689 0.03
0.4770 0.05
0.4412 0.05
Average Dist.
0.8932 +: 0.01
0.6849 + 0.01
0.5395 0.04
0.4461 = 0.05
Table 14: The average normalized final distance for compositional task evaluation MxT-Bench (as shown in
Table 1).
Sub-domain
Random
MI.P
Transformer (MTGv1)
Transformer (MTGv2)
ant sea_ceaa_haal hard diverse rentipode_reach hard diverse
0.9520 0.01
0.8288 0.02
0.6798 0.04
0.5365 = 0.05
0.8660 + 0.02
0.7546 1: 0.02
0.6130 + 0.03
0,5208 = 0.04
Average Dist.
0.8979 +: 0.01
0.7821 +: 0.02
0.6144 +: 0.04
0,5266 : 0.04
Table 15: The average normalized final distance for out-of-distribution evaluation on MxT-Bench shown in Table 1). The agents are diversified missing, mass, size randomization.
Method (data size)
Compositional (Morphology)
Compositional (Task)
Out-of-Distribation
MI.P (4K, randinit)
0.8162 + 0.02
0.7046 + 0.01
0.6674 + 0.03
MLP (4K, fine-tuning)
0.6715 : 0.02
0.5582 = 0.03
0.5694 = 0.04
Transformed (MTGv1) 1) (4K, randinit)
0,6021 1 0.03
0.5017 1 0.04
0.5222 1 0.06
Transformer (MTGv1) (4K, fine-tuning)
0,3499 + 0.04
0.3378 + 0.08
0,3401 = 0.08
Transformer (MTGv2) (4K, radinitt)
0,5504 + 0.03
0.4894 + 0.03
0.4980 + 0.05
Transformer (MTGv2) (4K, fine-tuning)
0.2863 + 0.02
0.3180 + 0.07
0,3115 + 0.07
MI.P (8K, randinit)
0.7044 = 0.02
0.4818 + 0.02
0.5247 + 0.04
MI.P (8K, fine-tuning)
0.5579 + 0.06
0.3709 + 0.01
0,4238 + 0.05
Transformer (MTGv1) (8K, randinit)
0.4001 = 0.04
0.2791 + 0.06
0,2762 = 0.04
Trameformer (MTGv1) (8K, fine-tuning)
0,2600 1 0.04
0.1926 1 0.03
0,2089 1 0.03
Transformer (MTGv2) (8K, randinit)
0.3758 + 0.03
0.1868 - 0.05
0.1868 + 0.05
Transformer (MTGv2) (8K, fine-tuning)
0.1871 + 0.01
0.1356 0.03
0.1498 + 0.03
MLP (12K, randinit)
0.5530 t 0.01
0.4256 1 0.03
0,4399 1 0.02
MLP (12K, finc-tuning)
0,4312 + 0.03
0.2982 - 0.02
0,3358 + 0.02
Transformer (MTGv1) (12K, randinit)
0,2914 1 0.03
0.1555 1 0.04
0.1759 1 0.03
Transformer (MTGv!) I) (12K, fine-tuning)
0.2042 + 0.01
0.0888 + 0.01
0.0894 + 0.01
Transformer (MTGv2) (12K, randinit)
0.2655 + 0.03
0.1149 0.02
0.1149 + 0.02
Transformer (MTGv2) (12K, fine-tuning)
0.1301 0.02
0.0562
0.0513 1 0.01
Table 16: The average normalized final distance among test environments in fine-tuning settings (compositional morphologytask or outto--distiibution evaluation).
30

---

Published as a conference paper at ICLR 2023
K WHY MORPHOLOGY-TASK GRAPH WORKS WELL?: ATTENTION ANALYSIS
In this section, we provide the attention analysis of Transformer with MTGv1 (Figure 13) and MTGv2 (Figure 12).
The experimental results reveal that despite slight differences, MTGv2 generalizes various morpholo- gies and tasks better than MTGv1. To find out the difference between those, qualitatively analyze the attention weights in Transformer. Figure 12 shows that MTGv2 consistently focuses on goal nodes over time, and activates important nodes solve the task; for instance, in centipede touch (top), MTGv2 pays attention to corresponding nodes (torsof) and goal0) at the beginning of the episode, and gradually sees other relevant nodes (torsol and torso2) to hold the movable ball. Furthermore, in ant twisters (bottom), MTGv2 firstly tries to raise the agent's 's legs to satisfy goall and goal2, and then focus on reaching a leg (goal0). Temporally-consisteen attention to goal nodes and dynamics attention to relevant nodes can contribute to generalization over goal-directed tasks and morphologies.
Figure 13 implies that MTGvl does not show such consistent activation to the goal-conditioned node; for instance, in centipede touch 3 (above), the goal information is treated feature of torso0, but there are no nodes that consistently activated with torso0. Morcover, in ant reach_ @each_handsup2. episode. Rather, MTGv1 tends to demonstrate some periodic patterns during the rollout as implied in prior works (Huang al., 2020; Kurin et al., 2021).
extra node
(bottom), MTGv1 does not keep focusing the agent's limbs during the
centipede_ 3
0.8
0.6
0.4
0.2
0.0
Time step ant_ reacre Saaachhhandsup2555
0.6
0.4
0.2
3000
0.0
Time step
Figure 12: Attention analysis of MTGx2 in centipede_toous touch (top) and reachh Sandd_pands (bottom; from twisters). From left to right, we visualize the attention weights of MTGv2 during the rollout. In contrast to MTGv1, MTGv2 consistently focuses goal nodes over time, and activates important nodes to solve the task.
31

---

Published as a conference paper at ICLR 2023
centipede_tous touch__
0.8
0.6
0.4
0.2
0.0
Time step p_.recchhhandssp..
bondd
Ims3
inc
krto
Impe
Int
Im:5
Inst
Lock
0.2
krto
Inch
1100
Enoo
Hed
knod
0.0
Time step
Figure 13: Attention analysis of MTGv1 in contipede_toucht and ant_reach reach_handsuu2_111 This tends to
demonstrate some periodic patterns during the rollout implied in prior works (Huang et al., 2020; Kurin et al., 2021).
32

---

Published as a conference paper at ICLR 2023
L ADDITIONAL RESULTS OF TASK GENERALIZATION
Although in Table and Figure 5, we examine the compositional task generalization and out-of- distribution generalization with reach hard tasks, where the goal distribution is farther than original reach tasks. While they seem more difficult unseen tasks (Furuta et al., 2021b), they also seem have some sort of task similarities between training dataset environments and those evaluation environments.
Another question might be how morphology-tas graph performs in more different unseen tasks. To evaluate compositional task generalization and -itt----istriuttion on the environments that have less similarity to the training datasets, we prepare push task, where the agents try move the box objects to the given goal position. See Table 7 for the environment division. For training datasets, we leverage In-Distribution division in Table 4. Because this task requires the sufficient interaction with the object, the nature of tasks seem quite different from training dataset environments (reach, touch, and twisters).
Table 17 shows the results of compositional task evaluation and Table 18 shows those of out-of- distribution evaluation. In contrast to Table I, the zero-shot performance seems limited. Transferring pre-trained control primitives significantly different tasks still remains as important future work.
However, as shown in Figure 14, morphology-tase graph works as better prior knowledge for downstream multi-tas task imitation learning, with the environments that have less similarity the pre-training datasets. As prior work suggested (Mandi et al., 2022), these results suggests that, in RL, jointly-learned multi-task model has strong inductive bias even for unseen and significantly different environments.
Sub-domain
Random
MI.P
Transformer (MTGv1)
Transformer (MTGv2)
ans_ centipede_ push push
10808 0.02
0.9995 +: 0.00
0.9212 0.14
0.8836 + 0.09
0.02
0.9377
womm_push
1.0026 0.04
0.6300 0.01
0.5579 0.10
0.4091 + 0.15
Average Dist.
1.0866 0.02 0.9051 0.00
0.8800 0.02
0.8263 +: 0.04
Table The average normalized final distance compositional task evaluation on MxT-Bench with unseen push task. See Table for the environment division.
Sub-domain
Random
MLP
Transformer (MTGv1)
Transformer (MTGv2)
aot_pushh_dvee diverse centipede puse_diverse
1.0814
0.02
1.0018 0.01
0.9906 0.01
0.8800 +
diverse
1.0848 0.02
1.0031 0.01
0.9997 0.00
0.9355 + 0.04
Average Dist.
1.0835 0.02
1.0021 0.02
0.9965 0.01
0.9155 + 0.03
Table 18: The average normalized final distance for out-of-distribution evaluation on MxT-Bench with unseen push task. See Table 7 for the environment division.
33

---

Published as a conference paper at ICLR 2023
Compositional (Task)
Out-of-Distribution
MLP (4k)
MLP (8k)
MLP (12k)
TF-MTGv1 (4k)
TF-MTGV1 (8k)
TF-MTGv1 (12k)
TF-MTGv2 (4k)
TF-MTGv2 (8k)
TF-MTGv2 (12k)
fireturing randint
finetrnig
randint
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
Average Normalized Distance
Average Normalized Distance
Figure 14:
Comparison of multi-tast goal-reaching performances on fine-tuning settings in the unseen push tasks. The results imply that morphology-task graph works as better prior knowledge for downstream multi-task imitation learning, even with the environments that have less similarity to the pre-training datasets.
Method (data size)
Compositional (Task)
Out-of-Distributiin
MLP (4K, randinit)
0.8042 0.05
0,8332 0.04
MLP (4K, fine-tuning)
0.7818 0.05
0.8108 + 0.02
Transformer (MTGvl) I) (4K, randinit)
0.6653 0.03
0.7191 + 0.02
Transformer (MTGvl) (4K, fine-tuning)
0.6398 0.05
0,6893 = 0.04
Transformer (MTGv2) (4K, randinit)
0.6276 1 0.05
0,6999 1 0.04
Transformer (MTGv2) (4K, fine-tuning)
0.5826 + 0.05
0.6779 + 0.03
MLP (8K, randinit)
0.7542 0.06
0.7751 = 0.03
MLP (8K, fine-tuning)
0.6847 0.04
0,7611 1 0.03
Transformer (MTGvl) (8K, randinit)
0.6027 0.03
0,6613 + 0.05
Transformer (MTGv1) (8K. fine-tuning)
0.5448 0.06
0,6169 - 0.05
Transformer (MTGv2) (8K, randinit)
0.5704 + 0.03
0.6286 1 0.0S
Transformer (MTGv2) (8K, fine-tuning)
0.4993 = 0.06
0,5772 + 0.06
MLP (12K, randinit)
0.7037 0.03
0.7452 + 0.03
MLP (12K, fine-tuning)
0.6678 0.04
0,7110 = 0.06
Transformer (MTGvl) (12K, randinit)
0.5925 1 0.04
0,6203 1 0.04
Transformer (MTGvl) (12K, fine-tuning)
0.4993 + 0.06
0,5158 + 0.08
Transformer (MTGx2) (12K, randinit)
0.5510 + 0.05
0.5454 1 0.11
Transformer (MTGv2) (12K, fine-tuning)
0.4132 0.04
0,4748 1 0.06
Table 19: The average normalized final distance among environments in fine-tuning settings (compositional morphology/tssk or out-o--distribution evaluation) with unseen push task.
M EXTENDED RELATED WORK
Morphology-Tane -ask Generalization In the previous literature, multi-task in RL has arisen from each MDP component; from different (1) dynamics (Hallak et al., 2015; Yu et al., 2019), (2) reward (Kaelbling, 1993; Andrychowiez et al., 2017), and (3) state/action space (Wang al., 2018; Huang et al., 2020). We think multi-morpholggy RI. covers (1) and (3), and multi-task (or multi-goal) RL does (2) and (1). Considering morphology-tase generalization, we could study the generalization in RL across all the MDP components (dynamics, reward, state space, action space). While, from a broader perspective, the notion of multi-task in RL may contain both morphological and goal diversity, in this paper, treat them separately, and multi-task just stands for multi-goa -gall settings.
Connection to Policy Distillation Policy distillation has been proposed and studied for a while (Rusu al., 2016; Parisotto al., 2016; Levine et al., 2016; Czarnecki et al., 2019), where single student policy distills the knowledge from multiple teacher policies to obtain better multitask learners. In contrast, we call our protocol behavior distillation, where the single student policy distills the knowledge from multi-source offline data (e.g. human teleoperation, scripted behavior (Singh et al., 2021), 1), play data (Lynch et al., 2019)), not limited to RL policy, since recent RL robotics research often leverage such a data-driven approach for scalability (Levine et al., 2020; Chen et al., 202 Gu al., 2021a; Lee et al., 2022; Reed et al., 2022; Zeng al., 2020; Shridhar al., 2022b).
34


Reviews:
[{"review_id": "N2aZWlBvnJ", "reviewer": "Reviewer_JKD7", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors created an benchmark called \"MxT-Bench\", which study the cross product of morphology and task generalization. Moreover, the authors proposed a unified graph-based representation for multi-morphology and multi-task. The authors provide a wide range of experiments and showcase the benchmark and models.", "summary_of_the_review": "I like this paper and think this is above acceptance threshold. However, I'm waiting to see other reviewers's feedback (who probably are more familiar with the field than I am), before giving it a very strong endorsement.", "strength_and_weaknesses": "I want to preface this by saying that I'm not really familiar with the field of controls across morphologies. \n\nStrength:\n- MxT bench is well designed and seems to be technically sound. This can be of interest to many different researchers.\n- The experiment study is quite thorough and offers some interesting insights. \n\nWeakness:\n- I feel like the analysis in 5.3 is not so strong. Especially figure 6, where I don't see much to be honest. \n- I wonder if a comparison of models learned across benchmarks can be useful. For example, does the cross morphology models learned in other benchmarks perform well in MxT?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JM_3UGtEmYs", "reviewer": "Reviewer_UDt3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose control graphs as a way to train agents with different morphologies and across different tasks. They develop MxT-Bench to train with different agent morphologies and across different tasks. They show a set of offline experiments that demonstrate the usefulness of their control graphs over MLPs and GNNs.", "summary_of_the_review": "The work is well-presented and easy to follow, and the direction of the work is quite fascinating. However, I'm concerned about the lack of demonstrated practical usefulness of control graphs. The experiments were only on relatively simple tasks with basic mujoco-styled agents. I would have liked to see more experiments where the agent requires temporal memory and with more real-world robotic embodiments. As the work stands now, it's hard for me to see the usefulness of using control graphs beyond pretty simple setups.", "strength_and_weaknesses": "Strengths:\n- The paper is well presented and the motivation is very clear.\n- The research question of designing architectures that generalize to different agent embodiments and tasks is very interesting, practically useful, and well-discussed.\n- The idea for control graphs and setting them up to be used in a transformer is pretty neat!\n- The analysis of the attention weights is well written and done.\n\nWeaknesses:\n- W1. The agent morphology and task setups are really limiting (agents being basic variants of ant, centipede, and claw agents, tasks being very limited to reach, touch, twister, and push tasks). It would be hard for me to accept the usefulness of control graphs without much more experimentation with different robotic setups and tasks. I would have liked to see something along the lines of trying control graphs with respect to parameterizable simulated robots and with respect to more real-world tasks. It doesn't seem like this work \n- W2. The work only deals with optimizing offline datasets.\n- W3. It appears that the Transformer CGv1 and CGv2 do not encode any temporal information, that is, there is no information flow from the previous time steps to the current time step, which would make the current formulation not work well in any tasks that require memory.\n- W4. In Table 1, performance for in-distribution seems lower than I would have expected. I would have expected a transformer to practically be able to overfit with near perfect training performance. Why isn't this happening?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "I7OgX_qzij3", "reviewer": "Reviewer_feBA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a continuous control benchmark for multitask reinforcement learning supporting three axes of change: creature morphology (supporting state and action spaces of different dimensions), task (e.g. reach some place or touch some point), and goal.\nThe paper further develops the graph-based representation pioneered by NerveNet, adding the support of the goal-directed setting. Finally, the authors investigate policy distillation behaviour on their own benchmark", "summary_of_the_review": "I think the paper has interesting findings and has the potential, but issues with the positioning and clarity I pointed out made me choose the score of 5. I am ready to increase the score after the discussion period if my concerns are addressed. \n\nUPD after rebuttal: The authors were actively engaged in the discussion and updated the paper to address my concerns. I would give the paper 7, but the scoring allows only 6 or 8, and I will update the score to 8.", "strength_and_weaknesses": "## Strengths\n\n* The proposed benchmark is a large benchmark supporting a variety of degrees of freedom: morphology, different mass/size for limbs, multiple tasks (using the authors' terminology), goal distribution.\n* The paper investigates policy distillation in the incompatible setting (mismatching state and action space dimensions across tasks) which, I don't think, was done before.\n* The benchmark seems to be super fast with single-agent PPO taking 5-30 minutes per environment\n\n## Weaknesses\n\nI think the positioning of the paper paper is its main weakness. I am not entirely sure what this paper is about. There are three potential storylines:\n\n* The paper is about the benchmark \n* The paper is a policy distillation study in a novel setting (mismatching state/action spaces across agents)\n* The paper is about the input representation\n\nI think, all three stories are equally legit for this paper, but each of the story will define the requirements for experiments, and jumping around the stories will confuse the reader.\n\nFor instance, the title of the paper says that the paper is about control graph, i.e. input representation. However, the intro's motivational section speaks about the generalist agent and morphology generalisation. If the paper is about morphology generalisation, then Representation Selection, which is Appendix F now, should be in the main body of the paper etc. I think, the main body can accommodate more information if the transformer basics are moved to the Appendix.\n\nI'm very confused about the terminology used throughout the paper: 'multi-task' is used in a very narrow meaning implying that the agents of the same morphology doing semantically different tasks (e.g. reaching some goal or moving limbs to specific positions). Usually, In RL, a task is a synonym of an MDP, and different morphology (even same morphology with different limb sizes/masses) constitute different tasks. As you mentioned, previous literature (e.g. Gupta et al., 2022, Metamorph) consider morphology/your definition of task separately and that 'broad generalisation over task and morphology jointly remains a long-standing problem'. Do you claim that there is a qualitative difference between generalising across morphologies and generalising across your definition of tasks? In other words, what is the motivation behind separating multitask into two separate groups?\n\nFinally, I find it concerning, that the paper that does policy distillation does not cite policy distillation papers, including Parisotto's \"Actor-Mimic\" and Rusu's \"Policy Distillation\". The \"Behaviour Distillation\" section includes only 2020/21 papers with a single exception of Levine et al 2016. Task generalisation with GNN/Transformers work should include papers on combinatorial optimisation (see Cappart's 'Combinatorial optimization and reasoning with graph neural networks' for references). GNNs for continuous control work misses Blake's \"Snowflake: Scaling GNNs to High-Dimensional Continuous Control via Parameter Freezing\".\n\nOther questions:\n* I am not sure I fully understand how the goals are encoded to either control graph v1 or v2. In v1, when you add the goal observation to the node features, are other nodes' features padded? Do you provide the task identifier somehow (e.g. if that's reach or touch or twisters). In v2, figure three shows the dashed edges from the limb nodes to goal nodes, how do you encode this information in the transformer case? Also, judging by Figure 3, the dimensionalities of the goal nodes for different tasks are different. How do you ensure that they are all the same before going to the GNN/transformer?\n* When you generate tasks, are they all solvable (whatever metric you choose for this). Can you provide plots with the return ranges for single-agent policies you use to collect the data? How did you select thresholds for goal-reaching?\n* Could you elaborate a bit more on the learnable position embeddings, how exactly is this implemented. When zero-shot transferred, how are position embeddings calculated?\n* How consistent is the behaviour in Figure 6 across the seeds?\n* What do you think the wider implications of the proposed input representation are? Can we draw any conclusions that can be used outside the continuous control setting?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0leXd9OQBvX", "reviewer": "Reviewer_zDig", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on multi-morphology and multi-task generalization in robotics. There are three main contributions to this paper. First, a comprehensive benchmark called MxT-Bench is developed for training and evaluating morphology-task generalization. This benchmark also supports the scalable procedural generation of both agents and tasks. Second, a new concept Control Graph is proposed as a universal interface representing observations, actions, and goals/tasks. Last, the authors study several offline supervised behavior distillation methods with different design choices and obtain conclusions about generalizable model design.\n\nCompared to existing works of universal controller and morphology/controller co-design, this paper provides a more comprehensive and systematic design, which has the potential to unify the tools and standards used in this area. \n", "summary_of_the_review": "I suggest accepting this paper. The morphology-task generalization topic becomes a prevalent trend recently. I believe the benchmark and unified interface proposed in this paper contribute a lot to this area.", "strength_and_weaknesses": "### **Strength**\n\n* **A systematic platform.** Universal controllers for different robots and co-design of morphology and controller are prevalent topics in robotics. Different papers propose many different simulators and protocols, making the development and comparison between algorithms difficult. The benchmark proposed by this paper looks promising to me. I think it will greatly contribute to this area by unifying the tools and standards used in it.\n\n* **Thorough analysis of experiments.** Beside the benchmark, this paper also explores multiple generalization settings, including multi-morphology, multi-task, and out-of-distribution downstream tasks. Some interesting results will definitely provide guidance for future work.\n\n\n### **Weaknesses**\n\nOne weakness in my opinion is that an overview of the structure of this paper is missing. Since this paper proposes many new concepts, understanding the connections between them could be difficult.\nFigure 1 tries to summarize all contents but still needs to be improved. For example, section numbers can be added to different components and main contributions can be highlighted.\nThe title can also be improved. I think the Control Graph is not the main contribution. A more precise title helps understand the content.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%