MetaMorph: Learning Universal Controllers with Transformers

Agrim Gupta; Linxi Fan; Surya Ganguli; Li Fei-Fei

Source-linked AI summary

MetaMorph: Learning Universal Controllers with Transformers

Agrim Gupta, Linxi Fan, Surya Ganguli, Li Fei-Fei

arXiv:2203.11931v1Mar 22, 2022cs.LGcs.NEcs.RO

TL;DR

Robotics often trains one robot for one task, while prior morphology-generalizing approaches have had limited diversity or complexity. MetaMorph uses a morphology-conditioned Transformer and joint pre-training across robot morphologies. The resulting policies generalize to unseen robot variations and support sample-efficient transfer, though the paper reports a performance drop for zero-shot generalization to new morphologies.

Problem
Robotics commonly trains a single robot for a single task, while prior morphology-generalizing work was limited in robot diversity, control complexity, or its assumed GNN inductive bias.
Source: §1 · p. 1
Method
MetaMorph jointly pre-trains a Transformer policy conditioned on module-level proprioceptive and morphology information, using dynamic replay buffer balancing across robots.
Source: §1 · p. 2
Results
Across all 3 environments, transfer showed strong zero-shot performance and 2 to 3 times higher sample efficiency than training from scratch.
Source: §5.4 · p. 9
Takeaways & Limitations
The authors identify zero-shot generalization to unseen robot variations as particularly useful where controllers need to be robust to hardware failures.
Source: §1 · p. 2
The approach assumes access to robots already optimized for locomotion, requiring a two-stage pipeline, and shows a performance drop in zero-shot generalization to new morphologies.
Source: C LIMITATIONS AND FUTURE WORK · p. 16

Abstract

from arXiv · show

Multiple domains like vision, natural language, and audio are witnessing tremendous progress by leveraging Transformers for large scale pre-training followed by task specific fine tuning. In contrast, in robotics we primarily train a single robot for a single task. However, modular robot systems now allow for the flexible combination of general-purpose building blocks into task optimized morphologies. However, given the exponentially large number of possible robot morphologies, training a controller for each new design is impractical. In this work, we propose MetaMorph, a Transformer based approach to learn a universal controller over a modular robot design space. MetaMorph is based on the insight that robot morphology is just another modality on which we can condition the output of a Transformer. Through extensive experiments we demonstrate that large scale pre-training on a variety of robot morphologies results in policies with combinatorial generalization capabilities, including zero shot generalization to unseen robot morphologies. We further demonstrate that our pre-trained policy can be used for sample-efficient transfer to completely new robot morphologies and tasks.

1 INTRODUCTION

Robotics largely follows a “one robot one task” paradigm, while modular systems offer a way to combine general-purpose building blocks into task-optimized morphologies. MetaMorph uses a Transformer conditioned on morphology to learn a universal controller, with the goal of generalizing across unseen robot variations and supporting transfer to new morphologies and tasks.

1 INTRODUCTION

Prior morphology-generalizing policies were limited to few manually varied robots, control complexity of at most 7 degrees of freedom, or a GNN inductive bias.The paper frames modular robot design spaces as an opportunity to bring large-scale pre-training and transfer learning to robotics.Source: §1 · p. 1
The target setting includes generalization to unseen dynamics and kinematics, plus sample-efficient transfer to new morphologies and tasks.The work instantiates this setting in UNIMAL, a 15–20-DoF design space for locomotion and mobile manipulation in complex stochastic environments.Source: §1 · pp. 1–2
MetaMorph conditions a Transformer on module-level proprioceptive and morphology information, jointly pre-training policies to generalize across robot embodiments.Dynamic replay buffer balancing adjusts experience collection according to each robot’s performance.Source: §1 · p. 2
The stated contributions include zero-shot generalization to unseen robot variations and attention-mask analysis that finds emergent motor synergies.The authors note that such generalization is particularly useful where controllers need robustness to hardware failures.Source: §1 · p. 2

2 RELATED WORK

Prior work on morphology-generalizing control has addressed parametric variation and limited sets of robot types, while other work studies designing or exploring modular morphologies. MetaMorph is positioned against these approaches by its diversity and scale of training robots, environment complexity, morphological conditioning, and generalization scope.

2 RELATED WORK

Earlier morphology-generalizing controllers mainly covered parametric variations of one or a few robot types, with GNNs used to handle differing action and state spaces across morphologies.Prior examples include planar agents with at most 7 DoFs; Kurin et al. (2021) also used Transformers for planar agents.Source: §2 · pp. 2–3
Work on modular design spaces includes task-specific morphology optimization, evolutionary search, RL for module parameters, and parameter-sharing controllers.Other approaches also explore morphologies without a task or reward specification and control self-reconfigurable modular robots.Source: §2 · p. 3

3 LEARNING A UNIVERSAL CONTROLLER

The paper formulates universal control over robots from a modular design space as multi-task reinforcement learning. It describes UNIMAL morphologies as articulated kinematic graphs and optimizes a shared policy across robot tasks.

3 LEARNING A UNIVERSAL CONTROLLER

Learning a universal controller over robots drawn from a modular design space is formulated as a multi-task reinforcement-learning problem.The optimization target is a policy shared across the robot set.Source: §3 · p. 3, §3.2 · p. 3

3.1 THE UNIMAL DESIGN SPACE

In UNIMAL, each robot morphology is a kinematic tree or directed acyclic graph whose vertices are modules and whose edges are motor-actuated joints.The root sphere represents the head; cylinders represent limbs, and joints are parameterized by axis, limits, and gear ratio.Source: §3.1 · p. 3

3.2 JOINT POLICY OPTIMIZATION

Each robot is modeled as an infinite-horizon discounted MDP, and the shared policy is optimized to maximize average expected return across tasks using PPO.The MDP tuple comprises states, actions, transition dynamics, reward, horizon, and discount factor.Source: §3.2 · p. 3

4 METAMORPH

MetaMorph addresses the infeasibility of separate policies across a vast modular robot design space by conditioning control on morphology. It combines module-level representations with global sensing and dynamically balances robot sampling during training.

4 METAMORPH

Separate policies are infeasible across exponentially many morphologies, so MetaMorph exploits shared modular structure to learn a universal controller.The UNIMAL design space is estimated to contain more than 1018 robots.Source: §4 · p. 3

4.2 MORPHOLOGY AWARE TRANSFORMER

Transformers parameterize the policy and critic; the policy encodes depth-first module tokens containing proprioceptive and morphological information, then combines encoder output with global sensing before per-joint decoding.Source: Figure 2 · p. 4, §4.2 · p. 4

4.1 FUSING PROPRIOCEPTIVE STATES AND MORPHOLOGY REPRESENTATIONS

Morphology serves as a task identifier, so the policy is conditioned on each robot's morphology rather than being morphology-agnostic.Source: §4.1 · p. 4

4.3 DYNAMIC REPLAY BUFFER BALANCING

Early episode lengths can favor passively stable robots, so dynamic replay balancing prioritizes robot sampling by performance; episode length is used in practice, and β=0 gives uniform sampling.Source: §4.3 · p. 5

5 EXPERIMENTS

The experiments evaluate MetaMorph across three locomotion environments, comparing joint pre-training with baselines and testing zero-shot generalization, transfer, and learned motor synergies. MetaMorph matches the per-morphology MLP average in two environments, generalizes to unseen variations, and transfers efficiently to new morphologies and tasks.

5.1 EXPERIMENTAL SETUP

The evaluation trains 100 UNIMAL robots in MuJoCo across flat, variable, and obstacle terrain, maximizing forward displacement over 1,000-timestep episodes.Variable terrain randomly interleaves hills, steps, and rubble; obstacles are cuboid-shaped.Source: §5.1 · p. 6

5.2 BASELINES AND ABLATIONS

MetaMorph matches the per-morphology MLP average in flat and obstacle environments, is slightly below it in variable terrain, and is 5× more sample-efficient than separate MLP training.The GNN baseline saturates 2–3 times below MetaMorph.Source: Figure 4 · p. 6, §5.2 · p. 7
Dynamic replay-buffer balancing is needed for good policies in 10–15% of robots across all three environments.Without balancing, those robots do not exceed the stated good-control threshold.Source: Figure 5 · p. 6, §5.2 · p. 7

5.3 ZERO-SHOT GENERALIZATION

Without dynamics randomization during training, MetaMorph still outperforms GNN and MetaMorph-NM across dynamics variations and environments, especially in variable terrain and obstacles.The training robots share armature and damping parameters.Source: Figure 6 · p. 7, §5.3 · p. 8
MetaMorph outperforms GNN and MetaMorph-NM across kinematic variations, though performance drops for the challenging unseen joint-angle variations.Joint-angle changes can alter possible motion and may require a different gait.Source: Figure 6 · p. 7, §5.3 · p. 8

5.4 SAMPLE EFFICIENT TRANSFER LEARNING

Fine-tuning pre-trained MetaMorph on new morphologies yields strong zero-shot performance and 2–3 times higher sample efficiency than training from scratch across all three environments.The comparison uses 100 test morphologies and averages reward progression over three runs.Source: Figure 7 · p. 8, §5.4 · p. 9
For new morphologies performing novel tasks, pre-training also gives strong zero-shot performance and 2–3 times greater sample efficiency than training from scratch.The transfers are variable terrain to Escape and obstacles to cylinder obstacles.Source: Figure 8 · p. 8, §5.4 · p. 9

5.5 EMERGENT MOTOR SYNERGIES

MetaMorph’s attention has low stable rank and oscillates between limb-group activation patterns, which the authors identify as a signature of motor synergies.They interpret this as the policy activating different motor synergies depending on the state.Source: §5.5 · p. 9

6 CONCLUSION

MetaMorph uses a Transformer to learn a universal controller for a large modular robot design space. The paper reports generalization to unseen variations, morphologies, and tasks after pre-training on diverse morphologies.

MetaMorph conditions a Transformer’s output on robot morphology to learn a universal controller for a large modular robot design space.Source: §6 · p. 9
Pre-training on diverse morphologies leads to policies that generalize to unseen variations in kinematics and dynamics, as well as new morphologies and tasks.Source: §6 · p. 9
The authors frame this work as a step toward large-scale pre-training and fine-tuning in robotics, a paradigm successful in vision and language.Source: §6 · p. 9

8 REPRODUCIBILITY STATEMENT

The authors released a PyTorch implementation of MetaMorph on GitHub, along with the robots and environments used for benchmarking.

The released GitHub repository includes a PyTorch implementation of MetaMorph.Source: §8 · p. 10
The repository also includes the robots and environments used for benchmarking.Source: §8 · p. 10

A IMPLEMENTATION DETAILS

The appendix describes the observation inputs and environment setup used for implementation. Observations combine per-module morphology and proprioception with global sensory information; the robot sets and kinematic variants are also specified.

A.1 INPUT OBSERVATIONS

Each module’s input includes morphology information, proprioceptive state, and additional global sensory information.Source: §4.1 · p. 4, §A.1 · p. 14
Morphology covers module shape, material, geometry, joint properties, and actuator properties; proprioception includes body pose and velocities plus joint position and velocity.Source: §A.1 · p. 14
Terrain input is a nonuniform 2D heightmap sampled from 1 m behind to 4 m ahead and 4 m to either side, with heights relative to the ground under the agent’s root.Source: §A.1 · p. 14

A.2 ENVIRONMENTS

The MuJoCo setup uses 100 selected UNIMALS from 300 evolved robots for training and a separate test set of 100, with no exact duplicate kinematic trees or parameters.Source: §A.2 · p. 14
Joint-angle variants were chosen to overlap the original range by at least 50%, helping prevent variants that could not be trained.Source: §A.2 · p. 14

A.3 TRAINING HYPERPARAMETERS

The training setup uses Transformers for policy and critic networks and PPO for joint training. The baselines and ablations use approximately matched parameter counts.

A.3 TRAINING HYPERPARAMETERS

Transformers parameterize both policy and critic networks, while a two-layer MLP encodes global sensory information.Source: §A.3 · p. 14
Agents are jointly trained with PPO, and dynamic replay buffer sampling uses α = 0.1 and β = 1.0.Source: §A.3 · p. 14
Baselines and ablations use the same hyperparameters and approximately 3.3 million parameters each.Source: §A.3 · p. 14

A.4 EVALUATION METHODOLOGY

Baseline and ablation experiments use three random seeds because RL performance can depend strongly on seed; zero-shot evaluation uses one pretrained seed, while transfer fine-tunes all three corresponding seeds.

A.4 EVALUATION METHODOLOGY

Three random seeds are used for baseline and ablation experiments, and their training curves show small standard deviations.Zero-shot evaluation uses a pretrained model from one seed; transfer fine-tunes models using each of the three pretraining seeds.Source: §A.4 · p. 15

B.1 BASELINES AND ABLATIONS

The appendix compares baselines and input ablations, examines learned position embeddings, and reports sensitivity to kinematic-tree traversal order. The findings highlight the value of morphology information and learned positional features, alongside an ordering sensitivity that can be mitigated with augmentation.

B.1 BASELINES AND ABLATIONS

MetaMorph-AO approximates Amorpheus inputs but adds morphology information, and the paper estimates Amorpheus would take about 30 GPU days versus 1.5 GPU days for MetaMorph.Direct comparison is not feasible because Amorpheus does not support variable joints between limbs in 3D morphologies, lacks exteroceptive observations, and collects experience and updates sequentially.Source: §B.1 · p. 15
MetaMorph outperforms the hand-designed position-encoding variant, while morphology-free variants perform significantly worse; the multi-task MLP also struggles to learn a good policy.MetaMorph-HPE is slightly better than MetaMorph-NPE, but still substantially worse than MetaMorph; the authors attribute the MLP difficulty to the diversity of the 100 training morphologies.Source: §B.1 · pp. 15–16, Figure 10 · p. 15

B.2 POSITION EMBEDDING

Learned position embeddings show a diagonal cosine-similarity pattern across seeds and environments, yet perform better than manually specified limb IDs, suggesting they encode additional useful information.The authors leave the precise information captured by these embeddings for future investigation.Source: §B.2 · p. 16, Figure 11 · p. 16

B.3 KINEMATIC TREE TRAVERSAL ORDER

Reversing the depth-first traversal order causes about a 75% zero-shot performance drop, while randomizing visitation among same-depth nodes during training makes the policy robust to such variations.The traversal convention visits same-depth nodes in MuJoCo XML order.Source: §B.3 · p. 16

C LIMITATIONS AND FUTURE WORK

The approach assumes access to a large collection of locomotion-optimized robots and therefore uses a two-stage process: create robots, then pre-train jointly. Future work includes more sample-efficient transfer and controllers for multiple skills.

LIMITATIONS AND FUTURE WORK

The current pipeline assumes a large collection of robots already optimized for locomotion, so robot creation and joint pre-training remain separate stages.Combining the two phases efficiently is identified as future work.Source: C LIMITATIONS AND FUTURE WORK · p. 16
Zero-shot generalization to new morphologies has a performance drop, motivating further work on sample-efficient transfer learning.Source: C LIMITATIONS AND FUTURE WORK · p. 16
The evaluated task suite focuses primarily on locomotion, leaving general-purpose multi-skill controllers as future work.Source: C LIMITATIONS AND FUTURE WORK · p. 16