Shared Modular Recurrence in Contextual MDPs
for Universal Morphology Control

Laurens Engwegen
Delft University of Technology
Delft, The Netherlands
l.r.engwegen@tudelft.nl
Max Weltevrede
Delft University of Technology
Delft, The Netherlands
Caroline Horsch
Delft University of Technology
Delft, The Netherlands
Daan Brinks
Delft University of Technology
Delft, The Netherlands
Wendelin Böhmer
Delft University of Technology
Delft, The Netherlands


Abstract

A universal controller for any robot morphology would greatly improve computational and data efficiency. Steps have been made towards such multi-robot control by utilizing contextual information about the properties of individual robots and exploiting their modular structure in the architecture of deep reinforcement learning agents. When the robots have highly dissimilar morphologies, however, this becomes a challenging problem, especially when the agent must generalize to new, unseen robots. In this paper, we posit that contextual features are often only partially available, but that they can be recovered through modular interactions. This can allow for better multi-robot control and generalization to contexts that are not seen during training. To this extent, we implement a transformer-based architecture with shared modular recurrence and evaluate its (generalization) performance on a large set of MuJoCo robots. The results show a substantial improvement in zero-shot generalization performance on robots with unseen dynamics, kinematics, and topologies, in four different environments.

1 Introduction↩︎

Reinforcement Learning (RL) has shown to be very promising for robotic control [1][3]. In an effort to close the gap with real-world applications, a lot of work has focused on RL agents that are able to generalize control to different tasks, e.g. manipulating different objects or acting in different environments. Recently, large datasets of robot trajectories have been established and are being used to train such generalizable agents by learning from demonstrations in an offline fashion with foundation models [4][6]. Several works show promising possibilities for a single model to adapt not only to different scenes and goals, but also to different embodiments that the model has seen demonstrations of [7], [8]. Zero-shot generalization to robots that were not seen during training, however, remains very challenging.

Different robot embodiments, or morphologies, can have advantages over each other, dependent on their task and the environments they act in. It is infeasible to train a policy from scratch for every robot: training (or even fine-tuning) a policy for every new robot that we are interested in requires expensive compute, excessive use of data, and often tedious hyperparameter tuning. A universal controller that can generalize control to any robot morphology could drastically improve efficiency. The UNIMAL design space [9] was developed for the evolution of diverse robots that can perform varied tasks. It contains more than 1000 different robots, a small subset of which is shown in Figure 1. The dataset of robots in the UNIMAL design space is suitable for the development and evaluation of RL agents that can generalize control to any robot.

By framing this problem as a multi-task RL problem [10], each robot can be considered as a separate task that has some specific contextual features, such as the mass of its different limbs or its topology. The goal in this multi-task setting is to learn a universal controller that can control any robot on basis of its observations and context. This is not only challenging because robots can have different action and state spaces, but also because robots with different morphologies and/or dynamics might learn tasks in different ways. The multi-task framework allows us to evaluate the performance of an agent during multi-robot training, and its zero-shot generalization performance [11] to new, unseen robots.

Figure 1: Example of robots that can be found in the UNIMAL design space [9].

Previous work on universal morphology control exploits contextual features that describe properties of the robot to better represent and distinguish tasks. In this paper, we reason that the context can be considered partially observable and that the available features do not encapsulate the true state needed for optimal (generalizable) control. We build upon modular architectures that were previously found to be effective for robotic control [12], [13], and address the essentially unobservable context with shared modular recurrence to retain the system’s compatibility with different action spaces while minimizing additional compute. This paper empirically shows that the proposed system improves multi-robot control on varied training morphologies, and enables better zero-shot generalization to unseen test robots. Code to reproduce all of the experiments is publicly available1.

2 Background↩︎

2.1 Contextual Markov Decision Process↩︎

Here, the problem of learning RL policies that are trained on a set of training robots and must generalize to unseen test robots is considered. This problem can be formulated as a Markov Decision Process (MDP) where the agent can only partially access the MDP during training. An MDP is a tuple \(( \mathcal{S}\), \(\mathcal{A}\), \(\mathcal{T}\), \(\mathcal{R}\), \(\rho )\), with state space \(\mathcal{S}\), action space \(\mathcal{A}\), transition function \(\mathcal{T}(s_{t+1}|s_t,a_t)\), mapping current state \(s_t \in \mathcal{S}\) and action \(a_t \in \mathcal{A}\) to a probability distribution over next states \(s_{t+1} \in \mathcal{S}\), reward function \(\mathcal{R}: \mathcal{S} \times \mathcal{A} \rightarrow \mathbb{R}\) and initial state distribution \(\rho (s_0)\). A Contextual Markov Decision Process (CMDP [14]) is an MDP where states can be decomposed, \(s_t = (s_t', c)\), into the underlying state \(s_t' \in \mathcal{S'}\) and a context \(c \in \mathcal{C}\). The context is sampled at the start of each episode and remains static until the episode ends. In multi-robot control, the context defines the robot to be controlled. This framework allows to evaluate zero-shot generalization, that is, generalization to contexts not seen during training [11].

The UNIMAL design space contains modular robots [9] that are simulated with the MuJoCo physics engine [15]. Such robots consist of a set of nodes, or limbs, that share the same state space and action space, i.e. \(\mathcal{S} = \{\mathcal{S}^i | i=1,...,N\}\) and \(\mathcal{A} = \{\mathcal{A}^i | i=1,...,N\}\) for robots with \(N\) limbs. Besides the underlying states, a node-level context, \(\mathcal{C} = \{\mathcal{C}^i | i=1, ..., N\}\), is provided in the state that describes information of the limbs (e.g. mass, initial position with respect to the parent limb, and the initial position of each joint attached to the limb). The node-level observation and context features that are provided to the agent are listed in Table 1 in Appendix 7. Previous methods exploited this modular structure in combination with modern architectures for effective multi-robot control [12], [13], [16], [17]. Rather than aiming to develop context-agnostic agents, [12] and [13] condition the agent on the available contextual information with methods coined MetaMorph and ModuMorph, respectively. In doing so, they did not only show improved performance during training, but also on generalization to unseen robots. However, the gap between the performance on training and testing robots remains substantial.

2.2 Partially Observable Context↩︎

In this paper, we recognize and leverage the fact that not all task-relevant contextual information is observable to the agent. In multi-agent control settings like MuJoCo, the missing context consists of structural and abstract properties that are difficult to encode directly in observations or cannot be processed effectively in the agent architecture. First, structural information about the robot is not directly accessible. While the agent receives local positional relationships (e.g., a limb’s position relative to its parent), the true adjacency structure of the morphology is not explicitly provided. As noted by [13], encoding such graph structure in modular architectures is not trivial due to challenges with positional encoding. Consequently, topology must be inferred indirectly. Moreover, some of the most critical contextual features are inherently abstract. For example, the functional influence that one limb exerts on another during coordinated movement is not directly observable nor easily encoded. Such effects of interaction cannot be captured through static observations alone.

Although the underlying state of a single robot in MuJoCo is treated as fully observable [15], the CMDP is more accurately modeled as partially observable [18]. Formally, the emission or observation function defined in partially observable MDPs [19], \(\phi: \mathcal{S} \rightarrow \mathcal{O}\), maps a state to an observation \(o_t \in \mathcal{O}\) that contains the underlying state and the observable context, \(c^+\), for every timestep \(t\): \(o_t = \phi ((s_t', c)) = (s_t', c^+)\). The unobserved portion of \(c\) can include both missing physical parameters and abstract features. In the single robot setting, such information never changes between episodes and can therefore effectively be treated as fully observable. Besides, even in the multi-robot setting, aforementioned methods that rely solely on explicitly provided context features can achieve good performance on training robots: the agent can simply overfit. The challenge arises when the agent must generalize to robots with (abstract) unobservable context features that were not present in the training set.

Incorporating memory into the agent architecture [20] can allow the agent to implicitly identify and quickly adapt to latent contextual information online. Because the underlying physical state is otherwise observable, any information inferred from history must be from the missing context. Crucially, such memory mechanism must preserve the modular structure of the agent’s architecture to remain compatible with the multi-robot setting in which input dimensionalities can vary. E.g., a naive encoding of the global action-observation history is not directly compatible with the state-of-the-art modular architectures.

3 Related Work↩︎

3.1 Universal Morphology Control↩︎

We build upon previous work aimed at learning an RL policy that can generalize control to different robot morphologies, even when state and action dimensionalities can change. Effective approaches utilized the modularity in robots [21] and introduced weight sharing across different modules [16]. Several works adopted graph neural networks [22], [23] or transformers [17], [24] as inductive biases to more explicitly infer relationships between different limbs or modules through message-passing. More recently, multi-robot training has been evaluated on a larger scale of different morphologies with the introduction of the UNIMAL design space [9]. [12] constructed training and testing sets of robots to evaluate multi-robot control, and showed the effectiveness of a modular transformer-based approach when contextual information is provided to the agent. By further exploiting this contextual information in the agent’s architecture, [13] demonstrated improved performance.

Steps towards robotic applications in the real world also exploit such modular transformer-based architectures for generalizable and transferrable control. Several recent works incorporate additional information about the robot topology through a graph encoding or sparse attention matrices [25], [26] for improved generalization in simulation and the real world. Fine-tuning and policy distillation approaches have also been proposed for generalization purposes [27], [28]. Lastly, more complex modular architectures were introduced for effective transfer to unseen robots in simulation and the real world [29], [30], although evaluating on smaller sets of, and relatively similar robot morphologies. None of these approaches consider this generalization problem as partially observable. Earlier works did suggest that “implicit system identification” with memory-based policies can benefit robotic control, but did not utilize a modular system, nor evaluate generalization on varied morphologies [31], [32]. We are the first to show that the addition of shared modular recurrence can improve zero-shot generalization on a diverse set of robot morphologies.

3.2 Neural Architectures↩︎

The effectiveness of the transformer architecture [24] for multi-robot control, lies in its capability to model pairwise dependencies between limbs with self-attention. Self-attention can be defined as \(A = \sigma (QK^T/\sqrt{d}) V\), with query, key and value matrices \(Q,K,V \in \mathbb{R}^{N \times d}\), for robots with \(N\) limbs and a hidden size of \(d\). Learnable parameters \(W_Q, W_K\) and \(W_V\) map the input \(X \in \mathbb{R}^{N \times d}\) to those matrices, i.e. \(Q = XW_Q\), \(K = XW_K\) and \(V = XW_V\), and \(\sigma(\cdot)\) is a row-wise softmax function. In addition to the attention mechanism, [13] utilize hypernetworks [33] to more explicitly condition the agent on the available node-wise contextual information. Briefly, they train a hypernetwork that is conditioned on the observable context to generate (1) the parameters of a node-wise encoder that produces \(X_V\) with which the value in the subsequent transformer encoder layer is calculated as \(V = X_VW_V\), (2) \(X_Q\) and \(X_K\), where the query and key matrices are now defined as \(Q = X_QW_Q\), \(K = X_KW_K\), respectively, and (3) the parameters of a node-wise decoder that projects the output of the transformer encoder. Recurrent neural networks (RNNs), like LSTM and GRU [34], [35], are useful architectures to deal with partially observable domains in deep RL. [36] combined LSTMs with attention to develop systems that can adapt to changes (permutations) of the input. Here, we build on these methods to develop a transformer-based architecture with shared modular recurrence and hypernetworks to improve multi-robot control.

4 Shared Modular Recurrence↩︎

4.1 Recurrent PPO↩︎

We aim to learn a universal control policy that is effective for any robot we can encounter in the UNIMAL design space by training on a limited set of \(K\) robots. An effective approach to RL in partially observable domains is to learn an encoding of the belief over the agent’s true state. This is often done by, at every timestep \(t\), encoding the action-observation history (AOH) \(\tau_t^k = (o_0^k, a_0^k, \dots, o_{t-1}^k, a_{t-1}^k, o_t^k)\) with an RNN, for (in the current domain) robot \(k\) the agent is controlling. In this way, the training objective can be formulated as finding parameters \(\theta\) for policy \(\pi_{\theta}(a_{t}^k | \tau_{t}^k )\) that maximize the (discounted) cumulative reward, averaged over all training robots: \(\mathrm{max}_{\theta} \frac{1}{K} \sum_{k=1}^{K} \mathbb{E}_{\pi_\theta}[\sum_{t=0}^{H} \gamma^{t}r_{t}^{k}]\), with task horizon \(H\). We implemented a recurrent version of Proximal Policy Optimization (PPO) [37] to optimize this objective.

To make recurrent training tractable at scale, we incorporate R2D2-style sequence-chunking and burn-in [38], originally proposed for (off-policy) DRQN [20], on the PPO rollout buffer. Rather than storing complete episodes for PPO updates during rollout, which is very expensive due to variable episode lengths up to \(H=1000\), we store overlapping trajectory chunks. For each chunk of transitions, the RNN hidden states are initialized with the hidden states that are stored during rollout. After that, a burn-in phase is applied to reconstruct meaningful hidden states for the current parametrization. We use a chunk size of \(m=80\) and a burn-in period of \(l=20\), as those values were reported by [38] to be effective. A default recurrent PPO setup for universal morphology control settings is insufficient: naively encoding global action-observation histories breaks both the compatibility with varying input dimensionalities and the structural invariances required for generalization. This motivates our architectural contribution.

4.2 Shared Recurrent Network↩︎

In order to retain the agent’s modularity, we adopt and adapt a recurrent architecture that processes components of the input separately [36]: every limb-level action and observation are processed individually through an RNN to encode local AOHs \(\tau_t^{i} = (o_0^{i}, a_0^{i}, \dots, o_{t-1}^{i}, a_{t-1}^{i}, o_t^{i})\) for every limb \(i\) (we omit the superscript that indicates the robot as our policy is controlling only one robot at a time). Since nodes share the same state space, the parameters of the RNN can be shared to increase the scalability of this approach. We only keep track of the hidden states of all limbs, \(h_t^{i} = \mathrm{RNN}(o_t^i, a_{t-1}^i, h_{t-1}^i)\), that encode the local AOHs, \(\tau^{i}\), which are initialized with zeros at the start of an episode. In this way, the agent can approximate the relevant history, necessary for implicit inference of the true state, for each limb individually. There are various ways in which this modular recurrence can be incorporated in the agent’s architecture. Here, we build upon state-of-the art architectures for universal morphology control: MetaMorph [12] and ModuMorph [13].

Figure 2: Recurrent MetaMorph (R-MeMo)

4.2.0.1 Recurrent MetaMorph.

MetaMorph exploited a modular transformer-based architecture in which limb-wise observations are processed through a shared embedding layer, before applying multi-head attention. The resulting embeddings are mapped to actions with a shared decoder. As shown in Figure 2, our recurrent version of MetaMorph (R-MeMo) first encodes the underlying state and previous action with an RNN. This is done separately for each limb, but with shared RNN parameters. The hidden state for each limb at the start of each episode chunk is stored in the rollout buffer for initialization during training. The observable context is mapped with an MLP, as this part of the observation remains constant throughout an episode. The embeddings are added together to establish a latent representation of \(\tau_t\).

Figure 3: Recurrent ModuMorph (R-MoMo)

4.2.0.2 Recurrent ModuMorph.

ModuMorph was built on MetaMorph by introducing a hypernetwork that takes context as input and generates the parameters for the embedding layer, the query and key of multi-head attention, and the parameters for the decoder layer. Our recurrent ModuMorph version (R-MoMo), shown in Figure 3, uses the same hypernetwork and fixed attention as in the original architecture, but incorporates the shared modular RNN before the transformer to approximate local AOHs from latent encodings of the observation and previous action. The RNN could also be inserted before the embedding layer. This would, however, result in larger inputs for the embedding layer (if we keep the size of the hidden states the same): 128 instead of 54 per limb. This, in turn, leads to an increase of more than 1M trainable parameters in the hypernetwork with respect to ModuMorph. This was avoided to ensure a fair comparison between ModuMorph and R-MoMo.

In both proposed architectures, a transformer receives the local (AOH) encodings from the RNN to infer relationships between different limbs and finally map them to limb-wise actions. The retained modular structure allows for robots with different input dimensionalities (i.e., different numbers of limbs). The introduction of shared modular recurrence only causes a modest increase in trainable parameters, as shown in Table 3 in Appendix 8. Augmenting MetaMorph and ModuMorph with MLPs of similar size caused stability problems, demonstrated in Figure 11 in Appendix 8, and was therefore avoided. Besides, the individual runs with additional parameters that did show stable learning did not achieve better performance.

5 Experiments↩︎

In this Section, experiments are performed with MetaMorph, ModuMorph, and their recurrent counterparts R-MeMo and R-MoMo, respectively. We use the MetaMorph version from [13], which they report to perform better than the original implementation.

5.1 Experimental setup↩︎

The training set of 100 robots, as constructed by [12], is used to train agents for multi-robot control. We then first evaluate the agent’s generalization performance on unseen variations of these training robots, where parameters that influence the dynamics and kinematics are altered (such as the damping of limbs or the angles joints can make). Subsequently, the performance on robots with unseen topologies is evaluated. The provided test set of robots with unseen topologies is randomly split into a validation (32 robots) and test (70 robots) set to experiment with different hyperparameters and evaluate generalization performance.

We only tuned one regularization hyperparameter that [13] found to have a big impact on performance. See Appendix 8 for further details. All other hyperparameter values were taken from [13] and are listed in Table 2. Parameters of the introduced shared modular recurrence were assigned default values (and, thus, not optimized), with a hidden and cell state size of 128 for the LSTM that was used as RNN. Ablations in which the LSTM is replaced by a GRU, and with the shared modular recurrence in the hypernetwork can be found in Appendix 9.

Agents are trained and evaluated in four different environments. In each of those, the agent has to maximize the robot’s locomotion distance. In Flat Terrain, the agent needs to traverse a flat surface, while in Incline the robots are to be controlled on a surface that is inclined by 10 degrees. Variable Terrain contains a sequence of hills, steps and rubble, interleaved with flat terrain. Those sequences are randomly generated at the start of each episode. Finally, Obstacles is a flat terrain with randomly generated obstacles. In the latter two environments, the agent receives a 2D heightmap of its close surrounding, in addition to proprioceptive and contextual observations, to be able to react to changes in terrain. For more details on the environments, we refer to [9].

Figure 4: Training performance on the 100 training robots in the different environments. The average return with a 95% confidence interval (CI) over 10 seeds is visualized.

5.2 Multi-Robot Training Performance↩︎

The training performance of the different methods on the 100 training robots, averaged over 10 seeds, is shown in Figure 4. Across all environments, the recurrent architectures (R-MeMo and R-MoMo) obtain either similarly high or higher returns than their non-recurrent baselines MetaMorph and ModuMorph. Particularly in the Incline environment, which is more difficult as dynamics play a more important role, shared modular recurrence results in better training performance. In general, ModuMorph seems to perform better than MetaMorph and R-MeMo on the training robots, illustrating the effectiveness of the hypernetworks conditioned on the available context during multi-robot training. Due to the sequential processing of RNNs, the training time of R-MeMo and R-MoMo increases with respect to their baselines, which is reported in Appendix 10. However, this is not the case during inference, as the agent then always processes observations sequentially.

5.3 Zero-Shot Generalization to Different Dynamics and Kinematics↩︎

[12] constructed a set of robots that have the same topologies as the training robots, but differ in a contextual feature, to evaluate zero-shot generalization to different dynamics or kinematics.

Figure 5: Performance on robots with changes in contextual features that result in unseen dynamics or kinematics. The average return with a 95% CI over 10 seeds is reported.

For each training robot, they created four test robots with variations in armature, damping, gear, density, limb shapes or joint angles, resulting in a new set of 2400 test robots. Figure 5 shows the performance of the four evaluated methods on the test robots with unseen dynamics and kinematics parameters. Over all changes, R-MoMo obtains a higher average return than ModuMorph, consistently throughout the different environments. Only in Obstacles confidence intervals overlap. Similarly, R-MeMo performs better than MetaMorph, and in some environments and parameter changes even outperforms ModuMorph. These results demonstrate improved generalization performance obtained through the shared modular recurrence across environments with different levels of complexity.

Figure 6: Test performance on robots with unseen topologies in four different environments. The average return with a 95% CI is shown.

5.4 Zero-Shot Generalization to Unseen Robot Topologies↩︎

[12] constructed training and test robot topologies, \(\mathcal{C}_{\mathrm{train}}\) and \(\mathcal{C}_{\mathrm{test}}\), for which \(\mathcal{C}_{\mathrm{train}} \cap \mathcal{C}_{\mathrm{test}} = \emptyset\). These robots are very diverse, as they have a variable amount of limbs (up to 12), with different context features, and a variable amount of joints per limb (up to 2). After training on the 100 training robots, the methods were evaluated on the 70 test robots with unseen topologies. The averaged returns of the different methods in the four environments are shown in Figure 6. In each of the environments, ModuMorph and/or R-MoMo dominate training, but R-MoMo significantly outperforms the other methods in zero-shot generalization to the test robots. Additionally, R-MeMo is competitive with ModuMorph on the unseen test robots, even though ModuMorph performs better during training. These results show that architectures with shared modular recurrence can learn policies that generalize much better to unseen test robots than their non-recurrent baselines across environments with varying levels of difficulty.

a

b

Figure 7: The difference in return between R-MoMo and ModuMorph (top) and the obtained returns (bottom) on all test robots. The returns are normalized for each environment and averaged over 10 seeds. The green dotted line indicates the average performance improvement..

Returns in each environment can range from very low (\(<0\)) to very high (\(>4000\)) values. Averaged returns over 70 test robots could therefore be misleading, as differences can be caused by only a small set of test robots. We therefore additionally report the difference in return between R-MoMo and ModuMorph for every test robot individually, next to the obtained returns, in Figure 7. To compare the methods across all environments, the average returns are normalized using the minimal (\(52\), \(-74\), \(71\), \(126\)) and maximal (\(5008\), \(3751\), \(3000\), \(2403\)) returns found in Flat Terrain, Incline, Variable Terrain, and Obstacles, respectively. Per-robot performance comparisons for every environment separately can be found in Appendix 10. This comparison shows a consistent improvement in a majority of the test robots, rather than only on a limited set of robots. Additionally, robots that ModuMorph performs better on are often also well controlled by R-MoMo. In contrast, ModuMorph struggles to control various robots across environments, while R-MoMo shows substantially fewer robots that are poorly controlled. Analyses in Appendix 10 validate the consistency of the improved policy throughout the entire episode (rather than only at certain timesteps) for the agents with shared modular recurrence.

5.5 Training Set Size↩︎

Figure 8: Test performance with smaller training set sizes, averaged over 10 seeds with 95% CI error bars.

To further validate the increased generalization performance and scalability of the introduced approach, we performed experiments with different training set sizes. Figure 8 shows the return obtained by the different methods for different amounts of robots in the training set: 25, 50, 75, and the complete set of 100 training robots in Flat Terrain and Incline. The improvement in zero-shot generalization performance of R-MoMo over ModuMorph is consistently visible for the different sizes of training sets, especially from 50 robots onwards. The difference increases when more training robots are available, which is specifically demonstrated in the more difficult Incline environment. This shows the consistency and scalability of R-MoMo with respect to the amount of training examples.

5.6 Single Contextual Features↩︎

We experimented with a scenario in which a very limited number of contextual features would be available to the agent, to find potential differences in robustness against this lack of information. ModuMorph and R-MoMo were trained and evaluated under this constraint. Test performance using single contextual features (or using all available features) in Flat Terrain and Incline is shown in Figure 9. Large error bars (e.g. for ModuMorph with \(\mathrm{body\_mass}\) in Flat Terrain) correspond to highly instable performance across different random seeds. These results show that R-MoMo frequently outperforms ModuMorph in the single feature setting. This observation is more evident in Incline, the more difficult environment. However, this advantage is not consistent across all features: some features are insufficiently informative to induce strong performance when used in isolation for both methods. When the provided contextual feature is informative, R-MoMo’s performance accelerates more than ModuMorph’s. This sensitivity likely stems from the hypernetwork, which conditions exclusively on the available context. Notably, ModuMorph’s performance in the sparse context regime does not drop for the first six features in both environments. A possible explanation is overfitting to the training robots, reducing the effect of the available feature. Future work could investigate such sparse context availability and feature importance in more detail.

a

b

Figure 9: Test performance of ModuMorph and R-MoMo on Flat Terrain and Incline when providing either all available contextual features or only a single contextual feature. Mean performance and standard deviation over 5 seeds is shown..

6 Discussion↩︎

This work explores the introduction of shared modular recurrence in transformer-based architectures for improved multi-robot control. It was hypothesized that modular recurrence could improve performance by enabling the agent to infer relevant latent context from local histories. The results have shown a consistent increase in both training and zero-shot generalization performance when such memory mechanism was introduced across different environments for robots with different dynamics, kinematics, and topologies. This clearly indicates that the RNN extracts some unobservable contextual information from the history for better policies. Moreover, the improvement is consistent over a large set of test robots and for agents trained on varying amounts of training robots.

Experiments in which only a single context feature is provided did show an increase in performance, although not consistently for all features. Specifically in a more difficult environment, informative contextual features can enable the recurrent architecture to learn better policies. The results indicated that the implicit inference of more contextual information is dependent on the quality or informativeness of available context features. This is an interesting observation, as it implies that learning of partially observable contexts greatly benefits from observing some features, but not others. We leave it to future work to characterize what kind of features assist this inference best, and to what extent this is an artifact of the architecture’s dependence on the hypernetwork.

A limitation of the explored recurrent architecture, is that hidden states have to be stored for each limb. Scaling to robots with a large number of limbs requires, therefore, more memory. Nonetheless, using episode chunking with a burn-in period during training only requires storing hidden states at the start of a chunk, minimizing the additional memory usage. An interesting direction for future research would be to investigate a more efficient memory mechanism in the architecture. Secondly, the sequential processing of RNNs results in longer training times. This disadvantage only holds for training, due to the inherent sequential processing during inference. Lastly, there is still room for improvement regarding the generalization gap. Nonetheless, the combination of shared modular recurrence with transformers has shown to be promising for multi-robot control and could be effective in other problems with partially observable contexts in graph-like structures.

This project was funded by the Delft AI initiative within the BIOlab. This work was also partially funded by the Dutch Research Council (NWO) project Reliable Out-of-Distribution Generalization in Deep Reinforcement Learning with project number OCENW.M.21.234. DB acknowledges support from an ERC Starting Grant (850818 - MULTI-VIsion) and the Convergence Health and Technology (Integrative Neuromedicine flagship).

7 Observations and Context↩︎

The observation provided to the agent in the current environments consists of various features for each limb. Table 1 lists all of these features with a short description, the dimensionality of the feature and whether the feature is part of the context. We refer to the MuJoCo documentation [15] for more details.

Table 1: Features of observations and context with description and dimensionality. \(*\) indicates that each limb contains this feature twice, as every limb can contain two joints.
Feature Description Dim. Context
\(\mathrm{body\_xpos}\) Current x,y,z position of each limb 3 No
\(\mathrm{body\_xvelp}\) Linear velocity of each limb 3 No
\(\mathrm{body\_xvelr}\) Angular velocity of each limb 3 No
\(\mathrm{body\_xquat}\) Orientation of each limb 4 No
\(\mathrm{qpos}\) Generalized coordinates of each joint 1* No
\(\mathrm{qvel}\) Generalized velocity of each joint 1* No
\(\mathrm{body\_pos}\) Initial x,y,z position of limb w.r.t. parent limb 3 Yes
\(\mathrm{body\_ipos}\) Initial x,y,z position of center of mass w.r.t. parent limb 3 Yes
\(\mathrm{body\_iquat}\) Inverse quaternion of limb orientation 4 Yes
\(\mathrm{geom\_quat}\) Quaternion of geom relative to the body 4 Yes
\(\mathrm{body\_mass}\) Limb mass 1 Yes
\(\mathrm{body\_shape}\) Limb shape 2 Yes
\(\mathrm{jnt\_pos}\) Initial (x,y,z) coordinate of each joint 3* Yes
\(\mathrm{joint\_range}\) Range of motion (lower and upper bound) of each joint 2* Yes
\(\mathrm{joint\_axis}\) Axis of rotation/translation of each joint (one-hot for x,y,z) 3* Yes
\(\mathrm{gear}\) Gear ratio for each joint 1* Yes

a

b

c

d

Figure 10: Validation performance on the 32 robots in the validation set for different values of the regularization parameter that defines the maximum approximate KL-divergence between the old and the new policy, averaged over 10 seeds with 95% CIs..

8 Implementation Details↩︎

8.1 Regularization↩︎

In our experiments, we use the same hyperparameter values as in ModuMorph (see Table 2). We only evaluate two different values for a regularization parameter on the validation set. [13] namely argued that this parameter can have a big impact on performance. This parameter defines the maximum approximate KL-divergence between the old and the new policy for every mini-batch before the update step. If this value is exceeded, the iteration of updates ends, and we return to sampling new trajectories. Figure 10 shows the average performance of the different methods with the two validated values (\(3 \times 10^{-2}\) and \(5 \times 10^{-2}\)) that were also evaluated in ModuMorph. In most cases, there is not a big difference in performance. For every method, the value that resulted in the highest average return on the validation set was used for the reported results in Section 5.

8.2 Hyperparameters↩︎

[13] provide a link to the datasets of robots and to their implementation on which we built. Table 2 lists all the hyperparameter values that are used for all four methods that are considered in this paper. These values correspond to those used by [13]. At the start of each episode rollout, a robot is sampled following the dynamic replay buffer balancing scheme introduced by [12] (and used in both MetaMorph and ModuMorph). Note that the RNN is only used in R-MeMo and R-MoMo.

Table 2: Hyperparameters used in the experiments with all methods.
Hyperparameter Value
Environment Total timesteps \(\{1\times10^8, 2\times10^8\}\)
Vectorized environments 32
PPO Discount factor (\(\gamma\)) 0.99
GAE parameter (\(\lambda\)) 0.95
PPO clip range (\(\epsilon\)) 0.2
Reward normalization Yes
Observation normalization Yes
Batch size 5120
Timesteps per rollout 2560
Updates per rollout 8
Gradient clipping 0.5
Clipped value function Yes
Value loss coefficient 0.5
Type Adam
Epsilon \(1 \times 10^{-5}\)
Initial learning rate \(3 \times 10^{-4}\)
Learning rate schedule Linear warmup, cosine decay
Warmup iterations 5
Transformer Number of layers 5
Number of attention heads 2
Embedding dimension 128
Feedforward dimension 1024
Activation function ReLU
Dropout 0.0
RNN Type LSTM
Number of layers 1
Embedding dimension 128
Dropout 0.0

8.3 Increased Trainable Parameters↩︎

Table 3 shows the number of trainable parameters for all of the methods. The shared modular recurrence introduces a small increase. We experimented with versions of MetaMorph and ModuMorph with a similar number of additional trainable parameters. For MetaMorph, the first embedding block is increased by two fully connected layers (with a hidden size of 256 and ReLU activation). In ModuMorph these layers were added after the first embedding block, where the RNN is placed in R-MoMo. Figure 11 shows that the extra trainable parameters cause instability issues during training. Therefore, we stick to the architectures with the original number of parameters for the other experiments reported in this paper.

Table 3: Number of trainable parameters for each method.
Trainable parameters
MetaMorph 3,312,515
R-MeMo 3,568,259
ModuMorph 5,232,774
R-MoMo 5,563,526

a

b

Figure 11: Training performance of MetaMorph and ModuMorph with additional trainable parameters (MetaMorph+ and ModuMorph+), comparable to the amount of additional parameters in the recurrent architectures, in Flat Terrain. The average return over 5 seeds with 95% confidence intervals is shown on the left. On the right, individual seeds are shown for MetaMorph+ and ModuMorph+, which demonstrates that none of the stable runs perform better than the average of the baselines..

9 Ablations↩︎

Figure 12: The HyperR-MoMo architecture variation, in which the AOH representation is learned in the hypernetwork instead of in the base network.

9.0.0.1 Modular Recurrence in Hypernetwork.

Instead of inserting the RNN (here: LSTM) in the base controller, it can also be introduced in the hypernetwork to learn to infer latent contextual information. In this case, we add the AOH representation to the mapped observable context, to form the input for the blocks that generate the embedding layer, the query and key, and the decoder layer. The original RNN in the base controller (i.e. before the Transformer) is removed. We refer to this variation as HyperR-MoMo and illustrate it in Figure 12. Note that the original hypernetwork only required a forward pass at the start of an episode, since the context does not change during an episode. HyperR-MoMo’s hypernetwork requires a forward pass at every timestep and can, thus, generate different parameters at every timestep.

9.0.0.2 Gated Recurrent Unit.

To assess the robustness of the performance gains with shared modular recurrence and to evaluate whether they are specific for LSTMs, we experimented with a GRU. Leaving the rest of the architecture unchanged, the LSTM is swapped with a GRU in R-MoMo. The training performance in Figure 13 and test performance in Figure 14 in Flat Terrain and Incline show that this GRU-MoMo is competitive with R-MoMo that uses an LSTM. These results demonstrate that the performance gains are not dependent on the usage of LSTM and provide additional evidence of the necessity of inferring latent contextual features from the history.

Figure 13: Training performance on the 100 training robots in Flat Terrain and Incline for ModuMorph, HyperR-MoMo, GRU-MoMo (R-MoMo with a GRU), and R-MoMo (with an LSTM). Average return with a 95% confidence interval over 10 seeds is shown.
Figure 14: Test performance on the 70 unseen test robots in Flat Terrain and Incline for ModuMorph, HyperR-MoMo, GRU-MoMo (R-MoMo with a GRU), and R-MoMo (with an LSTM). Average return with a 95% confidence interval over 10 seeds is shown.

The training and zero-shot generalization performance of this architecture variation in Flat Terrain and Incline is shown in Figures 13 and 14, respectively. In Flat Terrain, notably, HyperR-MoMo obtains lower training performance, even though its average test performance is between ModuMorph and R-MoMo. In Incline, HyperR-MoMo shows the lowest training performance, and a similar test performance as ModuMorph. The lower performance and large variance across runs is due to a few runs in which the agent did not learn at all. This could be caused by larger instability due to the fact that the base network can contain a different set of parameters at every timestep (instead of the same set of parameters for each individual robot, as previously discussed). These results indicate that this architecture variation might require additional hyperparameter optimization to reduce instability, but could be promising for universal morphology control. Specifically because its generalization performance is still better than ModuMorph.

10 Comparisons↩︎

10.1 Reward Per Timestep↩︎

To qualitatively ensure that the observed performance gains are not the result of a limited number of transitions with a large reward (for example at the very beginning of the episode), we analyzed the reward obtained by the agent at every timestep during an episode. In Figure 15, the obtained average reward per timestep for all test robots is shown for the first 500 timesteps of the episode. In general, the reward slowly increases at the start of the episode when the agent has to apply torques to induce locomotion. In Variable Terrain and Obstacles, the average reward per timestep decreases as soon as the terrain changes or obstacles arise, as these slow down the robot. The differences between the methods demonstrate that the performance improvement with shared modular recurrence over the baselines is consistent over timesteps during the episode.

Figure 15: The reward obtained at every timestep in the first half of the episode in the different environments. The obtained rewards are averaged over all test robots and 5 seeds and the standard deviation is shown.

Additionally, to investigate the agent’s usage of the RNN hidden states, we performed the same analysis for recurrent agents that receive hidden states consisting of only zeros during the first 100 timesteps of each episode. In Figure 16 we can now see that the recurrent agents execute a weak policy during these timesteps, but are able to recover up to the original reward per timestep observed in Figure 15. This demonstrates that the hidden states are encoding crucial information.

Figure 16: The reward obtained at every timestep in the first half of the episode in the different environments, where the recurrent agents receive zeros as hidden states during the first 100 timesteps. The obtained rewards are averaged over all test robots and 5 seeds and the standard deviation is shown.

10.2 Wall Clock Training Time↩︎

In Table 4 the wall clock time of the different methods for 100M environmental steps and training in Flat Terrain is reported (note that this includes logging to different destinations and storing model checkpoints). Experiments were performed on NVIDIA A40 GPUs. Chunking the collected episodes before training, and applying the burn-in strategy introduced by [38] for DRQN reduces the extra overhead of sequential processing by the RNNs. Still, training time is increased significantly. However, this difference disappears during inference as data is then processed sequentially for all methods. The trade-off between performance and the reduction of the chunk and/or burn-in size could be further exploited to reduce training time, which we leave for future work.

Table 4: The wall clock time of the different methods for 100M timesteps of training on Flat Terrain, averaged over 10 seeds.
Wall clock training time
MetaMorph \(16.3 \pm 0.5\) hours
R-MeMo \(21.8 \pm 0.9\) hours
ModuMorph \(21.9 \pm 0.6\) hours
R-MoMo \(26.0 \pm 1.1\) hours

10.3 Zero-Shot Generalization Performance Improvement Per Environment↩︎

The difference in performance on test robots between R-MoMo and ModuMorph is shown in Figures 17, 18, 19 and 20 for the Flat Terrain, Incline, Variable Terrain and Obstacles environments, respectively. These results show increased test performance for R-MoMo on a majority of the test robots across all environments. Moreover, R-MoMo generally struggles with fewer robots (e.g. return below 500) than ModuMorph.

a

b

Figure 17: The difference in return between R-MoMo and ModuMorph (top) and the obtained returns (bottom) on each of the 70 unseen test robots in the Flat Terrain environment. Returns are averaged over 10 seeds..

a

b

Figure 18: The difference in return between R-MoMo and ModuMorph (top) and the obtained returns (bottom) on each of the 70 unseen test robots in the Incline environment. Returns are averaged over 10 seeds..

a

b

Figure 19: The difference in return between R-MoMo and ModuMorph (top) and the obtained returns (bottom) on each of the 70 unseen test robots in the Variable Terrain environment. Returns are averaged over 10 seeds..

a

b

Figure 20: The difference in return between R-MoMo and ModuMorph (top) and the obtained returns (bottom) on each of the 70 unseen test robots in the Obstacles environment. Returns are averaged over 10 seeds..

References↩︎

[1]
S. Levine, C. Finn, T. Darrell, and P. Abbeel, “End-to-end training of deep visuomotor policies,” Journal of Machine Learning Research, vol. 17, no. 39, pp. 1–40, 2016.
[2]
D. Kalashnikov et al., “Scalable deep reinforcement learning for vision-based robotic manipulation,” in Conference on robot learning, 2018, pp. 651–673.
[3]
O. M. Andrychowicz et al., “Learning dexterous in-hand manipulation,” The International Journal of Robotics Research, vol. 39, no. 1, pp. 3–20, 2020.
[4]
A. Brohan et al., RT-1: Robotics transformer for real-world control at scale,” arXiv preprint arXiv:2212.06817, 2022.
[5]
B. Zitkovich et al., RT-2: Vision-language-action models transfer web knowledge to robotic control,” in Conference on robot learning, 2023, pp. 2165–2183.
[6]
A. O’Neill et al., “Open X-embodiment: Robotic learning datasets and RT-X models: Open X-embodiment collaboration\(^0\),” in 2024 IEEE international conference on robotics and automation (ICRA), 2024, pp. 6892–6903.
[7]
R. Doshi, H. Walke, O. Mees, S. Dasari, and S. Levine, “Scaling cross-embodied learning: One policy for manipulation, navigation, locomotion and aviation,” arXiv preprint arXiv:2408.11812, 2024.
[8]
Octo Model Team et al., “Octo: An open-source generalist robot policy,” in Proceedings of robotics: Science and systems, 2024.
[9]
A. Gupta, S. Savarese, S. Ganguli, and L. Fei-Fei, “Embodied intelligence via learning and evolution,” Nature Communications, vol. 12, no. 1, p. 5721, 2021.
[10]
N. Vithayathil Varghese and Q. H. Mahmoud, “A survey of multi-task deep reinforcement learning,” Electronics, vol. 9, no. 9, p. 1363, 2020.
[11]
R. Kirk, A. Zhang, E. Grefenstette, and T. Rocktäschel, “A survey of zero-shot generalisation in deep reinforcement learning,” Journal of Artificial Intelligence Research, vol. 76, pp. 201–264, 2023.
[12]
A. Gupta, L. Fan, S. Ganguli, and L. Fei-Fei, “Metamorph: Learning universal controllers with transformers,” International Conference on Learning Representations, 2022.
[13]
Z. Xiong, J. Beck, and S. Whiteson, “Universal morphology control via contextual modulation,” in International conference on machine learning, 2023, pp. 38286–38300.
[14]
A. Hallak, D. D. Castro, and S. Mannor, “Contextual markov decision processes,” arXiv preprint arXiv:1502.02259, 2015.
[15]
E. Todorov, T. Erez, and Y. Tassa, MuJoCo: A physics engine for model-based control,” in 2012 IEEE/RSJ international conference on intelligent robots and systems, 2012, pp. 5026–5033.
[16]
W. Huang, I. Mordatch, and D. Pathak, “One policy to control them all: Shared modular policies for agent-agnostic control,” in International conference on machine learning, 2020, pp. 4455–4464.
[17]
V. Kurin, M. Igl, T. Rocktäschel, W. Böhmer, and S. Whiteson, “My body is a cage: The role of morphology in graph-based incompatible control,” International Conference on Learning Representations, 2021, [Online]. Available: https://arxiv.org/abs/2010.01856.
[18]
D. Ghosh, J. Rahme, A. Kumar, A. Zhang, R. P. Adams, and S. Levine, “Why generalization in RL is difficult: Epistemic POMDPs and implicit partial observability,” Advances in Neural Information Processing Systems, vol. 34, pp. 25502–25515, 2021.
[19]
M. T. Spaan, “Partially observable markov decision processes,” in Reinforcement learning: State-of-the-art, Springer, 2012, pp. 387–414.
[20]
M. J. Hausknecht and P. Stone, “Deep recurrent Q-learning for partially observable MDPs,” in 2015 AAAI fall symposia, 2015, vol. 45, p. 141.
[21]
D. Pathak, C. Lu, T. Darrell, P. Isola, and A. A. Efros, “Learning to control self-assembling morphologies: A study of generalization via modularity,” Advances in Neural Information Processing Systems, vol. 32, 2019.
[22]
F. Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner, and G. Monfardini, “The graph neural network model,” IEEE transactions on neural networks, vol. 20, no. 1, pp. 61–80, 2008.
[23]
T. Wang, R. Liao, J. Ba, and S. Fidler, “Nervenet: Learning structured policy with graph neural networks,” in International conference on learning representations, 2018.
[24]
A. Vaswani et al., “Attention is all you need,” Advances in Neural Information Processing Systems, vol. 30, 2017.
[25]
A. Patel and S. Song, “Get-zero: Graph embodiment transformer for zero-shot embodiment generalization,” arXiv preprint arXiv:2407.15002, 2024.
[26]
C. Sferrazza, D.-M. Huang, F. Liu, J. Lee, and P. Abbeel, “Body transformer: Leveraging robot embodiment for policy learning,” arXiv preprint arXiv:2408.06316, 2024.
[27]
M. Przystupa et al., “Efficient morphology-aware policy transfer to new embodiments,” Reinforcement Learning Journal, 2025.
[28]
Z. Xiong, R. Vuorio, J. Beck, M. Zimmer, K. Shao, and S. Whiteson, “Distilling morphology-conditioned hypernetworks for efficient universal morphology control,” International Conference on Machine Learning, 2024.
[29]
N. Bohlinger et al., “One policy to run them all: An end-to-end learning approach to multi-embodiment locomotion,” arXiv preprint arXiv:2409.06366, 2024.
[30]
B. Li, H. Li, Y. Zhu, and D. Zhao, “Mat: Morphological adaptive transformer for universal morphology policy learning,” IEEE Transactions on Cognitive and Developmental Systems, vol. 16, no. 4, pp. 1611–1621, 2024.
[31]
W. Yu, J. Tan, C. K. Liu, and G. Turk, “Preparing for the unknown: Learning a universal policy with online system identification,” arXiv preprint arXiv:1702.02453, 2017.
[32]
X. B. Peng, M. Andrychowicz, W. Zaremba, and P. Abbeel, “Sim-to-real transfer of robotic control with dynamics randomization,” in 2018 IEEE international conference on robotics and automation (ICRA), 2018, pp. 3803–3810.
[33]
D. Ha, A. Dai, and Q. V. Le, “Hypernetworks,” International Conference on Learning Representations, 2017.
[34]
S. Hochreiter and J. Schmidhuber, “Long short-term memory,” Neural Computation, vol. 9, no. 8, pp. 1735–1780, 1997.
[35]
K. Cho, B. Van Merriënboer, D. Bahdanau, and Y. Bengio, “On the properties of neural machine translation: Encoder–decoder approaches,” in Proceedings of SSST-8, eighth workshop on syntax, semantics and structure in statistical translation, 2014, pp. 103–111.
[36]
Y. Tang and D. Ha, “The sensory neuron as a transformer: Permutation-invariant neural networks for reinforcement learning,” Advances in Neural Information Processing Systems, vol. 34, pp. 22574–22587, 2021.
[37]
J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” arXiv preprint arXiv:1707.06347, 2017.
[38]
S. Kapturowski, G. Ostrovski, J. Quan, R. Munos, and W. Dabney, “Recurrent experience replay in distributed reinforcement learning,” in International conference on learning representations, 2018.

  1. Code will be made public upon publication, but can be shared upon request.↩︎