DMAP: a Distributed Morphological Attention Policy for Learning to Locomote with a Changing Body
Abstract
Biological and artiﬁcial agents need to deal with constant changes in the real world.
We study this problem in four classical continuous control environments, augmented with morphological perturbations.
Learning to locomote when the length and the thickness of different body parts vary is challenging, as the control policy is required to adapt to the morphology to successfully balance and advance the agent.
We show that a control policy based on the proprioceptive state performs poorly with highly variable body conﬁgurations, while an (oracle) agent with access to a learned encoding of the perturbation performs signiﬁcantly better.
We introduce DMAP, a biologically-inspired, attention-based policy network architecture.
DMAP combines independent proprioceptive processing, a distributed policy with individual controllers for each joint, and an attention mechanism, to dynamically gate sensory information from different body parts to different controllers.
Despite not having access to the (hidden) morphology information, DMAP can be trained end-to-end in all the considered environments, overall matching or surpassing the performance of an oracle agent.
Thus DMAP, implementing principles from biological motor control, provides a strong inductive bias for learning challenging sensorimotor tasks.
Overall, our work corroborates the power of these principles in challenging locomotion tasks.
The code is available at the following link: DMAP
Introduction
Animals and humans are highly adaptive to changes in the external environment and in their own body.
For instance, animals can, after an early onset, locomote throughout their development despite dramatic changes in their body size and weight.
Animals can also robustly deal with irregular surfaces, and load changes [1, 2, 3, 4, 5].
Here, we focus on the problem of learning to walk when the body is subject to morphological perturbations, i e. changes in the length and thickness of body parts.
Remarkably, animals can rapidly adapt to these kind of perturbations.
For example, desert ants with elongated or shortened legs can immediately locomote, although misjudging the traveled distance [6].
To challenge artiﬁcial agents on such skills, we adapted four different locomotion environments (Ant, Half Cheetah, Walker,
Hopper) from the OpenAI Gym [7], in which robots have to learn to walk as fast as possible.
In our adaptive setting, the agent’s body parameters are sampled from a morphological perturbation space at the beginning of each episode and we sought to ﬁnd reinforcement learning policies that could deal with these varying bodies.
We hypothesize that principles of the sensorimotor system provide a strong inductive bias to robustly learn such a policy.
Speciﬁcally, we built on three well-known principles:
Independent
-connected networks control a
Independent low-level processing.
Somatosensory inputs are ﬁrst processed locally (e.g., by temporal ﬁltering) before they are integrated across body parts and represented in an topographically ordered fashion [8, 9, 10].
Gated connectivity.
Sensory information is dynamically gated in a task dependent manner [11, 12], and integrated with action and value signals in the brain [13, 10, 14].
Distributed control.
Different body parts are controlled in a distributed fashion [9, 10, 5].
Based on these principles, we propose the Distributed Morphological Attention Policy (DMAP) architecture (Figure 1).
We show that DMAP can robustly learn to locomote with changing bodies without access to the morphological parameters.
To assess the performance of DMAP, we compare it to several alternative agents: Simple, Oracle and rapid motor adaptation (RMA)
Experiments reveal that an agent with access only to the proprioceptive state, consisting of its joint angles and velocities ( Simple ), fails to learn an effective locomotion strategy for a large variety of body conﬁgurations.
In contrast, augmenting the observation with the current body parameters leads to a strong ceiling performance ( Oracle ).
Remarkably, for all different environments, DMAP is as good as
Oracle.
We also compare DMAP to a powerful, recently proposed two-step training technique called rapid motor adaptation ( RMA )
This algorithm trains a temporal convolutional network to infer a context encoding from a short history of proprioceptive state, imitating the Oracle’s morphology encoder.
While RMA can also reach Oracle performance, we show that end-to-end training of such a network does not succeed for more complicated environments (Ant and Walker).
In contrast, DMAP’s policy can be trained end-to-end without access to the morphological parameters and reaches or surpasses Oracle performance for all four environments (Ant, Half Cheetah, Walker, Hopper).
Overall, our work suggests that independent low-level processing, distributed control and dynamic gating mechanisms provide a strong inductive bias for sensorimotor learning.
Finally, we analyze the information gating mechanisms of DMAP.
The strength of gated connections reveals morphologyspeciﬁc patterns, such as bilateral and back-front symmetry.
The low dimensional embedding of the gating mechanism exhibits rotational dynamics, robustly across different morphologies.
Related work
Benchmarks for zero-shot adaptation.
The setup of the experiments described in this paper can be framed as a zero-shot adaptation problem, as the agents are evaluated starting from the ﬁrst moment in which they interact with an unseen context [16].
While classic RL benchmarks, such as Arcade [17] and rllab [18], assess the ability of an agent to maximize the cumulative reward in the training environment, more recent ones introduce perturbations in the dynamics [19, 20], in the observation function [21, 22] or in the reward function [23] at test time.
For a more complete list, we refer to the survey by Kirk et al.
about generalization in Deep Reinforcement Learning [16].
Mann et al.
showed that the performance of RL agents under morphological distributional shifts correlates with the ability of internal models to generalize [24].
We took their design of environments as an inspiration for the design of morphological changes in our work.
Context encoding and adaptation.
Algorithms which enable learning in changing environments are important for the sim-to-real transfer problem, where domain randomization represents a successful strategy to bridge the reality gap [25, 26, 27].
A notable approach to solving the zero-shot transfer problem consists in encoding the context into a hidden state, e.g., by online system identiﬁcation [28], or with a recurrent neural network (RNN) [18, 29].
However, while off-policy continuous control algorithms such as TD3 [30] and SAC [31] are the current state of the art among model-free RL algorithms, both in sample efﬁciency and ﬁnal performance [32], training a recurrent neural network off-policy presents relevant technical challenges [33].
The problem of adaptation to an unseen environmental conditions is central in the Meta-RL literature [34].
Both model-based (ReBAL, GrBal [35], MOLe [36]), and model-free (PEARL [37]) algorithms prove capable to accelerate adaptation to unseen contexts.
However, they require an adaptation phase in the test environment, which can last a few episodes [37] or some hours [35], seemingly suboptimal in our experimental setup (Appendix A.4).
As the focus of our research is on the development of a network architecture to enable robust locomotion in a zero-shot manner, we chose to compare it to RMA, a powerful system identiﬁcation approach by Kumar et al.
They showed that a context encoding can be successfully inferred from a short sequence of past states and actions, and this context encoding can guide a robust policy [15].
Distributed control.
An agent can learn to generalize to multiple body shapes by using modular policies to control each joint independently, e.g., by treating each joint [38] or groups of joints [39] as a node of the graph.
For example, Nervenet [40] achieves state-of-the-art performance on standard
RL benchmarks by leveraging different policies for each joint type.
However, it cannot generalize to different designs in a zero-shot manner (without ﬁne-tuning).
On the other hand, Shared Modular
Policies (SMP) [41] employs the same policy for each joint, enabling an agent to deal with variable observation and action space sizes.
Further solutions to handle bodies with variable connectivity graphs include transformer-based solutions, such as Amorpheus [42] and the recently proposed
AnyMorph [43].
While these works study the problem of controlling bodies with a variable number of segments and connectivity graph, we focus on perturbations which affect morphological parameters, but without changing the body structure.
Hierarchical RL A common approach to solve complex RL tasks is to decompose them in a hierarchy of subproblems [44].
Walking with different body shapes could be interpreted as the union of a long-horizon task (locomotion) and a short-horizon one (body shape identiﬁcation).
Algorithms to tackle multi-horizon control problems by separating high-level behavior from motor primitives include Hierarchical Actor-Critic [45], HeLMS [46] and HiPPO [47].
Adaptation to morphological perturbations requires an agent to rapidly identify which joints are affected and what compensatory torques are required.
We propose a Distributed Morphological
Attention Policy (DMAP) to address this problem (Figure 1).
With DMAP, we seek to facilitate this process by promoting the formation of communication pathways from individual body parts to the speciﬁc joint controllers.
The architecture is designed according to three principles of (biological) sensorimotor control (also mentioned in the introduction):
Independent low-level processing.
Each proprioceptive channel is ﬁrst processed independently, in order to obtain a channel-speciﬁc representation.
This is inspired by the processing in the proprioceptive pathway [9, 10, 11].
Gated connectivity.
An attention mechanism assigns the connection weight between the control policy of a joint and the features of each component of the proprioceptive state.
This is inspired by dynamic gating mechanisms in the brain [11, 12, 13].
Distributed control.
Independent networks control each of the agent’s joint, by outputting one single action scalar (corresponding to a joint’s torque).
This is inspired by the architecture of the spinal cord and the motor homunculus [8, 2, 10].
We implemented those principles in the following way (Figure 1 and for more details Appendix A.1):
3.1 Independent processing and attention-based feature encoder
A feature encoder receives in input a history of T proprioceptive states and actions (( s ( t  T ) , a ( t  T ) ) , ..., ( s ( t  1) , a ( t  1) )) , which we denote as H ( t ) 2 R ( N S + N A ) ⇥ T , N S and N A being the state and action space size, respectively.
Each row h ( t )
i 2 R T of H ( t ) is processed independently by the same Temporal Convolutional Network (TCN).
A learned linear transformation maps each temporally-ﬁltered input channel into the vectors k ( t ) i 2 R N A , representing the attention of each joint towards that input channel i , and v ( t ) i 2 R N V , a value encoding vector of size N V (a hyperparameter of the network).
The vectors k ( t ) i and v ( t ) i from different channels are stacked to form the matrices ˜ K ( t ) 2 R ( N S + N A ) ⇥ N A and V ( t ) 2 R ( N S + N A ) ⇥ N V .
The matrix K ( t ) is obtained by applying a softmax operation to each column of ˜ K ( t ) .
The context encoding vectors for each joint controller are the rows e ( t ) T of the matrix E ( t )
=
K ( t ) T V ( t ) , which are convex combinations of the features extracted from each sensory modality, weighted by their attention score (See diagram in Appendix A.1).
This attention mechanism enables each controller to focus on the relevant morphological information carried by speciﬁc sensory channels, while ignoring the irrelevant ones.
Furthermore, the attention matrix is conditioned to the recent transition history, making it adaptable both across episodes (and thus morphological states) and within a single episode.
3.2 Distributed joint controllers
Differently from standard policy networks, we adopt independent controllers ⌦ i , i 2 1 , .., N A , for each joint, implemented as distinct fully-connected networks.
Each network outputs the element a ( t ) i of the action vector a ( t ) based on the current state s ( t ) , the previous action a ( t  1) and the context encoding e ( t ) i .
Distributed action selection is a natural way to enable the emergence of sensorimotor pathways from the sensory data of different body parts to the controller of a joint.
As one controller needs to predict one single action value, we show in our experiments that smaller fullyconnected networks are sufﬁcient to successfully solve the considered tasks, limiting the complexity overhead stemming from distributed control (See hyperparameters in Appendix A.2).
Qualitatively, independent policy networks can more easily adapt their behavior when a morphological perturbation affects a localized area of the body, leaving the control policy of the other networks unaltered.
Morphological perturbation environments
We consider four locomotion environments of the PyBullet physics simulator [48]: Hopper, Walker,
Ant and Half-Cheetah, which are standard benchmarks for continuous control reinforcement learning algorithms [49, 32].
To test the ability of an agent to learn locomotion with a changing body, we introduce a set of morphological perturbations for each of them.
The perturbations consist in modifying the size and the length of one or more body parts (Appendix A.3).
Such perturbations strongly deteriorate the performance of an agent trained in the unperturbed environment [24].
The parameter  controls the size of the perturbation space C  , limiting the intensity of the perturbation between   and  .
The perturbation is deﬁned as the relative deviation of a morphological parameter from its original value.
For example, a perturbation of 0.5 applied to a leg thickness will make it up to 1.5 times thicker.
This construction results in widely varying body shapes for each of the four environments (Figure 2).
As in the unperturbed environments [32], the training is organized in episodes of at most 1000 time steps.
At the beginning of each episode, bodily parameters are sampled uniformly at random in the space of admissible perturbations (and ﬁxed throughout this episode).
This setup deﬁnes a Contextual
Markov Decision Process [16] M = h S , A , O , T , R ,  , C  i .
M is determined by the state space S , the observation function
O , the action space A , the transition dynamics T , the reward function R , the discount factor  and the context C  .
The observation only includes proprioceptive information, in the form of angular position and velocity of the joints, besides other non-visual variables (orientation of the body, contact forces with the ﬂoor; see Appendix A.3).
The observation space S does not contain explicit information about the current context, as identical observations can come from bodies with different parameters.
On the other hand, the context c  2 C  inﬂuences the transition dynamics p ( s t +1 | s t , a t , c  ) , modifying the effect of an action depending on the body conﬁguration.
This condition of partial observability makes the locomotion problem considerably harder.
However, conﬁrming the intuition that contextual parameters inﬂuence the transition dynamics, we show that an agent can adapt by extracting information from the movement history.
[Table]
Experiments
We trained all the agents with Soft Actor Critic (SAC) [31], using the open source library RLLib [50].
The experiments have been run on a local CPU cluster, for a total of approx.
100 000 cpu-hours.
We used similar hyperparameters to Rafﬁn et al.
[32], which provide state-of-the-art results in the base versions of the four PyBullet environments considered in our work (Appendix A.3).
We ran every experiment for 5 random seeds.
We found that the training of the critic network beneﬁts slightly from augmenting the input observation with the current morphological perturbation.
Therefore, we included it for all policy architectures, except for the Simple agent.
This addition does not pose any limitation to the deployment of the proposed algorithms, as the policy network alone is necessary at that stage.
The experiments involve ﬁve different policy networks (Simple, Oracle, RMA, TCN, DMAP), which differ in architecture, required input data and training procedure (Table 1).
During training, the morphology parameters are randomly sampled from a perturbation space, deﬁned by the parameter  , at the beginning of each episode.
An episode lasts 1000 time steps, unless an early termination condition is met (in case the agent falls).
In this time, the agent has to learn to adapt to its unseen body conﬁguration and run as fast as possible.
We evaluate the performance of a trained agent on its ability to adapt in a zero-shot manner to an unseen body morphology.
To do so, we use a test set of 100 body conﬁgurations for each perturbation intensity, which were not part of the body conﬁgurations used during the training phase.
The agent is tested on a single episode per test morphology, in which it only observes the proprioceptive states (without any information about the body state - except for the Oracle).
During the testing phase, the agent does not observe the rewards it collects.
We apply this testing protocol to the performance evaluation of all the trained policies and all the environments and perturbation levels.
To evaluate the agent’s ability to adapt to an unseen morphology, we consider the reward it accumulates in the complete test episode (without changing the policy; zero-shot setting).
We do so for all the body conﬁgurations in the test set.
Table 2 and the tables in Appendix A.4 list the zero-shot performance of all the trained agents in different test environments.
The rewards in bold are signiﬁcantly higher according to a paired t-test (p-value < 0.001).
We denote the tests as IID if the 100 body conﬁgurations were extracted from the morphological space used during the training (  train   test ), otherwise we denote them as OOD (  train <  test ).
5.1 A morphology encoding boosts learning in strongly perturbed environments
5 random seeds).
While the
Simple performs competitively in low perturbation environments (  = 0 . 1 ), Oracle learns better at higher perturbation levels.
To assess whether an agent can beneﬁt from an encoding of its morphology when learning a locomotion policy, we compare the performance achieved when providing only the proprioceptive state (Simple) and when augmenting it with the perturbation parameters (Oracle).
Speciﬁcally, this
Oracle follows the architecture of Kumar et al.
It extracts a morphology encoding from the raw perturbation vector and concatenates it to the observation vector, thus augmenting the input of the policy network and breaking the partial observability of the environment.
We found that, while in mildly perturbed environments Simple performs competitively with Oracle, learning a control policy solely based on the proprioceptive state becomes difﬁcult with increasing perturbation strength (Figure 3 and Appendix A.4).
In all our experiments with  2 { 0 . 3 , 0 . 5 } , learning a morphology representation improved performance (Table T4).
This demonstrates, as expected, that learning an environment encoding is a viable option to deal with strongly perturbed environments.
5.2 The morphology encoding can be regressed from experience
The morphology parameters c  2 C  inﬂuence the transition dynamics p ( s t +1 | s t , a t , c  ) by changing the physical properties of the body.
Therefore, they are implicitly represented in a sequence of transitions, from which they might be inferred.
To test this possibility, we consider a recently proposed algorithm for adaptive locomotion called rapid motor adaptation (RMA)
It leverages the latent representation of the body generated by a trained Oracle, which has access to the contextual information, to train a TCN network to regress the context encoding from a short history of T proprioceptive states and actions H ( t ) =
(( s ( t  T ) , a ( t  T ) ) , ..., ( s ( t  1) , a ( t
) )) with T = 30 transitions (see Appendix A.2 for hyperparameters).
[Table]
RMA achieves similar average reward to the Oracle (Table 2).
Thus, a sequence of 30 transitions of previous proprioceptive states and actions (approx.
0.5 s in our experiments) is sufﬁcient to extract an embedding of the context that yields high reward.
This demonstrates that RMA’s two-phase adaptation procedure is also viable for morphological perturbations.
5.3 DMAP enables end-to-end training of a morphology encoding network
Our RMA experiments demonstrated that a temporal convolutional network can regress the morphology encoding from a sequence of proprioceptive states and actions representing 0.5 seconds of transition history.
However, it is unclear whether RMA’s TCN encoder requires RMA’s two-step training through imitation to successfully extract a morphology encoding from a sequence of transi- tions.
In fact, if such an encoding helps maximizing the reward, the TCN might successfully learn to extract the same representation by end-to-end reinforcement learning.
Remarkably, end-to-end training of a TCN encoder proves successful for the Hopper and the Half
Cheetah.
However, we encountered strong instability for the Ant and Walker across the random seeds (Figure 4).
Surprisingly, in these two latter environments the addition of the TCN encoder is even harmful for the overall performance, as Simple achieves stronger results in the same environments (Appendix A.4).
Further analysis reveals that for certain random seeds, the training process is disrupted before the agent can learn how to walk (Appendix A.4).
We hypothesize that this inability to learn is due to the difﬁculty of extracting a meaningful encoding to directly optimize the reward without a suitable inductive bias .
For complex morphologies, the
TCN encoder might even act as a confounder for the policy network.
Indeed, replacing the TCN with DMAP leads to better stability during the training (Figure 4) and to a general performance improvement, particularly for the Ant and Walker (Appendix A.4).
Differently from TCN, DMAP performs on par with Oracle in all the considered environments, with the exception of Walker with  = 0 .
Oracle and DMAP also perform similarly when tested OOD, with a slight advantage for
DMAP (Figure 5 and Appendix A.4, A.5).
Finally, we compared the adaptation speed of RMA and
DMAP to unseen morphological perturbations.
We found that the two algorithms perform motor adaptation equally fast (Appendix A.6).
Thus, DMAP learns to robustly locomote in an end-to-end fashion, without access to the morphology parameters.
the larger perturbations falling outside of the training distribution, DMAP slightly outperforms Oracle, which receives the exact perturbation information.
TCN performs worse for more challenging morphologies (Ant, Walker).
5.4 Ablation: The morphology encoding is key in statically unstable environments
To isolate the contribution of the morphology encoding on the performance of DMAP, we performed test-time ablation by evaluating the trained, distributed DMAP policy, while removing the morphology encoding.
We found a performance drop for all the environments and perturbation levels (Table 3).
This performance decay is especially severe in Hopper and Walker (approx.
50% ), where the controller with ablated encoding fails to stabilize the agents, causing the early termination condition to be met sooner and thus yielding lower episode reward.
This suggests that the attention mechanism is fundamental for the dynamic stability of these agents
[Table] [Table]
5.5 How does attention gate sensorimotor information?
We hypothesized that DMAP enables the agent to learn which components of the input are relevant for each joint’s controller.
Next, we verify that such pathways emerge by analyzing the values of the attentional matrix K ( t ) , both in the steady and dynamic condition.
The steady attentional matrix K represents the strength of each input to output (i.e., sensorimotor) connection.
We computed morphology-speciﬁc patterns by averaging elements across time and episodes (Figure 6).
As in the spinal cord, the controllers typically attend to the sensory input of mechanically linked parts (e.g., foot controller to foot velocity, etc.).
Yet, this topography is not universal: for instance in the hopper, all controllers focus on the leg angle and the foot velocity.
Furthermore, the attention patterns seem to reﬂect the bilateral symmetry of the Ant and Walker and a front-back asymmetry due to the directionality of the locomotion task.
However, the attention maps developed by DMAP vary for different seeds, indeed even gait patterns can vary across seeds.
Furthermore, many state and action parameters are correlated during simple locomotion.
This makes interpreting why speciﬁc connections emerge in the attention maps challenging.
Despite this, across seeds, we often observe bilateral and back-front symmetry.
In the future, we want to study this with more challenging locomotion tasks that decorrelate the states (e.g. non-ﬂat terrains).
To analyze the temporal evolution of attention, we projected the attentional matrix K ( t ) onto a 3-dimensional space learned with UMAP [51].
Depending on the phase of the gait, the attentional matrix K ( t ) strengthens or weakens speciﬁc sensorimotor pathways (Figure 7 A).
The graphs reveal a rotational dynamics of the gain modulation typical of biological, legged locomotion [52, 53, 11].
Interestingly, in DMAP the attentional dynamics seem signiﬁcantly less tangled than that of the proprioceptive and action input (Figure 7 B).
We quantiﬁed this visual impression across different environments and found a robust effect (Appendix A.7).
We speculate that this untangling mechanism might contribute to the robustness of the policy to morphological perturbations.
Discussion
We proposed DMAP, a novel architecture that integrates computational principles of the sensorimotor system in reinforcement learning architectures.
We ﬁnd that DMAP enables the end-to-end training of a RL policies in environments with changing body parameters, without having access to the plots.
We refer to
A perturbation information.
Thereby, DMAP matches or even surpasses the performance of an Oracle when tested in a zero-shot manner on unseen body conﬁgurations.
Interestingly unlike DMAP a TCN cannot be trained end-to-end to achieve robust locomotion for more complex morphologies (Section 5.3).
Our work provides an example for biological principles providing excellent inductive biases for learning, consistent with the idea that innate, structured brain connectivity nurtures biological learning [54].
Furthermore the attention mechanism of DMAP solves this challenging task by establishing gait-phase-dependent gain modulation, robust across morphologies.
Notably, synaptic gains in the nervous system also change in a phase-dependent way [52, 53, 11].
Limitations and broader impact.
One characteristic of our current architecture design is that each joint controller observes the full proprioceptive state.
A future direction will be to limit the input to the local sensory data.
Such design, separating low-latency sensorimotor connection (reﬂexes) and slower connections to distal body parts, would be closer to the locomotor system of the animals [2].
Furthermore, in our experiments, we have used an oracle critic to learn an estimate of the cumulative reward and guide the training of the policy.
This gave slightly better results, however, a suitable inductive bias (like DMAP) or possibly inspired by the neural circuits responsible for reward prediction error, could remove this (training) limitation.
Our research aims to facilitate the training of control policies by transferring insights from neuroscience to artiﬁcial intelligence.
Powerful reinforcement learning algorithms might ﬁnd application in robotics, potentially enabling, in the long term, a large-scale deployment of autonomous agents.
Automation comes with great opportunities, and ethical concerns [55].
Acknowledgments and Disclosure of Funding
This research was supported by EPFL and the Swiss Government Excellence Scholarships to AMV.
We thank Ça˘glar Gülçehre, Mackenzie Mathis, Ann Huang, Viva Berlenghi, Lucas Stofﬂ, Bartlomiej Borzyszkowski, Mu Zhou and Haozhe Qi for their feedback on earlier versions of this manuscript.
We also thank Viva Berlenghi for contributing to the single leg perturbation experiments.