COMPOSER: Scalable and Robust Modular Policies for Snake Robots

Early works of snake robots model-based control mainly focus on reproducing the innate behavior of real snakes. Various approaches have been explored, including generating torque proportional to the curvature derivative of the body curve [6], applying a serpenoid traveling wave propagated from the head to the tail [7], utilizing gaits for shape based control [8, 9]. However, simplifications made in these model-based approaches constrain the versatility of the modular snake robot, hurting compliance and flexible locomotion ability. Many recent works in snake robots control have exhibited reduced reliance on the model, including co-optimization of morphology and control [16], locomote with neural-network and CPG-based control[17], coach-based reinforcement learning [12], reinforcement learning for snake robot tracking control[18]. However,  [16, 12, 18] use wheels to simplify snake robots and focus on in-plane control, leaving a gap in the exploration of complex tasks, such as manipulation and 3D locomotion in unstructured environments. Many prior studies have not fully leveraged the inherent modularity of snake robots. While Sartoretti et al. first explore snake robot modularity through the decentralized learning of shape-based locomotion parameters, its applicability is limited by reliance on a shape-based controller [19]. In contrast, our approach in this paper directly learns a control policy to generate actuator torque, which enables learning intricate manipulation and locomotion skills.

Decentralized control is common in biological motor control in living creatures. For example, a majority of neurons in the octopus are found in the arms, which can independently control basic motions without input from the brain, allowing fast responses. Using modular policy for decentralized control has demonstrated efficiency, robustness, and generalizability [20, 21, 22, 23]. Schilling et al. illustrate that within a decentralized control paradigm, reinforcement learning-based motor control not only attains high performance but also showcases enhanced robustness and improved generalizability [20]. Pigozzi et al. employ evolutionary computation to optimize a shared neural controller for Voxel-based Soft Robots  [21]. Whitman et al. use modular policy for reconfigurable robots [22]. In  [23], a shared modular policy is trained with messages passing.  [8] shows that fewer individual degrees of freedom coupled means more efficient locomotion on irregular terrains. However, decentralized control has primarily been applied to snake robots within model-based control [6, 8, 19]. As far as we know, this is the first work investigating learning a modular torque control policy to exploit the modularity of snake robots.

On one hand, many works in MARL follow the Centralized Training Decentralized Execution (CTDE) framework, where decentralized policies are trained in a centralized fashion with extra information. MADDPG adopts actor-critic structures and learns a centralized critic [24]. Value-decomposition (VD) methods decompose the joint Q-function as a function of agents’ local Q-functions [25, 26]. MAPPO demonstrates supreme performance in homogeneous multi-agent cooperative settings [27]. HATPRO and HAPPO achieve promising results in heterogeneous settings, without parameter sharing among agents [28]. On the other hand, some works employ centralized communication protocols during execution to share local information and augment the coordination among agents [29, 30, 31]. In this paper, we aim to utilize MARL to develop a modular policy tailored specifically for snake robots and continuum robots.

In this paper, we formulate snake robot control as a fully cooperative MARL problem. Each actuator of the snake robot is considered an individual agent. We consider a multi-agent goal-conditioned partially observable Markov decision process for a snake robot with $n$ actuators, defined by $\langle\mathcal{S},\mathcal{A},O,R,P,\mathcal{G},n,\gamma\rangle$. $\mathcal{S}$ is the state space. $s\in\mathcal{S}$ describes states of all $n$ agents. Given the same configuration of all actuators with the snake robot, each agent shares an identical observation space and action space. Local observation for agent $i$ with $s$ is $o_{i}=O(s;i)$. At each time step, agent $i$ executes action $a_{i}$ sampled from a distribution generated by policy $\pi_{\theta}(.|o_{i},g)$, which is parameterized by $\theta$ and shared by all agents. And all $n$ agents receive shared reward $r_{1},\cdots,r_{n}=R(s,\boldsymbol{a},g)$. $R(s,\boldsymbol{a},g)$ denotes the reward function and $g\in\mathcal{G}$ is a goal shared among all agents. $\boldsymbol{a}=(a_{1},a_{2},\ldots,a_{n})\in\mathcal{A}$ is the joint action of $n$ agents, which produces the next state $s^{\prime}$ with transition probability $P(s^{\prime}|s,\boldsymbol{a})$.

We aim to learn a shared modular torque control policy $\pi_{\theta}$ for the snake robot. Each agent makes decisions individually while all agents execute actions simultaneously. The coupling between agents of the snake robot requires precise coordination among all cooperative agents to successfully complete tasks. The modular torque control policy $\pi_{\theta}$ is trained to jointly maximize the expected return $J(\theta)=\mathbb{E}_{\boldsymbol{a_{t}},s_{t}}\left[\sum_{t}\gamma^{t}R(s_{t},\boldsymbol{a_{t}},g)\right]$, with the discount factor $\gamma$. $\pi_{\theta}$ is optimized with the experiences of all agents simultaneously.

In addition to using a shared torque control policy to output torque commands for each agent, we propose to use a shared modular imagination policy, $\pi^{H}_{\psi}$, parameterized by $\psi$ to predict per-step displacement for each agent. We use this imagination policy as an imaginary planner to decompose the long-horizon task into manageable steps and provide supplementary rewards for the torque control policy.

The imagination policy $\pi_{\theta}(.|o_{i},g)$ is conditioned on the local observation $o_{i}$ and goal, and it generates distribution over anticipated displacement $\Delta p_{i}\in\mathbb{R}^{3}$. The displacement vector $\Delta p_{i}$ points from agent $k$’s current position $p_{i}\in\mathbb{R}^{3}$ to an imagined next-step position $p_{i}^{\prime}\in\mathbb{R}^{3}$, $p_{i}^{\prime}=p_{i}+\Delta p_{i}$. $\boldsymbol{\Delta p}$ denotes the predicted snake whole body shape transition. Agents receive rewards by making good predictions so that the imagined next-step positions $p_{i}^{\prime}$ can be a reasonable subgoal which decomposes the long-horizon task. The reward function for imagination policy is defined to capture the proximity of the imagined next-step positions to the goal completion state, $R^{H}(\boldsymbol{p^{\prime}},g)=d(\boldsymbol{p^{\prime}},g)$, $\boldsymbol{p^{\prime}}=\{p_{i}^{\prime}\}_{i=1}^{n}$ is the joint next-step positions predicted by $n$ agents. For example, in the shape formation task of snake robots, we use the Wasserstein distance [32] between the imagined next-step position $\boldsymbol{p^{\prime}}$ and goal positions $g$ as the inverse of the reward for imagination policy, $\displaystyle W(\boldsymbol{p^{\prime}},g)=\inf_{\pi\in\mathcal{T}(\boldsymbol{p^{\prime}},g)}\int_{\mathbb{R}^{3}\times\mathbb{R}^{3}}|x-y|d\pi(x,y)$. The goal of the imagination policy is to maximize the accumulated reward $J(\psi)=\mathbb{E}_{\boldsymbol{p_{t}^{\prime}},s_{t}}\left[\sum_{t}\gamma^{t}R^{H}(\boldsymbol{p_{t}^{\prime}},g)\right]$.

The modular torque control policy receives rewards by both following the imagined displacement as prescribed by the imagination policy and completing the task. Agent $i$ executes action $a_{i,t}$ induced by the torque control policy $\pi_{\theta}(.|o_{i,t},g)$, and reaches position $p_{i,t+1}\in\mathbb{R}^{3}$. The “follow-the-imagination” reward $r_{f}$ is defined as the sum of errors between actual position $p_{i,t+1}$ agent $i$ reaches and the imagined next-step position $p_{i,t}^{\prime}$.

The reward function for the torque control policy can be described as $R=r_{t}+r_{f}$, where $r_{t}$ is a task-specific reward as shown in Table I in Section IV-A.

The shared imagination policy and the shared torque control policy are jointly trained following the framework of Multi-Agent PPO (MAPPO)[27]. For each policy, we train two separate neural networks, the actor network and the critic network, $\pi_{\theta}(.|o_{i},g)$ and $V_{\phi}(s)$ for the torque control policy, and $\pi^{H}_{\psi}(.|o_{i},g)$ and $V^{H}_{\omega}(s)$ for the imagination policy, respectively. The critic networks, $V_{\phi}$ and $V^{H}_{\omega}(s)$, map the joint state to value: $S\rightarrow\mathbb{R}$.

The actor network for the torque control policy $\pi_{\theta}$ is trained to maximize the objective

where $\displaystyle{r_{\theta,i}^{(k)}=\frac{\pi_{\theta}(a_{i}^{(k)}|o_{i}^{(k)})}{\pi_{\theta_{old}}(a_{i}^{(k)}|o_{i}^{(k)})}}$, $A_{i}^{(k)}$ is the estimated advantage computed using the Generalizable Advantage Estimation (GAE)[33], $H$ is the policy entropy, and $\sigma$ is the entropy coefficient hyperparameter. The critic network $V_{\phi}$ is trained to minimize the loss function

where $\hat{R}$ is the discounted reward-to-go, $B$ refers to the batch size, and $n$ refers to the number of agents. The actor network and critic network for the imagination policy, $\pi^{H}_{\psi}$ and $V^{H}_{\omega}$ are optimized similarly.

The coordinated behavior of a snake robot is heavily dependent on the cooperation of its highly coupled modules. We incorporate the self-attention [34] mechanism into the multi-agent training of the shared control policy, aiming to provide global information for each agent by augmenting inter-module communication.

The self-attention encoder takes joint observation sequence $\boldsymbol{o}=\{o_{i}\}_{i=1}^{n}\in\mathbb{R}^{n\times d}$ as input, where $n$ is the sequence length, and $d$ is the observation dimension. The self-attention encoder maps $\boldsymbol{o}$ to new latent space observation sequence of equal length $n$ as shown in Fig. 1. Attention is calculated as:

Where $Q,K,V$ are query, key, and value, $Q=XW_{q}\in\mathbb{R}^{n\times r},K=XW_{k}\in\mathbb{R}^{n\times r},V=XW_{v}\in\mathbb{R}^{n\times f}$, $W_{q},W_{k}\in\mathbb{R}^{d\times r},W_{v}\in\mathbb{R}^{d\times f}$ are linear transformations. By calculating the attention, the new local observation for agent $i$ incorporates global information in the form of the weighted sum of $n$ transformed local observations of $n$ agents.

Environments. We introduce five snake robot tasks, Goal Reaching, Block Pushing, Shape Formation, Tube Crossing, and Wall Climbing as shown in Fig. 2. Our simulation environment is adapted from SomoGym  [35]. The reward is calculated as the sum of instant rewards across an episode for each environment. Once the robot completes the task, which is determined by the task-specific success criterion as shown in Table I, or the maximum episode length is reached, the episode terminates.

The action space contains normalized torques for each actively controlled degree of freedom (DOF) of each actuator. The action $a_{i}$ is represented by two scalar $[a_{1},a_{2}]$ normalized to $[-1,1]$, corresponding to two independently actuated bending axes in one actuator. In the multi-agent setting, each agent takes action only based on its own observation including local positions, velocities, orientations, and applied torques. For monolithic baseline training, the control policy takes global observations as input and outputs the action including commands for $n$ agents.

Baselines. We compare our method with five baselines to investigate the effectiveness of COMPOSER. We apply the same hyper-parameters of baseline algorithms from their original papers. PPO is a strong on-policy RL method. We use PPO with zero padding as a monolithic trained baseline to learn a centralized policy. Following the setup in  [11], we zero-pad the states and actions to the maximum dimension across all agents with different lengths. MAT is an on-policy few-shot learner on unseen tasks regardless of changes in the number of agents [36]. MAT_dec is a CTDE-variant of MAT [36]. SMP is a modular policy to generate locomotion behaviors for different skeletal structures via message passing [23]. MAPPO is a strong multi-agent reinforcement learning baseline with competitive sample-efficiency [27].

We evaluate our COMPOSER framework in the following three aspects: learning efficiency in terms of task success rate, robustness against hardware failure, and zero-shot generalization to previously unseen robots.

Success Rate.

In Fig. 3, we illustrate the performance of our proposed method and the baselines across five snake robot tasks. The results are averaged over three seeds, and the standard errors are provided. In almost all tasks, COMPOSER outperforms monolithic trained zero-padding PPO and other multi-agent RL baselines including MAPPO, MAT, MAT_dec, and SMP. Monolithic trained PPO shows clear defects indicating the difficulties inherent in exploring a high-dimensional action space for snake robots. This observation validates the effectiveness of modular policy to tackle this high-dimensional challenge with decentralized control. For modular policy baselines, MAPPO is the strongest one while our method still outperforms it by a large margin, in terms of lower success rates and sample efficiency, showing the effectiveness of the self-attention mechanism and imagination policy. MAT features a complex model structure with both an encoder and a decoder, resulting in slow learning growth compared to COMPOSER. MAT_dec, being fully decentralized without the encoder for global observation embedding and the decoder that grants agents access to preceding agents’ actions in MAT, significantly underperformed in comparison to our proposed method. SMP does not manage to learn a successful control policy across all five tasks, which aligns with results reported in  [35], that TD3 has also shown no relevant improvement in nearly all training.

Robustness against Hardware Failures.

In Table II, we demonstrate the performance of different policies trained on a normally functioning snake robot for Goal Reaching when tested in scenarios involving partial hardware failures. This partial hardware failure includes agent action corruption and agent observation corruption. In the case of action corruption, faulty agents are assigned actions with a value of zero. In the case of observation corruption, the observation of faulty agents is assigned a value of zero. At the beginning of each episode, faulty agents are randomly selected based on the specified faulty probability, $p=n_{\rm fault}/n,$, where $n_{\rm fault}$ is the number of corrupted agents, and $n$ is the total agent number. Corrupted agents fail to follow the prescribed policy.

In Table II, we evaluate different policies on four scenarios, including zero corrupted agent, one corrupted action, two corrupted actions, and one corrupted observation. The corresponding model is evaluated on 100 episodes, and results include success rate, final distance to the goal, and episode length. The upper arrow means higher is better, and the down arrow means lower is better. The best performance of each task is marked in bold. We compare COMPOSER with monolithic trained baseline PPO and the strongest modular policy baseline MAPPO. COMPOSER demonstrates greater fault tolerance by having a shorter distance to the goal, fewer episode steps across all scenarios, and a higher success rate across action corruption scenarios, when compared to PPO and MAPPO, as shown in Table II. In the case of one corrupted observation, MAPPO demonstrates a slightly higher success rate. This outcome could be attributed to the fact that the corruption of a single agent’s local observation might lead to confusion in the observations of other agents, given the presence of the attention mechanism. In contrast, MAPPO is less affected by the corruption of local observations due to the lack of inter-agent communication.

Zero-shot Generalizability to Snake Robots with Different Number of Joints. To demonstrate zero-shot generalizability, we train the policy on a snake robot with 8 agents and directly apply the trained policy on previously unseen snake robots with 9 agents, 10 agents, and 11 agents. As shown in Table III, our modular policy demonstrates notable generalizability in comparison to PPO and MAPPO, particularly when applied to a 10-agent snake robot, where PPO and MAPPO almost always fail, but our proposed method has a 33% success rate. This shows the potential of our method to be applied to a scalable modular snake robot.

As shown in Fig. 4, we evaluate the performance of COMPOSER alongside its variants: one devoid of the imagination policy (COMPOSER_w/oI) and another excluding the self-attention mechanism (COMPOSER_w/oA) across five tasks. While COMPOSER_w/oI and COMPOSER_w/oA achieved a success rate similar to COMPOSER in simpler tasks such as Goal Reaching and Block Pushing, they fell notably short in tasks of increased complexity. In more challenging tasks, such as Shape formation which requires desired deformation, and Wall climbing and Shape Formation which necessitate effective interaction with the environment, the ablations demonstrate markedly worse performance. COMPOSER enjoys a higher success rate by leveraging both the attention mechanism for enhanced coordination and the imagination policy for more efficient planning.

In Fig. 5, we visualize the heatmap of the attention matrix $A=\displaystyle{\rm softmax}(\frac{QK^{T}}{\sqrt{d_{k}}}),A\in\mathbb{R}^{8\times 8},$ for an 8-agent snake. $A_{ij}$ is the attention score for agent $i$ attending to agent $j$. The first agent ($i=0$), denotes the snake head, and the last agent ($i=7$) denotes the snake tail. As shown in Fig. 5, our attention mechanism effectively captures different patterns for different tasks, and the attention dynamically moves to different body segments over time. This shows that our attention mechanism has the potential to encompass both task-related details and the underlying dynamics of the snake robot. In Goal Reaching (the first row in Fig. 5), the attention matrix consistently focuses on the snake head, which directly determines the shared reward and the success condition of the task. We also observe that a portion of the attention transitions from the snake head to the snake tail from $t=0$ to $t=90$, akin to the sinusoidal actuation pattern that generates locomotion through anisotropic friction [37]. In Wall Climbing (the second row in Fig. 5), the agents attend to different body segments that are interacting with the wall over time. In Tunnel Crossing (the third row in Fig. 5), the attention is distributed over more segments than the first two tasks at the early stage, because its whole body has to adapt to the tunnel shape. Then, the agents attend more to the head part as the head of the snake gets close to its targeted position. Notably, the colors across different agents are often nearly uniformly distributed within the same column of the heatmap, implying that $n$ agents achieve a shared consensus on global information through self-attention communication.