Original: C ONTROL G RAPH AS U NIFIED IO FOR M ORPHOLOGY -T ASK G ENERALIZATION
Revised: Control Graph as Unified IO for Morphology-Task Generalization

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: Abstract

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In order to align input-output (IO) interface among multiple tasks and diverse agent morphologies while preserving essential 3D geometric relations, we introduce control graph , which treats observations, actions and goals/task in a unified graph representation.
Revised: In order to align input-output (IO) interface among multiple tasks and diverse agent morphologies while preserving essential 3D geometric relations, we introduce control graph , which treats observations, actions and goals/task in a unified graph representation.

Original: 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.
Revised: 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.

Original: Through efficient representation and architecture selection on MxT-Bench, we find out that a control 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.
Revised: Through efficient representation and architecture selection on MxT-Bench, we find out that a control 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.

Original: 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.
Revised: 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 1 .

Original:  I NTRODUCTION
Revised:  Introduction

Original: The impressive success of large language models (Devlin et al., 2019; Radford et al., 2019; Bommasani et al., 2021; Brown et al., 2020; Chowdhery et al., 2022) has encouraged the other domains, such as computer vision (Radford et al., 2021; Gu et al., 2021b; Alayrac et al., 2022; Jaegle et 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.
Revised: The impressive success of large language models [25, 84, 8, 10, 19] has encouraged the other domains, such as computer vision [85, 42, 3, 52] or robotics [1, 51], to leverage the large-scale pre-trained model trained with massive data with unified input-output interface.

Original: 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 et al., 2020; Chen et al., 2022a).
Revised: These large-scale pre-trained models are innately multi-task learners: they surprisingly work well not only in the fine-tuning or fewshot transfer but also in the zero-shot transfer settings [86, 15].

Original: Learning a “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.
Revised: Learning a “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.

Original: In reinforcement learning (RL) for continuous control, various aspects are important for generalization.
Revised: In reinforcement learning (RL) for continuous control, various aspects are important for generalization.

Original: First, we care about “task” generalization.
Revised: First, we care about “task” generalization.

Original: 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 et al., 2019).
Revised: For instance, in robotic manipulation, we care the policy to generalize for different objects and target goal positions [58, 4, 103, 72].

Original: 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).
Revised: Recent advances in vision and language models also enable task generalization through compositional natural language instructions [54, 91, 1, 23].

Original: However, to scale the data, equally important is “morphology” generalization, where a single policy can control agents of different embodiment (Wang et al., 2018; Noguchi et al., 2021) and can thereby ingest experiences from as many robots in different simulators (Freeman et al., 2021; Todorov et al., 2012; Coumans & Bai, 2016) as possible.
Revised: However, to scale the data, equally important is “morphology” generalization, where a single policy can control agents of different embodiment [100, 76] and can thereby ingest experiences from as many robots in different simulators [31, 96, 22] as possible.

Original: Most  prior works (Mendonca et 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 a long-standing problem.��������� ��������������������� ����������������������������� ��������� �������������� ��������������������� ������������������������������� �������������������� ��������������������� �������������������������� �����������
Revised: Most ��������� ��������������������� ����������������������������� ��������� �������������� ��������������������� ������������������������������� �������������������� ��������������������� �������������������������� ����������� prior works [75, 44] only address either the task or morphology axis separately, and achieving broad generalization over task and morphology jointly remains a long-standing problem.

Original: This paper first proposes MxT-Bench 1 , the first multi-morphology and multi-task benchmarking environments, as a step toward building the massive diverse dataset for continuous control.
Revised: This paper first proposes MxT-Bench 2 , the first multi-morphology and multi-task benchmarking environments, as a step toward building the massive diverse dataset for continuous control.

Original: MxTBench provides various combinations of different morphologies (ant, centipede, claw, worm, and unimal (Gupta et al., 2022)) and different tasks (reach, touch, and twisters).
Revised: MxTBench provides various combinations of different morphologies (ant, centipede, claw, worm, and unimal [44]) and different tasks (reach, touch, and twisters).

Original: MxT-Bench is easily scalable to additional morphologies and tasks, and is built on top of Brax (Freeman et al., 2021) for fast behavior generation.
Revised: MxT-Bench is easily scalable to additional morphologies and tasks, and is built on top of Brax [31] for fast behavior generation.

Original: Next, we define unified IO representation for an architecture to ingest all the multi-morphology multi-task data.
Revised: Next, we define unified IO representation for an architecture to ingest all the multi-morphology multi-task data.

Original: Inspired by scene graph (Johnson et al., 2015) in computer vision that represents the 3D relational information of a scene, and by morphology graph (Wang et al., 2018; Chen et al., 2018; Huang et al., 2020; Gupta et al., 2022) that express an agent’s geometry and actions, we introduce the notion of control graph (CG) as an unified interface to encode observations, actions, and goals (i.e. tasks) as nodes in the shared graph representation.
Revised: Inspired by scene graph [55] in computer vision that represents the 3D relational information of a scene, and by morphology graph [100, 13, 49, 44] that expresses an agent’s geometry and actions, we introduce the notion of control graph (CG) as a unified interface to encode observations, actions, and goals (i.e. tasks) as nodes in the shared graph representation.

Original: 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).
Revised: 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).

Original: 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 et al., 2018; Ghasemipour et al., 2022), and we empirically shows it outperforms naive tokenization in our control-focused dataset.
Revised: In contrast to discretizing and tokenizing every dimension as proposed in recent work [53, 87], this unified IO limits data representation it can ingest, but strongly preserves 3D geometric relationships that are crucial for any physics control problem [100, 38], and we empirically shows it outperforms naive tokenization in our control-focused dataset.

Original: 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 8 local node observations, for optimal generalization through behavior distillation (Figure 1), where RL is essentially treated as a (single-task, low-dimensional) behavior generator (Gu et al., 2021a) 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).
Revised: Lastly, while conventional multi-task or meta RL studies generalization through on-policy joint training [103, 20], we perform efficient representation and architecture selection , over 11 combinations of unified IO representation and network architectures, and 8 local node observations, for optimal generalization through behavior distillation (Figure 1), where RL is essentially treated as a (single-task, low-dimensional) behavior generator [41] and multi-task supervised learning (or offline RL [33]) is used for imitating all the behaviors [93, 14, 87].

Original: Through offline distillation, we controllably and tractably evaluate two variants of CG representation, along with multiple network architectures (MLP, GNN (Kipf & Welling, 2017), Transformers (Vaswani et al., 2017)), and show that CGv2 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 %).
Revised: Through offline distillation, we controllably and tractably evaluate two variants of CG representation, along with multiple network architectures (MLP, GNN [60], Transformers [98]), and show that CGv2 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 %).

Original: 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 to continue growing diverse behavior datasets, designing different IO representations, and iterating more representation and architecture selection, and eventually optimize a single policy that can be deployed on any morphology for any task.
Revised: As the fields of vision and language move toward broad generalization [18, 8], we hope our work could encourage RL and continuous control communities to continue growing diverse behavior datasets, designing different IO representations, and iterating more representation and architecture selection, and eventually optimize a single policy that can be deployed on any morphology for any task.

Original: In summary, our key contributions are:
Revised: In summary, our key contributions are:

Original: • We develop MxT-Bench 2 as a test bed for morphology-task generalization with fast expert behavior generator.
Revised: • We develop MxT-Bench 3 as a test bed for morphology-task generalization with fast expert behavior generator.

Original: MxT-Bench supports the scalable procedural generation of both agents and tasks with minimal blueprints.
Revised: MxT-Bench supports the scalable procedural generation of both agents and tasks with minimal blueprints.

Original: 1 Pronounced as “mixed”-bench.

Original: It stands for “Morphology × Task”.

Original: 2 https:

Original: • We introduce control graph , a universal IO for control which treats the agent’s observations, actions and goals/tasks in a unified graph representation, while preserving the task structure.
Revised: • We introduce control graph , a universal IO for control which treats the agent’s observations, actions and goals/tasks in a unified graph representation, while preserving the task structure.

Original: • We study generalization through offline supervised behavior distillation, where we can efficiently try out various design choices; over 11 combinations of unified IO representation and network architectures, and 8 local node observations.
Revised: • We study generalization through offline supervised behavior distillation, where we can efficiently try out various design choices; over 11 combinations of unified IO representation and network architectures, and 8 local node observations.

Original: As a result, we find that Transformer with CGv2 achieves the best multi-task performances among other possible designs (MLP, GNN and Transformer with CGv1, Tokenized-CGv2, etc.)
Revised: As a result, we find that Transformer with CGv2 achieves the best multi-task performances among other possible designs (MLP, GNN and Transformer with CGv1, Tokenized-CGv2, etc.)

Original: in both in-distribution and downstream tasks, such as zero-shot transfer and fine-tuning for multi-task imitation learning.
Revised: in both in-distribution and downstream tasks, such as zero-shot transfer and fine-tuning for multi-task imitation learning

Original:  R ELATED W ORK
Revised:  Related Work

Original: ������������������������� ������������������� ��� ������������������������������ �
Revised: ������������������������� ������������������� ��� ������������������������������ �

Original: 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-aware Graph Neural Network (GNN) policies (Wang et al., 2018; Huang et al., 2020).
Revised: Morphology Generalization While, in RL for continuous control, the policy typically learns to control only a single morphology [94, 96], several works succeed in generalizing the control problem for morphologically different agents to solve a locomotion task by using morphologyaware Graph Neural Network (GNN) policies [100, 49].

Original: In addition, several work (Kurin et al., 2021; Gupta et al., 2022; Hong et al., 2022; Trabucco et al., 2022) have investigated the use ofthe Transformer (Vaswani et al., 2017).
Revised: In addition, several work [62, 44, 48, 97] have investigated the use of Transformer [98].

Original: Other work jointly optimize the morphology-agnostic policy and morphology itself (Pathak et 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).
Revised: Other work jointly optimize the morphology-agnostic policy and morphology itself [80, 43, 105, 47], or transfer a controller over different morphologies [24, 13, 46, 71].

Original: 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.
Revised: 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.

Original: We believe that goal-directed control is a key problem for an embodied single controller.
Revised: We believe that goal-directed control is a key problem for an embodied single controller.

Original: Concurrently, Feng et al.
Revised: Concurrently, Feng et al.

Original: (2022) propose an RL-based single controller that is applied to different quadruped robots and target poses in the sim-to-real setting.
Revised: [30] propose a RL-based single controller that is applied to different quadruped robots and target poses in the sim-to-real setting.

Original: In contrast, our work introduces the notion of control graph as a unified IO that represents observations, actions, and goals in a shared graph, and can handle more diverse morphologies to solve multiple tasks.
Revised: In contrast, our work introduces the notion of control graph as a unified IO that represents observations, actions, and goals in a shared graph, and can handle more diverse morphologies to solve multiple tasks.

Original: Task Generalization
Revised: Task Generalization

Original: In the previous research, task generalization has been explored in multi-task or meta RL literature (Wang et al., 2016; Duan et 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).
Revised: In the previous research, task generalization has been explored in multitask or meta RL literature [99, 27, 11, 95, 21, 67, 102, 63].

Original: Each task might be defined by the difference in goals, reward functions, and dynamics (Ghasemipour et al., 2019; Kalashnikov et al., 2021; Eysenbach et al., 2020), under shared state and action spaces.
Revised: Each task might be defined by the difference in goals, reward functions, and dynamics [37, 59, 28], under shared state and action spaces.

Original: 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).
Revised: Some works leverage graph representation to embed the compositionality of manipulation tasks [68, 108, 69, 61, 38], while others use natural language representation to specify diverse tasks [54, 91, 1, 50, 23].

Original: Despite the notable success of acquiring various task generalization, multi-task RL often deals with only a single morphology.
Revised: Despite the notable success of acquiring various task generalization, multitask RL often deals with only a single morphology.

Original: 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.
Revised: 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.

Original: Transformer for RL Recently, Chen et al.
Revised: Transformer for RL Recently, Chen et al.

Original: (2021a) and Janner et al.
Revised: [12] and Janner et al.

Original: (2021) consider offline RL as supervised sequential modeling problem and following works achieve impressive success (Reed et al., 2022; Lee et al., 2022; Furuta et al., 2022; Xu et al., 2022; Shafiullah et al., 2022; Zheng et al., 2022; Paster et al., 2022).
Revised: [53] consider offline RL as supervised sequential modeling problem and following works achieve impressive success [87, 64, 36, 101, 90, 107, 79].

Original: In contrast, our work leverages Transformer to handle topological and geometric information of the scene, rather than a sequential nature of the agent trajectory.
Revised: In contrast, our work leverages Transformer to handle topological and geometric information of the scene, rather than a sequential nature of the agent trajectory.

Original: Behavior Distillation Due to the massive try-and-error and large variance, training a policy from scratch in online RL is an inefficient process, especially in multi-task setting.
Revised: Behavior Distillation Due to the massive try-and-error and large variance, training a policy from scratch in online RL is an inefficient process, especially in multi-task setting.

Original: It is more efficient to use RL for generating single-task behaviors (often from low dimensions)
Revised: It is more efficient to use RL for generating single-task behaviors (often from low dimensions)

Original: (Gu et al., 2021a) and then use supervised learning to imitate all behaviors with a large single policy (Levine et al., 2016; Singh et al., 2021; Ajay et al., 2021; Chen et al., 2021b).
Revised: [41] and then use supervised learning to imitate all behaviors with a large single policy [65, 93, 2, 14].

Original: Several works tackle the large-scale behavior����������� ���������� ������������� ������ ����� ����� �������� ����������� ����� ����������� ��� �� �� distillation with Transformer architecture (Reed et al., 2022; Lee et al., 2022), or with representation that treats observations and actions in the same vision-language space (Zeng et al., 2020; Shridhar et al., 2022b).
Revised: Several works tackle the large-scale behavior distillation with Transformer architecture [87, 64], or with representation that treats observations and actions in the same vision-language space [106, 92]. Our work utilizes [Table]

Original: Our work utilizes similar pipeline, but focuses on finding the good representation and architecture to generalize across morphology and tasks simultaneously with proposed control graph.
Revised: ����������� ���������� ������������� ������ ����� ����� �������� ����������� ����� ����������� ��� �� �� similar pipeline, but focuses on finding the good representation and architecture to generalize across morphology and tasks simultaneously with proposed control graph.

Original:  P RELIMINARIES
Revised:  Preliminaries

Original: In RL, consider a Markov Decision Process with following tuple ( S , A , p , p 1 , r , γ ), which consists of state space S , action space A , state transition probability function p : S ×
Revised: In RL, consider a Markov Decision Process with following tuple ( S , A , p , p 1 , r , γ ), which consists of state space S , action space A , state transition probability function p : S ×

Original: A × S → [0 , ∞ ) , initial state distribution p 1 : S → [0 , ∞ ) , reward function r : S ×
Revised: A × S → [0 , ∞ ) , initial state distribution p 1 : S → [0 , ∞ ) , reward function r : S ×

Original: A → R , and discount factor γ ∈ [0 , 1) .
Revised: A → R , and discount factor γ ∈ [0 , 1) .

Original: Deep RL parameterizes a Markovian policy π : S → A with neural networks.
Revised: Deep RL parameterizes a Markovian policy π : S → A with neural networks.

Original: The agent seeks optimal policy π ∗ that maximizes the discounted cumulative rewards:
Revised: The agent seeks optimal policy π ∗ that maximizes the discounted cumulative rewards:

Original: [Equation]
Revised: [Equation]

Original: )] , (1) where p πt ( s t ) =
Revised: (1) where p πt ( s t ) =

Original: (cid:82)(cid:82) s 0: t ,a 0: t − 1 (cid:81)
Revised: (cid:82)(cid:82) s 0: t ,a 0: t − 1 (cid:81)

Original: t p ( s t | s t − 1 , a t − 1 ) π ( a t | s t ) and ρ π ( s ) =
Revised: t p ( s t | s t − 1 , a t − 1 ) π ( a t | s t ) and ρ π ( s ) =

Original: (1 − γ ) (cid:80) t γ t p πt ( s t = s ) represent time-aligned and time-aggregated state marginal distributions following policy π .
Revised: (1 − γ ) (cid:80) t γ t p πt ( s t = s ) represent time-aligned and time-aggregated state marginal distributions following policy π .

Original: 3. G RAPH R EPRESENTATION FOR M ORPHOLOGY -A GNOSTIC C ONTROL
Revised: 3.1 Graph Representation for Morphology-Agnostic Control

Original: Following prior continuous control literature (Todorov et al., 2012), we assume the agents have bodies modeled as simplified skeletons of animals.
Revised: Following prior continuous control literature [96], we assume the agents have bodies modeled as simplified skeletons of animals.

Original: An agent’s morphology is characterized by the parameter for rigid body module (torso, limbs), such as radius, length, mass, and inertia, and by the connection among those modules (joints).
Revised: An agent’s morphology is characterized by the parameter for rigid body module (torso, limbs), such as radius, length, mass, and inertia, and by the connection among those modules (joints).

Original: In order to handle such geometric and topological information, an agent’s morphology can be expressed as an acyclic tree graph representation G :
Revised: In order to handle such geometric and topological information, an agent’s morphology can be expressed as an acyclic tree graph representation G :

Original: = ( V , E ) , where V is a set of nodes v i ∈ V and E is a set of edges e ij ∈ E between v i and v j .
Revised: = ( V , E ) , where V is a set of nodes v i ∈ V and E is a set of edges e ij ∈ E between v i and v j .

Original: The node v i corresponds to i -th module of the agent, and the edge e ij corresponds to the hinge joint between the nodes v i and v j .
Revised: The node v i corresponds to i -th module of the agent, and the edge e ij corresponds to the hinge joint between the nodes v i and v j .

Original: Each joint may have 1-3 actuators corresponding to a degree of freedom.
Revised: Each joint may have 1-3 actuators corresponding to a degree of freedom.

Original: If a joint has several actuators, the graph G is considered a multipath graph.
Revised: If a joint has several actuators, the graph G is considered a multipath graph.

Original: This graph-based formulation can describe the various agents’ morphologies in a tractable manner (Wang et al., 2018; Huang et al., 2020).
Revised: This graph-based formulation can describe the various agents’ morphologies in a tractable manner [100, 49].

Original: We assume that node v
Revised: We assume that node v

Original: i observes local sensory input s it at time step t , which includes the information of limb i such as position, velocity, orientation, joint angle, or morphological parameters.
Revised: i observes local sensory input s it at time step t , which includes the information of limb i such as position, velocity, orientation, joint angle, or morphological parameters.

Original: To process these node features and graph structure, a morphology-agnostic policy can be modeled as node-based
Revised: To process these node features and graph structure, a morphology-agnostic policy can be modeled as node-based

Original: GNN (Kipf & Welling, 2017; Battaglia et al., 2018), which takes a set of local observations { s it } | V | 
Revised: GNN [60, 6], which takes a set of local observations { s it } | V | i

Original: i = as an input and emit the actions for actuators of each joint { a et } | E | e =1 .
Revised: =1 as an input and emit the actions for actuators of each joint { a et } | E | e =1 .

Original: The objective function of morphology-agnostic RL is the average of Equation 1 among given morphologies.
Revised: The objective function of morphology-agnostic RL is the average of Equation 1 among given morphologies.

Original: �������������� ���������� ������� ��������� ���� �������������������� ��� ��������� ���� ������������
Revised: �������������� ���������� ������� ��������� ���� �������������������� ��� ��������� ���� ������������

Original: [Figure]
Revised: [Figure]

Original: 3. G OAL - CONDITIONAL RL
Revised: 3.2 Goal-conditional RL

Original: In goal-conditional RL (Kaelbling, 1993; Schaul et al., 2015) the agent aims to find an optimal policy π ∗ ( a | s, s g ) conditioned on goal s g ∈ G , where G stands for goal space that is sub-dimension of state space S (e g. XYZ coordinates, velocity, or quaternion).
Revised: In goal-conditional RL [57, 88], the agent aims to find an optimal policy π ∗ ( a | s, s g ) conditioned on goal s g ∈ G , where G stands for goal space that is sub-dimension of state space S (e g. XYZ coordinates, velocity, or quaternion).

Original: The desired goal s g is sampled from the given goal distribution p ψ : G → [0 , ∞ ) , where ψ stands for task in the task space Ψ
Revised: The desired goal s g is sampled from the given goal distribution p ψ : G → [0 , ∞ ) , where ψ stands for task in the task space Ψ

Original: (e.g. reaching the agent’s leg or touching the torso to the ball).
Revised: (e.g. reaching the agent’s leg or touching the torso to the ball).

Original: The reward function can include a goal-reaching term, that is often modeled as r ψ ( s t , s g )
Revised: The reward function can include a goal-reaching term, that is often modeled as r ψ ( s t , s g )

Original: =
Revised: =

Original: − d ψ ( s t , s g ) , where d ψ ( · , · ) is a task-dependent distance function, such as Euclidean distance, between the sub-dimension of interest in current state s t and given goal s g .
Revised: − d ψ ( s t , s g ) , where d ψ ( · , · ) is a task-dependent distance function, such as Euclidean distance, between the sub-dimension of interest in current state s t and given goal s g .

Original: Some task ψ give multiple goals to the agents.
Revised: Some task ψ give multiple goals to the agents.

Original: In that case, we overload s g to represent a set of goals; { s ig } N ψ i =1 , where N ψ is the number of goals that should be satisfied in the task ψ .
Revised: In that case, we overload s g to represent a set of goals; { s ig } N ψ i =1 , where N ψ is the number of goals that should be satisfied in the task ψ .

Original: 3. M ORPHOLOGY -T ASK G ENERALIZATION
Revised: 3.3 Morphology-Task Generalization

Original: This paper aims to achieve morphology-task generalization, where the learned policy should generalize over tasks and morphologies simultaneously.
Revised: This paper aims to achieve morphology-task generalization, where the learned policy should generalize over tasks and morphologies simultaneously.

Original: The optimal policy should generalize over morphology space M , task Ψ , and minimize the distance to any given goal s g ∈ G .
Revised: The optimal policy should generalize over morphology space M , task Ψ , and minimize the distance to any given goal s g ∈ G .

Original: Mathematically, this objective can be formulated as follows:
Revised: Mathematically, this objective can be formulated as follows:

Original: [Equation]
Revised: [Equation]

Original: where the graph representation of morphology m ∈ M is denoted as G m =
Revised: where the graph representation of morphology m ∈ M is denoted as G m =

Original: ( V m , E m ) , and s m := { s it } | V m | i =1 and a m := { a et } | E m | e =1 stand for a set of local observations and actions of morphology m .
Revised: ( V m , E m ) , and s m := { s i t } | V m | i =1 and a m := { a e t } | E m | e =1 stand for a set of local observations and actions of morphology m .

Original: 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 Ψ , and goal space G simultaneously.
Revised: 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 Ψ , and goal space G simultaneously.

Original:  M ETHOD 4. M X T-B ENCH AS A T EST B ED FOR M ORPHOLOGY -T ASK G ENERALIZATION
Revised:  Method 4.1 MxT-Bench as a Test Bed for Morphology-Task Generalization

Original: To overcome these shortcomings in the existing RL environments, we develop MxT-Bench , which has a wide coverage over both tasks and morphologies to test morphology-task generalization, with the functionalities for procedural generation from minimal blueprints (Figure 4).
Revised: To overcome these shortcomings in the existing RL environments, we develop MxT-Bench , which has a wide coverage over both tasks and morphologies to test morphology-task generalization, with the functionalities for procedural generation from minimal blueprints (Figure 4).

Original: MxT-Bench is built on top of Brax (Freeman et al., 2021) and Composer (Gu et al., 2021a), for faster iteration of behavior distillation with hardware-accelerated environments.
Revised: MxT-Bench is built on top of Brax [31] and Composer [41], for faster iteration of behavior distillation with hardware-accelerated environments.

Original: Beyond supporting multi-morphology and multi-task settings, the scalability of MxT-Bench helps to test the broader range of morphology-task generalization since we can easily generate out-of-distribution tasks and morphologies, compared to manually-designed morphology or task specifications.
Revised: Beyond supporting multi-morphology and multi-task settings, the scalability of MxT-Bench helps to test the broader range of morphology-task generalization since we can easily generate out-of-distribution tasks and morphologies, compared to manually-designed morphology or task specifications.

Original: 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.
Revised: 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.

Original: Through MxT-Bench, we can easily spawn agents that have different numbers of bodies, legs, or different sizes, lengths, and weights.
Revised: Through MxT-Bench, we can easily spawn agents that have different numbers of bodies, legs, or different sizes, lengths, and weights.

Original: Moreover, we can also import the existing complex morphology used in previous work.
Revised: Moreover, we can also import the existing complex morphology used in previous work.

Original: For instance, we include 60+ morphologies that are suitable for goal-reaching, adapted from Gupta et al.
Revised: For instance, we include 60+ morphologies that are suitable for goal-reaching, adapted from Gupta et al.

Original: (2022) designed in MuJoCo.
Revised: [44] designed in MuJoCo.

Original: In the task axis, we design reach, touch, and twisters 3 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.
Revised: In the task axis, we design reach, touch, and twisters 4 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.

Original: In the touch task, agents aim to create and maintain contact between a specified torso and a movable ball.
Revised: In the touch task, agents aim to create and maintain contact between a specified torso and a movable ball.

Original: The touch task requires reaching behavior, while maintaining a conservative momentum to avoid kicking the ball away from the agent.
Revised: The touch task requires reaching behavior, while maintaining a conservative momentum to avoid kicking the ball away from the agent.

Original: Twisters tasks are the multi-goal problems; for instance, the agents should satisfy XY-position for one leg, and Z height for another leg.
Revised: Twisters tasks are the multi-goal problems; for instance, the agents should satisfy XY-position for one leg, and Z height for another leg.

Original: We pre-define 4 variants of twisters with max 3 goals (see
Revised: We pre-define 4 variants of twisters with max 3 goals (see

Original: Appendix C.2 for the details).
Revised: Appendix C.2 for the details).

Original: Furthermore, we could easily specify both initial and goal position distribution with parameterized distribution.
Revised: Furthermore, we could easily specify both initial and goal position distribution with parameterized distribution.

Original: In total, we prepare 180+ environments combining the morphology and task axis for the experiments in the later section.
Revised: In total, we prepare 180+ environments combining the morphology and task axis for the experiments in the later section.

Original: See Appendix C for further details.
Revised: See Appendix C for further details.

Original: 4. B EHAVIOR D ISTILLATION
Revised: 4.2 Behavior Distillation

Original: Toward broader generalization over morphologies and tasks, a single policy should handle the diversity of the scene among morphology M , task Ψ , and goal space G simultaneously.
Revised: Toward broader generalization over morphologies and tasks, a single policy should handle the diversity of the scene among morphology M , task Ψ , and goal space G simultaneously.

Original: Multitask online RL from scratch, however, is difficult to tune, slow to iterate, and hard to reproduce.
Revised: Multi-task online RL from scratch, however, is difficult to tune, slow to iterate, and hard to reproduce.

Original: Instead, we employ behavior cloning on RL-generated expert behaviors to study morphology-task generalization.
Revised: Instead, we employ behavior cloning on RL-generated expert behaviors to study morphology-task generalization.

Original: To obtain rich goal-reaching behaviors, we train a single-morphology single-task policy using PPO (Schulman et al., 2017) with a simple MLP policy, which is significantly more efficient than multi-morphology training done in prior work (Gupta et al., 2022; Kurin et al., 2021; Huang et al., 2020).
Revised: To obtain rich goal-reaching behaviors, we train a single-morphology single-task policy using PPO [89] with a simple MLP policy, which is significantly more efficient than multi-morphology training done in prior work [44, 62, 49].

Original: Since MxT-Bench is built on top of Brax (Freeman et al., 2021), a hardwareaccelerated simulator, training PPO policies can be completed in about 5 ∼ 30 minutes per environment (on NVIDIA RTX A6000).
Revised: Since MxT-Bench is built on top of Brax [31], a hardware-accelerated simulator, training PPO policies can be completed in about 5 ∼ 30 minutes per environment (on NVIDIA RTX A6000).

Original: We collectthe many behaviors per morphology-task combination from the expert policy rollout.
Revised: We collect many behaviors per morphology-task combination from the expert policy rollout.

Original: We then train a single policy π θ with a supervised learning objective:
Revised: We then train a single policy π θ with a supervised learning objective:

Original: [Equation]
Revised: [Equation]

Original: where D m,ψ is an expert dataset of morphology m and task ψ .
Revised: where D m,ψ is an expert dataset of morphology m and task ψ .

Original: Importantly, offline behavior distillation protocol runs (parallelizable) single-task online RL only once, and allows us to reuse the same fixed data to try out various design choices, such as model architectures or local features of control graph, which is often intractable in multi-task online RL.
Revised: Importantly, offline behavior distillation protocol runs (parallelizable) single-task online RL only once, and allows us to reuse the same fixed data to try out various design choices, such as model architectures or local features of control graph, which is often intractable in multi-task online RL.

Original: 4. C ONTROL G RAPH
Revised: 4.3 Control Graph

Original: To learn a single policy that could solve various morphology-task problems, it is essential to unify the input-output interface among those.
Revised: To learn a single policy that could solve various morphology-task problems, it is essential to unify the input-output interface among those.

Original: Inspired by the concept of scene graph (Johnson et al., 2015) and morphological graph (Wang et al., 2018), we introduce the notion of control graph (CG) representation that incorporates goal information while preserving the geometric structure of the task.
Revised: Inspired by the concept of scene graph [55] and morphological graph [100], we introduce the notion of control graph (CG) representation that incorporates goal information while preserving the geometric structure of the task.

Original: Control graph could express the agent’s observations, actions, and goals/tasks in an unified graph space.
Revised: Control graph could express the agent’s observations, actions, and goals/tasks in an unified graph space.

Original: Although most prior morphology-agnostic RL has focused on locomotion (running) with the reward calculated from the velocity of the center of mass, control graph formulation could naturally extend morphology-agnostic RL to multi-task goal-oriented settings: including static single positional goals (reaching), multiple-goal problems (twister-game) and object interaction tasks (ball-touching).
Revised: Although most prior morphology-agnostic RL has focused on locomotion (running) with the reward calculated from the velocity of the center of mass, control graph formulation could naturally extend morphology-agnostic RL to multi-task goal-oriented settings: including static single positional goals (reaching), multiplegoal problems (twister-game) and object interaction tasks (ball-touching).

Original: In practice, we develop two different way to inform the goals and tasks (Figure 3); control graph v1 (CGv1) 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.
Revised: In practice, we develop two different way to inform the goals and tasks (Figure 3); control graph v1 (CGv1) 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.

Original: For instance, in touch task, CGv1 includes XY position of the movable ball as an extra node features of the body node.
Revised: For instance, in touch task, CGv1 includes XY position of the movable ball as an extra node features of the body node.

Original: Moreover, control graph v2 (CGv2) considers given goals as additional disjoint nodes of morphological graph representation.
Revised: Moreover, control graph v2 (CGv2) considers given goals as additional disjoint nodes of morphological graph representation.

Original: These control graph strategies enable the policy to handle a lot of combination of tasks and morphologies simultaneously.
Revised: These control graph strategies enable the policy to handle a lot of combination of tasks and morphologies simultaneously.

Original: Transformer for Control Graph While control 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 
Revised: Transformer for Control Graph While control 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 modules and goals dynamically. We mainly employ Transformer as a policy architecture because it can process morphological graph as a fully-connected graph

Original: [Table]
Revised: [Table]

Original: relationship between the agent’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 relationships between the node beyond manually-encoded geometric structure (Kurin et al., 2021).
Revised: and achieve notable performance by leveraging the hidden relationships between the node beyond manually-encoded geometric structure [62].

Original: Transformer first encodes control graph to the latent representation vector z 0 with shared single-layer MLP and learnable position embedding (PE) (in case of CGv1, we omit s g , but include it to corresponding node observation s i instead):
Revised: Transformer first encodes control graph to the latent representation vector z 0 with shared single-layer MLP and learnable position embedding (PE) (in case of CGv1, we omit s g , but include it to corresponding node observation s i instead):

Original: [Equation]
Revised: z 0 =

Original:  then multi-head attention (MHA) (Vaswani et al., 2017) and layer normalization (LayerNorm) (Ba et al., 2016) are recursively applied to latent representation z l at l -th layer,
Revised: [ MLP ( s 1 ) , . . , MLP ( s | V m | ) , MLP ( s g )] + PE , then multi-head attention (MHA) [98] and layer normalization (LayerNorm) [5] are recursively applied to latent representation z l at l -th layer,

Original: [Equation]
Revised: [Equation]

Original: Before decoding the action per module from the last-layer latent representation z L , we employ the residual connection of the node features (Kurin et al., 2021),
Revised: Before decoding the action per module from the last-layer latent representation z L , we employ the residual connection of the node features [62],

Original: [Equation]
Revised: [Equation]

Original: where shared MLP has a single layer and tanh activation to clip the output within the range of [-1, 1].
Revised: where shared MLP has a single layer and tanh activation to clip the output within the range of [-1, 1].

Original: We mask out the output from the goal modules or modules that have no actuators to ensure action size as | E m | .
Revised: We mask out the output from the goal modules or modules that have no actuators to ensure action size as | E m | .

Original: See Appendix A for further details of implementation.
Revised: See Appendix A for further details of implementation.

Original:  E XPERIMENTS
Revised:  Experiments

Original: We first evaluate the multi-task performance of control graph representation (CGv1, CGv2) 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 out-of-distribution generalization (either morphology or task is unseen) on MxT-Bench (Section 5.1).
Revised: We first evaluate the multi-task performance of control graph representation (CGv1, CGv2) 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 out-of-distribution generalization (either morphology or task is unseen) on MxT-Bench (Section 5.1).

Original: Then, we investigate whether control graph could contribute to obtaining better control prior for multi-task fine-tuning (Section 5.2).
Revised: Then, we investigate whether control graph could contribute to obtaining better control prior for multi-task fine-tuning (Section 5.2).

Original: Lastly, we examine how control graph works well by visualizing attention weights (Section 5.3).
Revised: Lastly, we examine how control graph works well by visualizing attention weights (Section 5.3).

Original: The results are averaged among 4 random seeds.
Revised: The results are averaged among 4 random seeds.

Original: See Appendix A for the hyperparameters.
Revised: See Appendix A for the hyperparameters.

Original: We also examine the other axis of representations or architectures (Appendix B, F, G) and test the effect of dataset size and the number of morphology-task combinations (Appendix H).
Revised: We also examine the other axis of representations or architectures (Appendix B, F, G) and test the effect of dataset size, the number of morphology-task combinations, and model size (Appendix H).

Original: Evaluation Metric Goal-reaching tasks are evaluated by the distance to the goals at the end of episode (Pong et al., 2018; 2020; Ghosh et al., 2021; Choi et al., 2021; Eysenbach et al., 2021).
Revised: Evaluation Metric Goal-reaching tasks are evaluated by the distance to the goals at the end of episode [82, 83, 39, 17, 29].

Original: However, this can be problematic in our settings, because the initial distance or the degree of goalreaching behaviors might be different among various morphologies and tasks.
Revised: 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.

Original: We measure the performance of the policy π by using a normalized final distance metric ¯ d ( M , Ψ; π ) over morphology M and task space Ψ with pre-defined max/min value of each morphology m and task ψ ,
Revised: We measure the performance of the policy π by using a normalized final distance metric ¯ d ( M , Ψ; π ) over morphology M and task space Ψ with pre-defined max/min value of each morphology m and task ψ ,

Original: [Equation]
Revised: [Equation]

Original: where s mT is the last state of the episode, d i,m,ψ max is a maximum, and d i,m,ψ min is a minimum distance of i -th goal s ig with morphology m and task ψ .
Revised: where s mT is the last state of the episode, d i,m,ψ max is a maximum, and d i,m,ψ min is a minimum distance of i -th goal s ig with morphology m and task ψ .

Original: We use a distance threshold to train the expert PPO 
Revised: We use a distance threshold to train the expert PPO policy as d i,m,ψ min , and average distance from initial position of the scene as d i,m,ψ max . Equation 3 is normalized around the range of [0, 1] and the smaller, the better. See Appendix C for the details.

Original:    
Revised:    

Original:  "WFSBHF/PSNBMJ[FE%JTUBODF 5'$(W	L  5'$(W	L5'$(W	L  5'$(W	L5'$(W	L  5'$(W	L.-1	L  .-1	L.-1	L  $PNQPTJUJPOBM	.PSQIPMPHZ  GJOFUVOJOHSBOEJOJU 
Revised:  "WFSBHF/PSNBMJ[FE%JTUBODF 5'$(W	L5'$(W	L  5'$(W	L5'$(W	L  5'$(W	L5'$(W	L  .-1	L.-1	L  .-1	L  $PNQPTJUJPOBM	.PSQIPMPHZ  GJOFUVOJOHSBOEJOJU        

Original: [Figure] "WFSBHF/PSNBMJ[FE%JTUBODF $PNQPTJUJPOBM	5BTL  [Figure]
Revised: "WFSBHF/PSNBMJ[FE%JTUBODF $PNQPTJUJPOBM	5BTL  GJOFUVOJOHSBOEJOJU        

Original:  0VUPG%JTUSJCVUJPO GJOFUVOJOHSBOEJOJUpolicy as d i,m,ψ min , and average distance from initial position of the scene as d i,m,ψ max . Equation 3 is normalized around the range of [0, 1] and the smaller, the better. See Appendix C for the details.
Revised: "WFSBHF/PSNBMJ[FE%JTUBODF 0VUPG%JTUSJCVUJPO GJOFUVOJOHSBOEJOJU

Original: 5. B EHAVIOR D ISTILLATION ON M X T-B ENCH
Revised: 5.1 Behavior Distillation on MxT-Bench

Original: We systematically evaluate three types of morphology-task generalization; in-distribution, compositional, and out-of-distribution generalization through behavior distillation.
Revised: We systematically evaluate three types of morphology-task generalization; in-distribution, compositional, and out-of-distribution generalization through behavior distillation.

Original: As baselines, we compare
Revised: As baselines, we compare

Original: MLP, GNN (Wang et al., 2018) with CGv1, Transformer with CGv1 or CGv2, and tokenized CGv2, similar to Reed et al.
Revised: MLP, GNN [100] with CGv1, Transformer with CGv1 or CGv2, and tokenized CGv2, similar to Reed et al.

Original: (2022) (see Appendix G for the details).
Revised: [87] (see Appendix G for the details).

Original: In in-distribution settings, we prepare 50 environments and 60 unimal environments adapted from
Revised: In in-distribution settings, we prepare 50 environments and 60 unimal environments adapted from

Original: Gupta et al.
Revised: Gupta et al.

Original: (2022) (see Appendix C.3 for the details) for both training and evaluation.
Revised: [44] (see Appendix C.3 for the details) for both training and evaluation.

Original: Compositional and out-of-distribution settings evaluate zero-shot transfer.
Revised: Compositional and out-of-distribution settings evaluate zero-shot transfer.

Original: In compositional settings, we test morphology 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.
Revised: In compositional settings, we test morphology 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.

Original: In out-of-distribution settings, we also leverage 50 environments as a training set, and define 27 environments with diversified morphologies and unseen task as an evaluation set.
Revised: In out-of-distribution settings, we also leverage 50 environments as a training set, and define 27 environments with diversified morphologies and unseen task as an evaluation set.

Original: The proficient behavioral data contains 12k transitions per each environment.
Revised: The proficient behavioral data contains 12k transitions per each environment.

Original: See Appendix E for the details of environment division.
Revised: See Appendix E for the details of environment division.

Original: Table 1 reveals that CGv2 achieves the best multi-task goal-reaching performances among other possible combinations in all the aspects of generalization.
Revised: Table 1 reveals that CGv2 achieves the best multi-task goal-reaching performances among other possible combinations in all the aspects of generalization.

Original: Comparing average normalized distance,
Revised: Comparing average normalized distance,

Original: CGv2 improves the multi-task performance against the second best, CGv1, by 23% in in-distribution evaluation.
Revised: CGv2 improves the multi-task performance against the second best, CGv1, by 23% in in-distribution evaluation.

Original: Following previous works (Kurin et al., 2021; Gupta et al., 2022), Transformer with CGv achieves better goal-reaching behaviors than GNN.
Revised: Following previous works [62, 44], Transformer with CGv1 achieves better goal-reaching behaviors than GNN.

Original: In compositional and out-of-distribution zero-shot evaluation, CGv2 outperforms other choices by 14 ∼ 18%.
Revised: In compositional and out-of-distribution zero-shot evaluation, CGv2 outperforms other choices by 14 ∼ 18%.

Original: Moreover, the compositional zero-shot performance of CGv2 is comparable with the performance of CGv1 in in-distribution settings.
Revised: Moreover, the compositional zero-shot performance of CGv is comparable with the performance of CGv1 in in-distribution settings.

Original: These results imply CGv2 might be the better formulation to realize the morphology-task generalization.
Revised: These results imply CGv might be the better formulation to realize the morphology-task generalization.

Original: 5. D OES C ONTROL G RAPH O BTAIN B ETTER P RIOR FOR C ONTROL ?
Revised: 5.2 Does Control Graph Obtain Better Prior for Control?

Original: 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-Bench.
Revised: 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-Bench.

Original: We adopt the same morphology-task division for compositional and out-of-distribution evaluation in Section 5.1.
Revised: We adopt the same morphology-task division for compositional and out-of-distribution evaluation in Section 5.1.

Original: 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.
Revised: 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.

Original: The same as zero-shot transfer results in Section 5.1, CGv2 outperforms other baselines, and is better than the second best, CGv1 by 50 ∼ 55%.
Revised: The same as zero-shot transfer results in Section 5.1, CGv2 outperforms other baselines, and is better than the second best, CGv1 by 50 ∼ 55%.

Original: Furthermore, CGv2 could work even with a small amount of dataset (4k, 8k); for instance, in compositional morphology evaluation (left in Figure 5), CGv2 trained with 8k transitions still outperforms competitive combinations with 12k transitions, which indicates better sample efficiency (see Appendix J for the detailed scores).
Revised: Furthermore, CGv2 could work even with a small amount of dataset (4k, 8k); for instance, in compositional morphology evaluation (left in Figure 5), CGv2 trained with 8k transitions still outperforms competitive combinations with 12k transitions, which indicates better sample efficiency (see Appendix J for the detailed scores).

Original: These results suggest CGv2 significantly captures the structure of control graph as prior knowledge for downstream tasks.
Revised: These results suggest CGv2 significantly captures the structure of control graph as prior knowledge for downstream tasks.

Original: 5. H OW D OES C ONTROL G RAPH W ORK W ELL ?
Revised: 5.3 How Does Control Graph Work Well?

Original: The experimental results suggest, despite the slight difference, CGv2 generalizes various morphologies and tasks better than CGv1.
Revised: The experimental results suggest, despite the slight difference, CGv2 generalizes various morphologies and tasks better than CGv1.

Original: To find out the difference between those, we qualitatively analyze�������������������� ��������� the attention weights in Transformer.
Revised: To find out the difference between those, we qualitatively analyze the attention weights in Transformer.

Original: Figure 6 shows that CGv2 consistently focuses on goal nodes over time, and activates important nodes to solve the task; for instance, in ant_twisters, CGv2 firstly  tries to raise the agent’s legs to satisfy goal1 and goal2, and then focus on reaching a leg to goal0.
Revised: Figure 6 shows that CGv2 consistently focuses on goal nodes over time, and activates important nodes to solve the task; for instance, in ant_twisters, CGv2 firstly [Figure] tries to raise the agent’s legs to satisfy goal1 and goal2, and then focus on reaching a leg to goal0.

Original: Temporally-consistent attention to goal nodes and dynamics attention to relevant nodes can contribute to generalization over goal-directed tasks and morphologies.
Revised: Temporally-consistent attention to goal nodes and dynamics attention to relevant nodes can contribute to generalization over goal-directed tasks and morphologies.

Original: In contrast, CGv1 does not show such consistent activation to goal-conditioned node, rather it demonstrates some periodic patterns as implied in prior works (Kurin et al., 2021; Huang et al., 2020).
Revised: In contrast, CGv1 does not show such consistent activation to goal-conditioned node, rather it demonstrates some periodic patterns as implied in prior works [62, 49].

Original: See Appendix K for the full results.
Revised: See Appendix K for the full results.

Original:  D ISCUSSION AND L IMITATION
Revised:  Discussion and Limitation

Original: While the experimental evaluation on MxT-Bench implies that control graph is a simple and effective method to distill the diverse proficient behavioral data into a generalizable single policy, there are some limitations.
Revised: While the experimental evaluation on MxT-Bench implies that control graph is a simple and effective method to distill the diverse proficient behavioral data into a generalizable single policy, there are some limitations.

Original: For instance, we focus on distillation from expert policies only, but it is still unclear whether control graph works with moderate or random quality behaviors in offline RL (Fujimoto et al., 2019; Levine et al., 2020; Fujimoto & Gu, 2021).
Revised: For instance, we focus on distillation from expert policies only, but it is still unclear whether control graph works with moderate or random quality behaviors in offline RL [33, 66, 32].

Original: Combining distillation with iterative data collection (Ghosh et al., 2021; Matsushima et al., 2021) or online fine-tuning (Li et al., 2022) would be a promising future work.
Revised: Combining distillation with iterative data collection [39, 74] or online fine-tuning [70] would be a promising future work.

Original: 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  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).
Revised: 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 [72, 38, 16].

Original: Since control graph only uses readily accessible features in any simulator and could be automatically defined through URDFs or MuJoCo XMLs, inthe future work we aim to keep training our best control 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), 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-task families.
Revised: Since control 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 control graph architecture policy on additional data from more functional and realistic control behaviors from other simulators like MuJoCo [103, 94], PyBullet [91], IsaacGym [16, 81], and Unity [56], and show it achieves better scaling laws than other representations [87] on broader morphology-task families.

Original:  C ONCLUSION
Revised:  Conclusion

Original: The broader range of behavior generalization is a promising paradigm for RL.
Revised: The broader range of behavior generalization is a promising paradigm for RL.

Original: To achieve morphologytask generalization, we propose control graph, which express the agent’s modular observations, actions, and goals as a unified graph representation while preserving the geometric task structure.
Revised: To achieve morphologytask generalization, we propose control graph, which expresses the agent’s modular observations, actions, and goals as a unified graph representation while preserving the geometric task structure.

Original: As a test bed for morphology-task generalization, we also develop MxT-Bench, which enables the scalable procedural generation of agents and tasks with minimal blueprints.
Revised: As a test bed for morphology-task generalization, we also develop MxT-Bench, which enables the scalable procedural generation of agents and tasks with minimal blueprints.

Original: Fast-generated behavior datasets of MxT-Bench with RL allows efficient representation and architecture selection through supervised learning, and CGv2, variant of control graph, achieves the best multi-task performances among other possible designs (MLP, GNN and Transformer with CGv1, and tokenized-CGv2, etc), outperforming them in in-distribution evaluation (by 23 %), zero-shot transfer among compositional or out-of-distribution evaluation (by 14 ∼ 18 %) and fine-tuning for downstream multi-task imitation ( ∼ 55 %).
Revised: Fast-generated behavior datasets of MxT-Bench with RL allow efficient representation and architecture selection through supervised learning, and CGv2, variant of control graph, achieves the best multi-task performances among other possible designs (MLP, GNN and Transformer with CGv1, and tokenized-CGv2, etc), outperforming them in in-distribution evaluation (by 23 %), zero-shot transfer among compositional or out-ofdistribution evaluation (by 14 ∼ 18 %) and fine-tuning for downstream multi-task imitation (50 ∼%).

Original: We hope our work will encourage the community to explore scalable yet incremental approach for building a universal controller.
Revised: We hope our work will encourage the community to explore scalable yet incremental approaches for building a universal controller.

Revised: Acknowledgments

Revised: This work was supported by JSPS KAKENHI Grant Number JP22J21582.

Revised: We thank Mitsuhiko

Revised: 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

Original:  A PPENDIX A D ETAILS OF I MPLEMENTATION
Revised: Appendix A Details of Implementation

Original: The hyperparameters we used are listed in Table 2.
Revised: The hyperparameters we used are listed in Table 2.

Original: We implement MLP and Transformer policy with
Revised: We implement MLP and Transformer policy with

Original: Jax (Bradbury et al., 2018) and Flax (Heek et al., 2020).
Revised: Jax [9] and Flax [45].

Original: For the implementation of GNN policy, we use a graph neural network library, Jraph (Godwin et al., 2020)
Revised: For the implementation of GNN policy, we use a graph neural network library, Jraph [40].

Original: [Table]
Revised: [Table]

Original: B A RCHITECTURE S ELECTION : P OSITION E MBEDDING
Revised: B Architecture Selection: Position Embedding

Original: As a part of architecture selection [Table] we investigate whether position embedding (PE) contribute to the generalization.
Revised: As a part of architecture selection, we investigate whether position embedding (PE) contributes to the generalization.

Original: Our empirical results in Table 3 suggest that, multi-task goal reaching performance seems comparable between those.
Revised: Our empirical results in Table 3 suggest that, multi-task goal reaching performance seems comparable between those.

Original: However, in more diverse morphology domains [Table] PE plays an important role.
Revised: However, in more diverse morphology domains, PE plays an important role.

Original: Therefore, we include PE into a default design
Revised: Therefore, we include PE into a default design

Original: [Table]
Revised: [Table]

Original: C D ETAILS OF M X T-B ENCH C. M ORPHOLOGY
Revised: C Details of MxT-Bench C.1 Morphology

Original: We prepare 4 base blueprints; ant, centipede, claw, and worm, for procedural scene generation (Figure 7).
Revised: We prepare 4 base blueprints; ant, centipede, claw, and worm, for procedural scene generation (Figure 7).

Original: Base ant has 1 torso and 2 legs with 2 joints per limb, and we could procedurally generate the agents with different number of limbs.
Revised: Base ant has 1 torso and 2 legs with 2 joints per limb, and we could procedurally generate the agents with different number of limbs.

Original: Base centipede has 2 bodies and 4 legs with 2 joints per limb, and we could procedurally generate the agents with different number of bodies.
Revised: Base centipede has 2 bodies and 4 legs with 2 joints per limb, and we could procedurally generate the agents with different number of bodies.

Original: Base claw has 1 torso and 2 legs with 4 joints per limb (each leg consists of 3 modules), and we could procedurally generate the agents with different number of limbs.
Revised: Base claw has 1 torso and 2 legs with 4 joints per limb (each leg consists of 3 modules), and we could procedurally generate the agents with different number of limbs.

Original: Base worm has 2 bodies and no legs, and we could procedurally generate the agents with different number of bodies.
Revised: Base worm has 2 bodies and no legs, and we could procedurally generate the agents with different number of bodies.

Original: Furthermore, we develop the functionality for morphology diversification, with missing, mass, and size parameters (Figure 8).
Revised: Furthermore, we develop the functionality for morphology diversification, with missing, mass, and size parameters (Figure 8).

Original: Missing randomization lacks one module at one leg.
Revised: Missing randomization lacks one module at one leg.

Original: This might be an equivalent situation that one leg is broken.
Revised: This might be an equivalent situation that one leg is broken.

Original: Mass randomization changes the default mass of each module with specified scales.
Revised: Mass randomization changes the default mass of each module with specified scales.

Original: The appearance does not change, but certainly the dynamics would differ.
Revised: The appearance does not change, but certainly the dynamics would differ.

Original: Size randomization changes the default length and radius of each module with specified scales.
Revised: Size randomization changes the default length and radius of each module with specified scales.

Original: C. T ASK
Revised: C.2 Task

Original: We prepare 4 base tasks with parameterized goal distributions; reach, touch, twisters, and push, for procedural task generation (Figure 9).
Revised: We prepare 4 base tasks with parameterized goal distributions; reach, touch, twisters, and push, for procedural task generation (Figure 9).

Original: Reach task requires the agents to put their one leg to the given goal position (XY).
Revised: Reach task requires the agents to put their one leg to the given goal position (XY).

Original: The variant, reach_hard task, represents that the goal distribution is farther than reach task.
Revised: The variant, reach_hard task, represents that the goal distribution is farther than reach task.

Original: Touch task requires the agents to contact their body or torso to the movable ball (i.e. movable ball is an goal).
Revised: Touch task requires the agents to contact their body or torso to the movable ball (i.e. movable ball is an goal).

Original: Twisters is a multi-goal problem; the agents should satisfy given goals at the same time.
Revised: Twisters is a multi-goal problem; the agents should satisfy given goals at the same time.

Original: There are two basic constraints; reach and handsup.
Revised: There are two basic constraints; reach and handsup.

Original: Handsup requires the agents to raise their one leg to the given goal Z height.
Revised: Handsup requires the agents to raise their one leg to the given goal Z height.

Original: Twisters has some combinations like reach_handsup, reach_hard_handsup, reach2_handsup, or reach_handsup2.
Revised: Twisters has some combinations like reach_handsup, reach_hard_handsup, reach2_handsup, or reach_handsup2.

Original: 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.
Revised: 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.

Original: Push task requires the agents to move the box object to the given goal position (XY).
Revised: Push task requires the agents to move the box object to the given goal position (XY).

Original: Since it has richer interaction with the object, push might be more difficult task than those three.
Revised: Since it has richer interaction with the object, push might be more difficult task than those three.

Original: We use this task in Appendix L.
Revised: We use this task in Appendix L.

Original: C. C USTOM M ORPHOLOGY
Revised: C.3 Custom Morphology

Original: MxT-Bench also supports custom morphology import used in previous work.
Revised: MxT-Bench also supports custom morphology import used in previous work.

Original: For instance, Gupta et al.
Revised: For instance, Gupta et al.

Original: (2022) propose unimal agents that are generated via evolutional strategy and designed for MuJoCo.
Revised: [44] propose unimal agents that are generated via evolutional strategy and designed for MuJoCo.

Original: Since they are not manually designed, their morphologies seem more diverse than our ant, centipede, claw, and worm.
Revised: Since they are not manually designed, their morphologies seem more diverse than our ant, centipede, claw, and worm.

Original: We inspect unimals whether they are suitable for goal-reaching, and include 72 morphologies from there (and use 60 morphologies for the experiments).
Revised: We inspect unimals whether they are suitable for goal-reaching, and includemorphologies from there (and use 60 morphologies for the experiments).

Original: Figure 10 shows some example of unimals.
Revised: Figure 10 shows some example of unimals.

Original: C. C OMPARISON TO E XISTING B ENCHMARK
Revised: C.4 Comparison to Existing Benchmark

Original: Since the current RL community has not paid much attention to embodied control so far, there are no suitable benchmarks to quantify the generalization over various tasks beyond single locomotion tasks and morphologies at the same time.
Revised: Since the current RL community has not paid much attention to embodied control so far, there are no suitable benchmarks to quantify the generalization over various tasks beyond single locomotion tasks and morphologies at the same time.

Original: In addition, the scalability to various morphologies or tasks seems to be required for the benchmark, because we should avoid “overfitting” to manually-designed tasks.
Revised: In addition, the scalability to various morphologies or tasks seems to be required for the benchmark, because we should avoid “overfitting” to manually-designed tasks.

Original: As summarized in Figure 4, MuJoCo (Todorov et al., 2012) or DM Control (Tassa et al., 2018), the most popular benchmarks in the continuous control domain, could not evaluate task or morphology generalization; they only have a single morphology for a single task as a pre-defined environment.
Revised: As summarized in Figure 4, MuJoCo [96] or DM Control [94], the most popular benchmarks in the continuous control domain, could not evaluate task or morphology generalization; they only have a single morphology for a single task as a pre-defined environment.

Original: (Yu et al., 2019) propose a robot manipulation benchmark for meta RL, but it does not care about the morphology.
Revised: Yu et al. [103] propose a robot manipulation benchmark for meta RL, but it does not care about the morphology.

Original: Furthermore, while it has quite a diverse set of tasks, the scalability of environments seems to be limited.
Revised: Furthermore, while it has quite a diverse set of tasks, the scalability of environments seems to be limited.

Original: In contrast, previous morphology-agnostic 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.
Revised: In contrast, previous morphology-agnostic RL works [49, 62, 48, 97] have a set of different morphologies adapted from MuJoCo agents.

Original: Gupta et al.
Revised: Gupta et al.

Original: (2022) also provide a much larger set of agents that are produced via joint-optimization of morphology and task rewards by the evolutionary strategy (Gupta et al., 2021), with kinematics and dynamics randomization.
Revised: [44] also provide a much larger set of agents that are produced via joint-optimization of morphology and task rewards by the evolutionary strategy [43], with kinematics and dynamics randomization.

Original: However, those works only aim to solve single locomotion tasks, i.e. running forward as fast as possible.
Revised: However, those works only aim to solve single locomotion tasks, i.e. running forward as fast as possible.

Original: D D ETAILS OF E XPERT D ATA G ENERATION
Revised: D Details of Expert Data Generation

Original: We train the single-task PPO (Schulman et al., 2017) at each environment to obtain the expert policy.
Revised: We train the single-task PPO [89] at each environment to obtain the expert policy.

Original: For the reward function, we basically adopt dense reward (except for push task) − d ψ ( s m , s g ) until agents satisfy a given condition, 

Revised: For the reward function, we basically adopt dense reward (except for push task) − d ψ ( s m , s g ) until agents satisfy a given condition,

Original: [ d ψ ( s m , s g ) ≤ d m,ψ min ] .
Revised: [ d ψ ( s m , s g ) ≤ d m,ψ min ] .

Original: See Appendix E for the threshold d m,ψ min .
Revised: See Appendix E for the threshold d m,ψ min .

Original: After convergence, we collect the proficient behaviors.
Revised: After convergence, we collect the proficient behaviors.

Original: Unless otherwise specified, we use 12k transitions per environment.
Revised: Unless otherwise specified, we use 12k transitions per environment.

Original: The average normalized final distance of those datasets are mostly less than 0.1.
Revised: The average normalized final distance of those datasets are mostly less than 0.1.

Original: E E NVIRONMENTS D IVISION
Revised: E Environments Division

Original: Throughout the experiments, we test a lot of morphology-task combinations to investigate the in-distribution generalization, compositional generalization for morphology and task, and out-ofdistribution generalization.
Revised: Throughout the experiments, we test a lot of morphology-task combinations to investigate the in-distribution generalization, compositional generalization for morphology and task, and out-ofdistribution generalization.

Original: In this section, we list up the combination of environments used in the experiments.
Revised: In this section, we list up the combination of environments used in the experiments.

Original: Table 4 explains the combinations for the experiments of in-distribution generalization, and compositional generalization for morphology (both zero-shot transfer and fine-tuning) in Table 1 and Figure 5.
Revised: Table 4 explains the combinations for the experiments of in-distribution generalization, and compositional generalization for morphology (both zero-shot transfer and fine-tuning) in Table 1 and Figure 5.

Original: For compositional morphology evaluation, we use Morph-Train division as training dataset and
Revised: For compositional morphology evaluation, we use Morph-Train division as training dataset and

Original: Morph-Test division as evaluation environments.
Revised: Morph-Test division as evaluation environments.

Original: For the evaluation of dataset size and the number of morphology-task combinations (Appendix H), we use In-Distribution, 12 Env, and 25 Env division as train datasets and test environments.
Revised: For the evaluation of dataset size and the number of morphology-task combinations (Appendix H), we use In-Distribution, 12 Env, and 25 Env division as train datasets and test environments.

Original: Table 6 also explains the combinations for the experiments of in-distribution generalization with unimals (Gupta et al., 2022) in Table 1 and Table 11.
Revised: Table 6 also explains the combinations for the experiments of in-distribution generalization with unimals [44] in Table 1 and Table 11.

Original: Table 5 provides the combinations for the experiments of compositional generalization for task and out-of-distribution generalization (both zero-shot transfer and fine-tuning) in Table 1 and Figure 5.
Revised: Table 5 provides the combinations for the experiments of compositional generalization for task and out-of-distribution generalization (both zero-shot transfer and fine-tuning) in Table 1 and Figure 5.

Original: We use them as test environments, and for training datasets, we leverage In-Distribution division in
Revised: We use them as test environments, and for training datasets, we leverage In-Distribution division in

Original: Table 4.
Revised: Table 4.

Original: 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 L for the details).
Revised: 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 L for the details).

Original: Tablealso shows the environment division [Figure] [Table] [Figure]
Revised: Tablealso shows the environment division for both zero-shot transfer and fine-tuning.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure][Table] [Table] [Table]
Revised: [Figure]

Original: (cid:34) ant_

Original: [Table]
Revised: [Table]

Original: _5_mass_0.5_1.0_1.

Original: (cid:34) ant_ [Table] _5_mass_1.0_3.0_3. (cid:34) ant_push_5_size_0.9_1.0_
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: (cid:34) centipede_push_diverse centipede_push_3_b_r_ 1. (cid:34) centipede_push_3_b_l_ 1. [Table] (cid:34)
Revised: [Figure]

Original: centipede_push_3_

Original: [Table]
Revised: [Table]

Original: _r_ [Table] [Table] (cid:34) centipede_push_
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Table] _l_ [Table] [Table]
Revised: [Figure]

Original: (cid:34)

Original: [Table]
Revised: [Table]

Original: _push_ [Table] _b_r_ [Table]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: (cid:34) centipede_

Original: [Table] [Table] [Table] [Table]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: (cid:34) centipede_push_4_ [Table] _l_ 1.
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: (cid:34) centipede_push_ [Table] [Table] [Table]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: (cid:34) centipede_push_4_ [Table] _l_ 1. (cid:34) centipede_push_
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: _size_ [Table] (cid:34) centipede_push_3_mass_
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: (cid:34) centipede_

Original: [Table]
Revised: [Table]

Original: _3_mass_ [Table][Table] (cid:34) centipede_push_3_mass_1.0_3.0_3.
Revised: [Table]

Original: (cid:34)

Original: [Table]
Revised: [Table]

Original: F R EPRESENTATION S ELECTION : L OCAL O BSERVATIONS
Revised: F Representation Selection: Local Observations

Original: In this section, we explore the set of node features that is most suitable for behavior distillation in multi-morphology-multi-task settings, and provide the design principle of node features.
Revised: In this section, we explore the set of node features that is most suitable for behavior distillation in multi-morphology-multi-task settings, and provide the design principle of node features.

Original: Previous morphology-agnostic studies have worked on searching network architectures that could efficiently process the encoded morphological graph, such as GNN (Wang et al., 2018), GNN for tree structure (Huang et al., 2020), and various Transformers (Kurin et al., 2021; Gupta et al., 2022; Trabucco et al., 2022; Hong et al., 2022).
Revised: Previous morphology-agnostic studies have worked on searching network architectures that could efficiently process the encoded morphological graph, such as GNN [100], GNN for tree structure [49], and various Transformers [62, 44, 97, 48].

Original: However, it is still unclear what kind of observations in the system could contribute to morphology-task generalization as the local features for each node of control graph.
Revised: However, it is still unclear what kind of observations in the system could contribute to morphology-task generalization as the local features for each node of control graph.

Original: It could be valuable to identify how each feature contributes to goal-reaching or generalization and which sets of node features perform best.
Revised: It could be valuable to identify how each feature contributes to goal-reaching or generalization and which sets of node features perform best.

Original: There are a lot of observable candidates per module in the agent’s system, such as Cartesian position ( p ), Cartesian velocity ( v) , quotation ( q ), angular velocity ( a ), joint angle ( ja ), joint range ( jr ), limb id ( id ), joint velocity ( jv ), relative position ( rp ), relative velocity ( rv ), and morphological information ( m ).
Revised: There are a lot of observable candidates per module in the agent’s system, such as Cartesian position ( p ), Cartesian velocity ( v) , quotation ( q ), angular velocity ( a ), joint angle ( ja ), joint range ( jr ), limb id ( id ), joint velocity ( jv ), relative position ( rp ), relative rotation ( rr ), and morphological information ( m ).

Original: Morphological information contains module’s shape, mass, inertia, actuator’s gear, and dofindex, etc.
Revised: Morphological information contains module’s shape, mass, inertia, actuator’s gear, and dofindex, etc.

Original: To shed light on the importance of the local representation per module, we execute an extensive ablation of node feature selections and prepossessing.
Revised: To shed light on the importance of the local representation per module, we execute an extensive ablation of node feature selections and prepossessing.

Original: In prior works, Huang et al.
Revised: In prior works, Huang et al.

Original: (2020) and the following (Kurin et al., 2021; Hong et al., 2022) used { p , v , q , a , ja , jr , id } and Gupta et al.
Revised: [49] and the following [62, 48] used { p , v , q , a , ja , jr , id } and Gupta et al.

Original: (2022) used { p , v , q , a , ja , jr , jv , rp , rv , m }.
Revised: [44] used { p , v , q , a , ja , jr , jv , rp , rr , m }.

Original: Considering the intersection of those, we define { p , v , q , a , ja , jr } as base_set and test the combination to other observations ( jv , id , rp , rv , m ).
Revised: Considering the intersection of those, we define { p , v , q , a , ja , jr } as base_set and test the combination to other observations ( jv , id , rp , rr , m ).

Original: We prepare In-Distribution division in Table 4 and evaluate in-distribution generalization with CGv
Revised: We prepare In-Distribution division in Table 4 and evaluate in-distribution generalization with CGv

Original: IO.
Revised: IO.

Original: Table 8 shows that, while some additional features ( id , rp , m ) contribute to improving the goal-reaching performance, the most effective feature seems morphological information ( m ).
Revised: Table 8 shows that, while some additional features ( id , rp , m ) contribute to improving the goal-reaching performance, the most effective feature seems morphological information ( m ).

Original: These results suggest base_set contains sufficient features for control [Table] and raw morphological properties ( m ) serves better task specification than manually-encoded information such as limb id ( id ).
Revised: These results suggest 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 ).

Original: The key observations are; (1) morphological information is critical for morphology-task generalization, and (2) extra observation might disrupt the goal-reaching performances, such as relative velocity between parent and child node ( rv ).
Revised: The key observations are; (1) morphological information is critical for morphology-task generalization, and (2) extra observation might disrupt the goal-reaching performances, such as relative rotation between parent and child node ( rr ).

Original: Throughout the paper, we use base_set-m for node features of control graph representation
Revised: Throughout the paper, we use base_set-m for node features of control graph representation

Original: [Table]
Revised: [Table]

Original: used base_set-id and Gupta used base_set-jv-rp-rv-m .

Original: While some additional features ( id , rp , m ) contribute to improving the goal-reaching performance, the most effective feature seems morphological information.
Revised: , m ) contribute to improving the goal-reaching performance, the most effective feature seems morphological information.

Original: These results suggest base_set contains sufficient features for control.
Revised: These results suggest base_set contains sufficient features for control.

Original: and raw morphological properties ( m ) serves better task specification than manually-encoded information such as limb id ( id )
Revised: and raw morphological properties ( m ) serves better task specification than manually-encoded information such as limb id ( id ).

Original: G A RCHITECTURE S ELECTION : T OKENIZED C ONTROL G RAPH
Revised: G Architecture Selection: Tokenized Control Graph

Original: In the recent literature, offline RL is considered as upervised sequential modeling problem Chen et al. (2021a), and some works (Janner et al., 2021; Reed et al., 2022) tokenize the continuous observations and actions, like an analogy of vision transformer (Dosovitskiy et al., 2020); treating input modality like a language.
Revised: In the recent literature, offline RL is considered as supervised sequential modeling problem [12], and some works [53, 87] tokenize the continuous observations and actions, like an analogy of vision transformer [26]; treating input modality like a language.

Original: As a part of offline architecture selection, we examine the effectiveness of tokenization.
Revised: As a part of offline architecture selection, we examine the effectiveness of tokenization.

Original: We mainly follows the protocol in Reed et al.
Revised: We mainly follow the protocol in Reed et al.

Original: ; we first apply mu-law encoding (Oord et al., 2016) to control graph representation:
Revised: [87]; we first apply mu-law encoding [77] to control graph representation:

Original: [Equation]
Revised: [Equation]

Original: with µ = 100 and M = 256 .
Revised: with µ = 100 and M = 256 .

Original: This pre-processing could normalize the input to the range of [0 , 1] .
Revised: This pre-processing could normalize the input to the range of [0 , 1] .

Original: Then, we discretize pre-processed observations and actions with 1024 bins.
Revised: Then, we discretize pre-processed observations and actions with 1024 bins.

Original: To examine the broader range of design choice, we prepare the following 6 variants of tokenized control graph: whether layer normalization is added after embedding function (LN) (Bhatt et al., 2019; Parisotto et al., 2019; Furuta et al., 2021a; Chen et al., 2021a), or outputting discretized action (D), smoothed action by taking the average of bins (DA), and continuous values directly (C).
Revised: To examine the broader range of design choice, we prepare the following 6 variants of tokenized control graph: whether layer normalization is added after embedding function (LN) [7, 78, 34, 12], or outputting discretized action (D), smoothed action by taking the average of bins (DA), and continuous values directly (C).

Original: As shown in Table 9, predicting continuous value reveals the better performance than discretized action maybe due to the approximation error.
Revised: As shown in Table 9, predicting continuous value reveals the better performance than discretized action maybe due to some approximation errors.

Original: However, the performance of Token-CGv2 (C) or
Revised: However, the performance of Token-CGv2 (C) or

Original: Token-CGv2 (C, LN) is still lower than CGv2 itself.
Revised: Token-CGv2 (C, LN) is still lower than CGv2 itself.

Original: This might be because tokenization looses some morphological invariance among nodes.
Revised: This might be because tokenization looses some morphological invariance among nodes.

Original: In Table 1, we adopt Token-CGv2 (C) for comparison.[Figure]
Revised: In Table 1, we adopt Token-CGv2 (C) for comparison.

Original: [Table]
Revised: [Table]

Original: H D OES T RANSFORMER WITH CG V 2 SCALE UP WITH DATASET OR MORPHOLOGY - TASK SIZE ?
Revised: H Does Transformer with Control Graph scale up with dataset/morphology-task/model size?

Original: The important aspect of the success in large language models is the scalability to the size of training data and the number of tasks for joint-training (Raffel et al., 2020; Brown et al., 2020).
Revised: The important aspect of the success in large language models is the scalability to the size of training data, the number of tasks for joint-training, and the number of parameters [86, 10].

Original: One natural question is whether the similar trend holds even in RL.
Revised: One natural question is whether the similar trend holds even in RL.

Original: Figure 11 suggests that the performance can get better if we increase the number of datasets.
Revised: Figure 11 suggests that the performance can get better if we increase the number of datasets and the number of parameters.

Original:  In contrast, when we increase the number of environments, the performance degrades while Transformer with CGv2 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
Revised: The performance of Transformer with 0.4M parameters is equivalent to that of MLP with 3.1M. In contrast, when we increase the number of environments, the performance degrades while Transformer withCGv1 and CGv2 surpasses the degree of degradation, which seems inevitable trends in multi-task RL [104, 63] and an important future direction towards generalist controllers.

Original: [Figure]
Revised: [Figure]

Original: I P ERCENTAGE OF I MPROVEMENT
Revised: I Percentage of Improvement

Original: For clarification [Figure] we compute the percentage of improvement between two average normalized final distances (defined in Equation 3) 
Revised: For clarification, we compute the percentage of improvement between two average normalized final distances (defined in Equation 3) ¯

Original: ¯ d 1 and ¯ d 2 as follows ( ¯ d 1 < ¯ d 2 ):
Revised: d 1 and ¯ d 2 as follows ( ¯ d 1 < ¯ d [Equation]

Original: [Equation]
Revised: [Equation]

Original: J A DDITIONAL R ESULTS
Revised: J Additional Results

Original: In this section, we provide the detailed performance of in-distribution generalization (Table 10 and
Revised: In this section, we provide the detailed performance of in-distribution generalization (Table 10 and

Original: Table 11), compositional morphology and task generalization (Table 12 and Table 13)
Revised: Table 11), compositional morphology and task generalization (Table 12 and Table 13)

Original: [Table]
Revised: [Table]

Original: and out-ofdistribution generalization (Table 14).
Revised: and out-ofdistribution generalization (Table 14).

Original: For fine-tuning experiments, we summarized the detailed scores of Figure 5 in Table
Revised: For fine-tuning experiments, we summarized the detailed scores of Figure 5 in Table

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: K A DDITIONAL A TTENTION A NALYSIS
Revised: K Additional Attention Analysis

Original: In this section, we provide the full results of attention analysis of Transformer (with CGv1 and CGv2) (Figure 12 and Figure 13).
Revised: In this section, we provide the full results of attention analysis of Transformer (with CGv1 and CGv2) (Figure 12 and Figure 13).

Original: The experimental results reveal that despite slight differences, CGv2 generalizes various morphologies and tasks better than CGv1.
Revised: The experimental results reveal that despite slight differences, CGv2 generalizes various morphologies and tasks better than CGv1.

Original: To find out the difference between those, we qualitatively analyze the attention weights in Transformer.
Revised: To find out the difference between those, we qualitatively analyze the attention weights in Transformer.

Original: Figure 12 shows that CGv2 consistently focuses on goal nodes over time, and activates important nodes to solve the task; for instance, in centipede_touch (top),
Revised: Figure 12 shows that CGv2 consistently focuses on goal nodes over time, and activates important nodes to solve the task; for instance, in centipede_touch (top),

Original: CGv2 pays attention to corresponding nodes (torso0 an goal0) at the beginning of the episode, and gradually sees other relevant nodes (torso1 and torso2) to hold the movable ball.
Revised: CGv2 pays attention to corresponding nodes (torso0 an goal0) at the beginning of the episode, and gradually sees other relevant nodes (torso1 and torso2) to hold the movable ball.

Original: Furthermore, in ant_twisters (bottom), CGv2 firstly tries to raise the agent’s legs to satisfy goal1 and goal2, and then focus on reaching a leg (goal0).
Revised: Furthermore, in ant_twisters (bottom), CGv2 firstly tries to raise the agent’s legs to satisfy goal1 and goal2, and then focus on reaching a leg (goal0).

Original: Temporally-consistent attention to goal nodes and dynamics attention to relevant nodes can contribute to generalization over goal-directed tasks and morphologies.
Revised: Temporally-consistent attention to goal nodes and dynamics attention to relevant nodes can contribute to generalization over goal-directed tasks and morphologies.

Original: Figure 13 implies that CGv1 does not show such consistent activation to the goal-conditioned node; for instance, in centipede_touch_3 (above), the goal information is treated as an extra node feature of torso0, but there are no nodes that consistently activated with torso0.
Revised: Figure 13 implies that CGv1 does not show such consistent activation to the goal-conditioned node; for instance, in centipede_touch_3 (above), the goal information is treated as an extra node feature of torso0, but there are no nodes that consistently activated with torso0.

Original: Moreover, in ant_reach_handsup2_3 (bottom), CGv1 does not keep focusing on the agent’s limbs during the episode.
Revised: Moreover, in ant_reach_handsup2_3 (bottom), CGv1 does not keep focusing on the agent’s limbs during the episode.

Original: Rather, CGv1 tends to demonstrate some periodic patterns during the rollout as implied in prior works (Huang et al., 2020; Kurin et al., 2021). ��������� ����������������� �������������������� ���������to solve the task.
Revised: Rather, CGv1 tends to demonstrate some periodic patterns during the rollout as implied in prior works [49, 62]. ��������� ����������������� �������������������� ���������

Original: ��������� ����������������� ��������� ��������������������
Revised: ��������� ����������������� ��������� ��������������������

Original: L A DDITIONAL R ESULTS OF T ASK G ENERALIZATION
Revised: L Additional Results of Task Generalization

Original: Although in Table 1 and Figure 5, we examine the compositional task generalization and out-ofdistribution generalization with reach_hard tasks, where the goal distribution is farther than original reach tasks.
Revised: Although in Table 1 and Figure 5, we examine the compositional task generalization and out-ofdistribution generalization with reach_hard tasks, where the goal distribution is farther than original reach tasks.

Original: While they seem more difficult unseen tasks (Furuta et al., 2021b), they also seem to have some sort of task similarities between training dataset environments and those evaluation environments.
Revised: While they seem more difficult unseen tasks [35], they also seem to have some sort of task similarities between training dataset environments and those evaluation environments.

Original: Another question might be how control graph performs in more different unseen tasks.
Revised: Another question might be how control graph performs in more different unseen tasks.

Original: To evaluate compositional task generalization and out-of-distribution on the environments that have less similarity to the training datasets, we prepare push task, where the agents try to move the box objects to the given goal position.
Revised: To evaluate compositional task generalization and out-of-distribution on the environments that have less similarity to the training datasets, we prepare push task, where the agents try to move the box objects to the given goal position.

Original: See Table 7 for the environment division.
Revised: See Table 7 for the environment division.

Original: For training datasets, we leverage
Revised: For training datasets, we leverage

Original: In-Distribution division in Table 4.
Revised: In-Distribution division in Table 4.

Original: 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).
Revised: 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).

Original: Table 16 shows the results of compositional task evaluation and Table 17 shows those of out-ofdistribution evaluation.
Revised: Table 16 shows the results of compositional task evaluation and Table 17 shows those of out-ofdistribution evaluation.

Original: In contrast to Table 1, the zero-shot performance seems limited.
Revised: In contrast to Table 1, the zero-shot performance seems limited.

Original: Transferring pre-trained control primitives to significantly different tasks still remains as important future work.
Revised: Transferring pre-trained control primitives to significantly different tasks still remains as important future work.

Original: However, as shown in Figure 14, control 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.
Revised: However, as shown in Figure 14, control 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.

Original: As prior work suggested (Mandi et al., 2022), these results suggests that, in RL, jointlylearned multi-task model has a strong inductive bias even for unseen and significantly different environments.
Revised: As prior work suggested [73], these results suggests that, in RL, jointly-learned multi-task model has a strong inductive bias even for unseen and significantly different environments.

Original: [Table] [Table]
Revised: [Table] [Table]

Original:    
Revised:    

Original:  "WFSBHF/PSNBMJ[FE%JTUBODF 5'$(W	L  5'$(W	L  5'$(W	L  5'$(W	L  5'$(W	L  5'$(W	L  .-1	L  .-1	L  .-1	L  $PNQPTJUJPOBM	5BTL  GJOFUVOJOH SBOEJOJU    
Revised:  "WFSBHF/PSNBMJ[FE%JTUBODF 5'$(W	L  5'$(W	L  5'$(W	L  5'$(W	L  5'$(W	L  5'$(W	L  .-1	L  .-1	L  .-1	L  $PNQPTJUJPOBM	5BTL  GJOFUVOJOH SBOEJOJU    

Original:  "WFSBHF/PSNBMJ[FE%JTUBODF 0VUPG%JTUSJCVUJPO GJOFUVOJOH SBOEJOJU
Revised:  "WFSBHF/PSNBMJ[FE%JTUBODF 0VUPG%JTUSJCVUJPO GJOFUVOJOH SBOEJOJU

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Figure] Out [Figure] of

Original: [Figure]
Revised: [Figure]

Revised: Distribution

Revised: MLP (4K, randinit) 0.

Original: [Table] 
Revised: [Table] [Table] [Table] [Table]

Revised: [Table]