Body Transformer: Leveraging Robot Embodiment for Policy Learning

Authors: carmelo  sferrazza
dun-ming  huang
fangchen  liu
jongmin  lee
pieter  abbeel
u c berkeley
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 13:58:14

Abstract: In recent years, the transformer architecture has become the de-facto standard for machine learning algorithms applied to natural language processing and computer vision. Despite notable evidence of successful deployment of this architecture in the context of robot learning, we claim that vanilla transformers do not fully exploit the structure of the robot learning problem. We propose Body Transformer (BoT), an architecture that exploits the robot embodiment by providing an inductive bias that guides the learning process. We represent the robot body as a graph of sensors and actuators, and rely on masked attention to pool information through the architecture. The resulting architecture outperforms the vanilla transformer, as well as the classical multilayer perceptron, with respect to task completion, scaling properties, and computational efficiency when representing either imitation or reinforcement learning policies.

Main paper: 
 i. introduction: 
 For most of their correcting and stabilizing actions, physical agents exhibit motor responses that are spatially correlated with the location of the external stimuli they perceive [9] . This is the case of a surfer, where the lower body, i.e. feet and ankles, is mostly responsible for counteracting the imbalance induced by the wave under the board [23] . In fact, humans present feedback loops at the level of the spinal cord's neural circuitry that are specifically responsible for the response of single actuators [24] .Corrective localized actuation is a main factor for efficient locomotion [5] . This is particularly important for robots too, where however, learning architectures do not typically exploit spatial interrelations between sensors and actuators. In fact, robot policies have mostly been exploiting the same archi-tectures developed for natural language or computer vision, without effectively leveraging the structure of the robot body.This work focuses on transformer policies, which show promise to effectively deal with long sequence dependencies and seamlessly absorb large amount of data. The transformer architecture [27] has been developed for unstructured natural language processing (NLP) tasks, e.g., language translations, where the input sequences often map to reshuffled output sequences. In contrast, we propose Body Transformer (BoT), an architecture that augments the attention mechanism of transformers by taking into account the spatial placement of sensors and actuators across the robot body.BoT models the robot body as a graph with sensors and actuators at the graph's nodes. Then, it applies a highly sparse mask at the attention layers, preventing each node from attending beyond its direct neighbors. Concatenating multiple BoT layers with the same structure leads to information being pooled throughout the graph, thus not compromising the representation power of the architecture.Our contributions are listed below:• We propose the BoT architecture, which augments the transformer architecture with a novel masking that leverages the morphology of the robot body. • We incorporate this novel architecture in an imitation learning setting, showing that the inductive bias provided by BoT leads to I) better steady-state performance and generalization and II) stronger scaling properties. • We show how BoT improves existing online reinforcement learning (RL) algorithms, where BoT outperforms MLP and vanilla transformer baselines.• We analyze the computational advantages of BoT, by showing how reformulating the scaled dot product in the computation of the attention operation leads to near-200% runtime and floating point operations (FLOPs) reduction.II. BODY TRANSFORMER Robot learning policies that employ the vanilla transformer architecture as a backbone typically neglect the useful information provided by the embodiment structure. In contrast, here we leverage this structure to provide a stronger inductive bias to transformers, while retaining the representation power of the original architecture. We propose Body Transformer (BoT), which is based on masked attention (see Section B-C), where at each layer in the resulting architecture, a node can only attend to information from itself and its direct neighbors. As a result, information flows according to the graph structure, with the upstream layers reasoning according to local information and the downstream layers pooling more global information from the farther nodes. 0 1 0 0 0 0 0 0 0 0 1 0 1 0 1 0 1 0 1 0 0 1 0 1 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 1 1 0 0 0 0 0 0 0 0 1 1 1 0 1 0 1 0 1 0 0 1 1 1 0 0 0 0 0 0 0 0 1 1 0 0 0 0 0 0 0 1 0 0 1 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 1 1 0 0 0 0 0 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 1 1 0 0 0 0 0 0 0 0 1 1 Fig. 2 : Formulation of Embodiment Mask. The mask M is constructed by adding a diagonal of 1s to the embodiment graph's adjacency matrices. Here, we visualize a simple example of a mask M for an arbitrary agent's embodiment where n = 10.We present below the various components of the BoT architecture (see also Figure 1 ): (i) a tokenizer that projects the sensory inputs into the corresponding node embedding, (ii) a transformer encoder that processes the input embeddings and generates output features of the same dimension, and (iii) a detokenizer that decodes the features to actions (or values, for RL critic's training).Tokenizer. We map the observation vector to a graph of local observations. In practice, we assign global quantities to the root element of the body, and local quantities to the nodes representing the corresponding limbs, similarly to previous GNN approaches [13, 14, 11, 12] . Then, a linear layer projects the local state vector into an embedding vector. Each node's state is fed into its node-specific learnable linear projection, resulting in a sequence of n embeddings, where n represents the number of nodes (or sequence length). This is in contrast to the existing works [14, 11, 12] that use a single shared learnable linear projection to deal with varying number of nodes in multi-task RL.BoT Encoder. We use a standard transformer encoder [27] with several layers as a backbone, and present two variants of our architecture:• BoT-Hard masks every layer with a binary mask M that reflects the structure of the graph. Specifically, we construct the mask as M = I n + A, where I n is the identity matrix of dimension n, and A is the adjacency matrix corresponding to the graph (see Figure 2 for an example).Concretely, this allows each node to attend to itself and its direct neighbors, and introduces considerable sparsity in the problem, which is particularly appealing from a computational perspective as highlighted in Section III-D. • BoT-Mix interleaves layers with masked attention (constructed as in BoT-Hard) with layers with unmasked attention. This is similar to the concurrent work in Buterez et al. [3] , with the distinctions that, I) we find it more effective in our experimental setting to have a masked attention layer as the first layer, II) our mask M is not equivalent to the adjacency matrix, allowing a node to additionally attend to itself at every layer of the architecture. Detokenizer. The output features from the transformer encoder are fed into linear layers that project them to the actions associated with the node's limb, which are assigned based on the proximity of the corresponding actuator with the limb. Once again, these learnable linear projection layers are separate for each of the nodes. When BoT is employed as a critic architecture in reinforcement learning settings, as in the experiments presented in Section III-B, the detokenizer would output values rather than actions, which are then averaged across body parts.We assess the performance of BoT across imitation learning and reinforcement learning settings. We keep the same structure as in Figure 1 and only replace the BoT encoder with the various baseline architectures to single out the effect of the encoder. Particularly, across the various experiments listed in this section, we present the following baselines and variations: (i) an MLP that stacks all embedding vectors as its input, (ii) a vanilla unmasked transformer encoder, (iii) BoT-Hard that only uses masked self-attention layers, (iv) BoT-Mix that alternates between masked and unmasked self-attention layers. All comparisons are made across models with a similar number of trainable parameters.With the following experiments, we aim to answer the following questions:• Does masked attention benefit imitation learning in terms of performance and generalization? • Does BoT exhibit a positive scaling trend compared to a vanilla transformer architecture? • Is BoT compatible with the RL framework, and what are sensible design choices to maximize performance? • Can BoT policies be applied to a real-world robotics task?• What are the computational advantages of masked attention?We evaluate the imitation learning performance of the BoT architecture in a body-tracking task defined through the MoCapAct dataset [29] , which comprises action-labeled humanoid state trajectories with over 5M transitions, spanning a total of 835 tracking clips. For each architecture, we train a deterministic behavioral cloning (BC) policy. We evaluate mean returns normalized by the length of a clip, in addition to the normalized length of an episode, which terminates when the tracking error goes beyond a threshold. We run the evaluations both on the training and the (unseen) validation clips.We report results in the table shown in Figure 3a , where BoT consistently outperforms the MLP and transformer baselines. Remarkably, the gap with these architectures further increases on the unseen validation clips, demonstrating the generalization capabilities provided by the embodiment-aware inductive bias. We also report the performance on the training set obtained by a tailored multi-clip policy presented in Wagener et al. [29] with the MoCapAct dataset. While the multi-clip policy is competitive with the vanilla transformer baseline, it is strongly outperformed by our architecture. This is a particularly remarkable result, as the comparison presents conditions more favorable to the baseline, which features a more flexible stochastic policy, was optimized in a recurrent fashion tailored to the tracking task, and was trained on a larger set of rollouts.As shown in Figure 3b , we also find that BoT-Hard exhibits strong scaling capabilities, as its performance keeps improving with the number of trainable parameters compared to the transformer baseline, both on the training and validation clips. This further indicates a tendency for BoT-Hard to not overfit to the training data, which is induced by the embodiment bias. Additional experiments and comparisons are reported in Section F.We evaluate the RL performance of BoT and baselines using PPO [22] on 4 robotic control tasks in Isaac Gym [15] : Humanoid-Mod, Humanoid-Board, Humanoid-Hill, and A1-Walk.All humanoid environments are built on top of the classical Humanoid environment in Isaac Gym, where we modify the observation space to increase the number of distributed sensory information (see details in Section C) and include contact forces at all limbs. Humanoid-Mod features the classical run-ning task on flat ground, while in Humanoid-Hill we replaced the flat ground with an irregular hilly terrain. Humanoid-Board is a newly designed task, where the task is for the humanoid to keep balancing on a board placed on top of a cylinder. Finally, we adapt the A1-Walk environment, which is part of the legged_gym repository [20] , where the task is for a Unitree A1 quadruped robot to follow a fixed velocity command.Figure 4 presents the average episode return of evaluation rollouts during training for MLP, Transformer, and BoT (Hard and Mix). The solid curve corresponds to the mean, and the shaded area to the standard error over three seeds. The result shows that BoT-Mix consistently outperforms both the MLP and vanilla transformer baselines in terms of sample efficiency and asymptotic performance, highlighting the efficacy of integrating body-induced biases into the policy network architecture. Meanwhile, BoT-Hard performs better than the vanilla transformer on simpler tasks (A1-Walk and Humanoid-Mod), but shows relatively inferior results in hard-exploration tasks (Humanoid-Board and Humanoid-Hill). Given that the masked attention bottlenecks information propagation from distant body parts, BoT-Hard's strong constraints on information communication may hinder efficient RL exploration: In Humanoid-Board and Humanoid-Hill, it may be useful for information about sudden changes in ground conditions to be transmitted from the toes to the fingertips in the upstream layers. For such tasks, BoT-Mix strikes a good balance between funneling information through the embodiment graph and enabling global pooling at intermediate layers to ensure efficient exploration. In contrast, in A1-Walk or Humanoid-Mod, the environment's state changes more regularly, thus the strong body-induced bias can effectively reduce the search space, enabling faster learning with BoT-Hard.The Isaac Gym simulated locomotion environments are widely popular for sim-to-real transfer of RL policies without requiring adaptation in the real-world [20] . To verify that our architecture is suitable for real-world applications, e.g., running in real time, we deploy one of the BoT policies trained above to a real-world Unitree A1 robot, adapting the codebases in Zhuang et al. [37] and Wu et al. [31] . This is showcased in the video at https://bodytransformer.site, demonstrating the feasibility of our architecture for real-world deployment. We note that for simplicity we did not make use of teacherstudent training or memory mechanisms [16] as common in the locomotion literature, which are both known to further improve the transfer by resulting in more natural gaits. Connections between body parts of a physical agent are often sparse, and so are the pre-computable masks M for embodiment graphs in BoT. Masked attention mechanisms can benefit from this sparsity, as their computational cost can be reduced by ignoring unnecessary computations of those matrix elements that will eventually be masked out. Large-purpose deep learning libraries such as PyTorch feature largely optimized matrix multiplication and attention routines (e.g., FlashAttention [6] ), but do not leverage possible sparsity in the masked attention mechanism, ascribed by Buterez et al. [3] to missing use cases so far. For a fair computational comparison, we re-implement the scaled dot product in Equation (1) using CPU-based NumPy and evaluate on a single sample and attention head, being their batched and multi-head versions further parallelizable on GPUs.Our custom implementation comprises two major changes in the computation of the masked attention mechanism from its vanilla implementation in PyTorch. First of all, we only perform the dot product computation for elements that will not be masked, resulting in fewer matrix multiplication-induced FLOPs from the computation of1 √ d k QK T .Secondly, we only use unmasked values to compute the softmax in Equation ( 1), also resulting in reduced FLOPs.We measure the average runtime of each implementation of the attention mechanism across 10,000 set of randomly generated Q, K, V , and M . For each randomization, the generated masks M have a diagonal of 1s and sparsity equal to that used in the MoCapAct experiments (β = 0.908) 1 .In Figure 5a , we present scaling results of our implementation against vanilla attention across sequence lengths (number of nodes) {4, 8, 16, 32, 48, 64, 96, 128}. We observe potential speed-ups of up to 206% for 128 nodes (e.g., in the order of humanoids with dexterous hands [26] ). Overall, this shows that the body-induced bias in our BoT architecture not only enhances the overall performance of physical agents but also benefits from the naturally sparse masks that the architecture imposes. With adequate parallelization, this approach may significantly reduce the training time of learning algorithms as shown above.A depiction of the amount of FLOPs required across different sequence lengths is also presented in Figure 5b . Further details and derivation about these experiments can be found in Section I.In this work, we presented a novel graph-based policy architecture, Body Transformer, which exploits the robot body structure as an inductive bias. Our experiments show how masked attention, which is at the core of BoT, benefits both imitation and reinforcement learning algorithms. Additionally, the architecture exhibits favorable scaling and computational properties, making it appealing for applications on highdimensional systems.Here, we used transformers to process sequences of distributed sensory information from the same timestep. However, transformers have been shown to excel at processing information across time too. We leave the extension of BoT to the temporal dimension as future work, as it promises to further improve real world deployment of robot policies, such as the one demonstrated on the Unitree A1 robot.A limitation of our approach is the fact that its computational advantages are currently not fully exploited by modern deep learning libraries, and we hope that this work may stimulate future developments in this direction. In addition, our architecture requires a minimum amount of transformer layers to make sure that the architecture does not lose representation power in modeling long-range relations, which generally increases the amount of required trainable parameters.Originally developed for NLP applications [27] , transformers have been successfully applied across domains, for example in computer vision [8] and audio processing [7] . Several works have shown applications of transformers as a means to represent robot policies [35, 4, 36] , demonstrating its core advantages in this setting, i.e., variable context length, handling long sequences [19] and multiple modalities [18, 25] . However, these approaches use transformers as originally developed for unstructured or grid-like inputs, such as language or images, respectively. In this work, we leverage the robot embodiment by appropriately adapting the transformer attention mechanism.Graph Neural Networks (GNNs). GNNs [32] are a class of learning architectures that can process inputs in the form of a graph [21] . While early versions of these architectures featured explicit message-passing schemes along the graph [2, 10] , more recent architectures mostly feature attention-based approaches. In fact, the vanilla transformer, with its variable context length, inherently supports fully connected graphs. However, state-of-the-art performance on graph interpretation benchmarks is only achieved via modifications of the original transformer architecture, for example by means of learned graph encodings and attention biases [34] . A contemporaneous work, Buterez et al. [3] , following a similar idea as in the work from Veličković et al. [28] , utilizes masked attention layers, where each node only attends to its neighbors, and interleaves such layers with unmasked attention layers. In this work, we exploit masked attention in a policy learning setting, by additionally proposing an architecture that only comprises layers where each can attend to itself and its direct neighbors, resulting in naturally growing context over layers, i.e., the outputs of the first layers are computed using more local information compared to those of the last layers.Exploiting body structure for policy learning. Graph neural networks have been explored by several works as a way to obtain multi-task RL policies that are effective across different robot morphologies. Earlier works focused on message passing algorithms [30, 13] , and were later outperformed by vanilla transformers [14, 11] and transformer-based GNNs that make use of learned encoding and attention biases [12] . All these approaches were only demonstrated in simulated benchmarking scenarios and not applied to a real-world robotics setting. Compared to previous work, we additionally show that introducing bottlenecks in the attention mask fully exploits the embodiment structure and benefits policy learning also for tasks achieved by a single agent, leading to better performance and more favorable scaling.Transformer, a foundational architecture in modern machine learning applications as well as in our work, is powered by the self-attention mechanism [27] . Self-attention weighs the values corresponding to each element of the sequence with a score that is computed from pairs of keys and queries extracted from the same sequence. Thus, it is able to identify relevant pairs of sequence elements in the model output.Concretely, the self-attention output vector is computed through the following matrix operation:Attention(Q, K, V ) = softmax QK T √ d k V,where Q, K, V (respectively query, key, and value matrices) are learnable linear projections of the sequence elements' embedding vectors, and d k is the dimension of the embedding space. As embedding pairs with higher correspondence will have a higher score (from the computation of QK T ), the corresponding value vector of an embedding's associated key will receive a higher weight in the attention mechanism output.In this work, we model the agent embodiment as a graph whose nodes are sensors and actuators, and their connecting edges reflect the body morphology. While message-passing GNNs are suitable inductive biases for this formulation, they tend to suffer from oversmoothing and oversquashing of representations, preventing effective long-range interactions and discouraging network depth [1] .More recently, self-attention was proposed as an alternative to message-passing [34, 14] . While the standard self-attention mechanism amounts to modeling a fully connected graph, a popular transformer-based GNN, Graphormer [34], injects node-edges information through graph-based positional encodings [17, 33, 34, 12] and by biasing the scaled dot-product, i.e.,EQUATION)where B is a learnable matrix that depends on graph features, e.g., shortest path or adjacency matrix, and can effectively encode the graph structure in a flexible manner.The attention mechanism can be altered [27] with a binary mask M ∈ {0, n×n (where n is the sequence length), which is equivalent to replacing the elements of B in (1):B i,j = 0 M i,j = 1 -∞ M i,j = 0 ,where i and j denote row and column indices. This operation effectively results in zeroing out the contribution of the pairs indicated with zeros in the mask M to the computation of the attention. coefficient β, that is, the number of FLOPs that a vanilla approach requires will be c(n) times the number of FLOPs a custom masked approach requires: Therefore, even though these implementations share the same asymptotic bound O(n 2 d k ), the custom masked implementation's amount of FLOPs scales better than the vanilla implementation.Return Normed Length NormedNote that it is possible to further optimize our implementation by sparsifying the multiplication Softmax(S)V ; this is left as a direction of future work, and requires the use of sparse array libraries, which was not in the scope of this analysis. A mask with sparsity β has βnzero elements. When β = 1 -1 n , the mask reduces to the identity In. In practice, the maximum degree of a vertex (i.e. node) in robots will be approximately constant, making the computational complexity of masked attention grow linearly with the number of nodes. 
 
 Appendix:  


Decision: Strengths: 
+ A novel architecture that incorporates robot morphologies for robot learning.
+ Experiments include both simulated and real-robot results.
+ Body Transformer achieves significant gains in performance and efficiency.


Weaknesses:
- The RL results are averaged across a limited number of seeds.
- Body Transformer needs to be tailored to each embodiment, which is non-trivial.
- Some details regarding the mask and real-robot deployment are missing.

Post rebuttal:  
The additional experiments and author responses have addressed the reviewers' concerns.

Reviews: review: title: title: novel architecture, insufficient evaluation
paper_summary: paper_summary: The paper proposes a transformer-based architecture, body transformer, which encodes the robot morphology as a graph of sensors and actuators. They demonstrate that body transformer based policy network achieves high performance in both imitation learning and RL setting. Furthermore, its performance scales well with the number of parameters in the imitation learning setting. The paper also highlights that the proposed method reduce the computation by 50% compared to a vanilla transformer, thanks to the efficient implementation of the masked attention layers. Additionally, they achieve zero-shot sim-to-real transfer in quadruped robot locomotion.
main_review: main_review: Strengths:

- The proposed approach is straightforward and easy to follow.
- The approach is validated in both simulation and real-world robots.
- It effectively leverages the sparsity of the graph structure to reduce computation.

Weaknesses:

The main area that needs improvement in my opinion is the evaluation, which is currently insufficient to fully assess the efficacy of the proposed method:

- In Fig.3(a), the “maximum value achieved throughout three trajectories” is reported. Why over only 3 trajectories? Compared to Fig. 7, it seems to be a typo and actually meant “seeds”.  Reporting the maximum value alone doesn’t provide a complete picture of performance. I think that including mean and variance, and also increasing the number of seeds, would strengthen the results.
- The RL results In Fig. 4(a) are based on 3 seeds, which is insufficient. Increasing the number of seeds would enhance the reliability of the result.
- Fig. 3(b) seems to be based on a single seed.

The real-world deployment demonstrates feasibility but lacks details on the success rate or failure scenarios. Providing this information would give a more comprehensive understanding of the practical application of the proposed approach.
questions: questions: - The metric on the left side in Fig. 7 is ambiguous, “maximum mean value recorded throughout all evaluation epochs across three seeds”. Can you elaborate?
- Can you provide more details about the real-world deployment, such as success rate, failure scenarios, or any challenges encountered?
- How does the epoch on the x-label in Fig. 4(a) correspond to environment steps?
limitations: limitations: 3
review_summary: review_summary: ** Updated post-rebuttal **: The proposed architecture is straightforward and relatively efficient to compute. The clarity of the paper and the quality of the experiments have improved following the rebuttal.
title: title: This is an interesting paper that proposes a very simple yet effective idea.
paper_summary: paper_summary: This paper proposes a new architecture called "Body Transformer (BoT)." BoT represents the robot's body structure as a graph, treating sensors and actuators as nodes. This architecture utilizes a masked attention mechanism, where each node only references information from itself and its adjacent nodes, making robot learning more efficient than traditional transformers. Key contributions include demonstrating that BoT outperforms conventional methods in imitation and reinforcement learning with improved computational efficiency. BoT is shown to be particularly effective for complex tasks and high-dimensional systems, and it has been validated for real-world robot control applications.
main_review: main_review: ## Quality:
### Strengths:
Methodology: The paper employs a comprehensive experimental setup to validate the proposed Body Transformer (BoT) architecture across various tasks, including imitation learning and reinforcement learning.
Detailed analysis: It provides in-depth analysis and comparisons with existing models, demonstrating significant improvements in performance and computational efficiency.
Real-world validation: The real-world experiments on a Unitree A1 robot enhance the practical relevance and applicability of the proposed method.

### Weaknesses:
Implementation details: Some implementation specifics, such as hyperparameters and training details, might be lacking, which can hinder reproducibility.

## Clarity:
### Strengths:
Clear structure: The paper is well-structured, with distinct sections covering background, methodology, experiments, and conclusions.
Figures and Tables: Effective use of figures and tables to illustrate the architecture, experimental setups, and results, aiding in comprehension.

### Weaknesses:
Fewer explanations about mask M: Some sections, especially those explaining the novel masking mechanism, could benefit from more detailed, step-by-step explanations.

## Originality:
### Strengths:
The concept of leveraging the robot's physical structure through a graph representation and integrating it with transformer architecture is innovative.


## Significance:
### Strengths:
Impact on robotics: The proposed architecture has the potential to significantly advance the field of robot learning by improving the efficiency and performance of learning policies.
Broad applicability: The approach is versatile and can be applied to various robotic tasks and configurations, enhancing its significance.

### Weaknesses:
Generalizability: When applying the proposed method to robots with varying physical characteristics (such as humanoid robots with different numbers of joints), it appears essential to tailor the learning process to accommodate each specific robot.
questions: questions: ## Construction of Mask M:
In this study, the mask M utilized by the proposed model is determined based on the robot's morphology, specifically the arrangement of sensors and actuators. While the experiments conducted in this research employed Unitree A1 and Humanoid in Issac Gym, a detailed explanation of how each sensor and actuator was mapped to the mask M is necessary.
Additionally, it should be noted that sensors on a robot do not only gather information specific to their physical location but may also collect data from distant spatial areas (e.g., cameras or IMUs). Similarly, actuators might operate joints located in different spatial positions. It would be beneficial to discuss how mask M should be constructed in such cases.

## Assumptions on the Sparsity of Mask M:
In the computational analysis, the sparsity of mask M is assumed to be β=0.908, based on the MoCapAct example. However, when the number of sensor and actuator nodes on a robot increases, can this assumption truly hold? When considering a single robot in a detailed graph representation, the sparsity might be reduced. I would appreciate a discussion on this point.
limitations: limitations: 2
review_summary: review_summary: This paper proposes a new architecture called "Body Transformer (BoT)." BoT represents the robot's body structure as a graph, where each node only references information from itself and its adjacent nodes, making robot learning more efficient than traditional transformers. The proposed method is highly intriguing. I am confident that addressing a few specific points of improvement will significantly enhance the quality of this paper.
title: title: Body Transformer: Leveraging Robot Embodiment for Policy Learning
paper_summary: paper_summary: The paper introduces the Body Transformer (BoT), a novel architecture that leverages the physical structure of robots to improve policy learning. By representing the robot's body as a graph of sensors and actuators and using masked attention, BoT integrates the spatial relationships within the robot's body into the learning process. This approach outperforms traditional transformer architectures and multilayer perceptrons in terms of task completion, scalability, and computational efficiency in both imitation and reinforcement learning contexts. The paper includes comprehensive experiments demonstrating BoT's effectiveness and computational advantages.
main_review: main_review: The paper provides a thorough theoretical foundation for the proposed approach and supports its claims with extensive experimental results. The experiments convincingly demonstrate that BoT outperforms both MLP and vanilla transformer baselines in imitation learning tasks, showing superior performance and generalization capabilities. Additionally, BoT's ability to scale effectively with an increasing number of parameters highlights its robustness and efficiency. In the context of reinforcement learning, BoT exhibits improved sample efficiency and asymptotic performance, particularly with the BoT-Mix variant, which balances local and global information pooling.

However, the reliance on transformer layers may introduce a significant number of trainable parameters, potentially increasing the computational burden. The paper could also benefit from a more detailed discussion on the practical implications of deploying BoT in real-world scenarios, including the challenges of sim-to-real transfer and the potential need for adaptation mechanisms.
questions: questions: How does BoT handle the potential computational burden introduced by the additional transformer layers, especially in resource-constrained environments?

Can the authors provide more insights into the practical implications and challenges of deploying BoT in real-world robotic applications?

Could the authors elaborate on the limitations related to the minimum number of transformer layers required to maintain representation power?
limitations: limitations: 2
review_summary: review_summary: The BoT presents a novel and effective approach to leveraging robot embodiment in policy learning, with significant contributions in both theoretical development and experimental validation.