Body Transformer: Leveraging Robot Embodiment for Policy Learning
Authors: Carmelo Sferrazza
Dun-Ming Huang
Fangchen Liu
Jongmin Lee
Pieter Abbeel
Venue: CoRL 2024 (2024)
Keywords: Robot Learning
Graph Neural Networks
Imitation Learning
Reinforcement Learning

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.

Body Transformer: Leveraging Robot Embodiment for Policy Learning
Carmelo Sferrazza Dun-Ming Huang
Fangchen Liu UC Berkeley
Jongmin Lee
Pieter Abbeel
Abstract: In recent years, the transformer architecture has become the de facto standard for machine learning algorithms applied natural language process- ing and computer vision. Despite notable evidence of successful deployment of this architecture in the context of robot learning, we claim that vanilla trans- formers do not fully exploit the structure of the robot learning problem. There- fore, we propose Body Transformer (BoT), an architecture that leverages the robot embodiment by providing an inductive bias that guides the learning pro- cess. We represent the robot body graph of sensors and actuators, and rely masked attention pool information throughout the architecture. The re- sulting architecture outperforms the vanilla transformer, well the classical multilayer perceptron, in terms of task completion, scaling properties, and com- putational efficiency when representing either imitation or reinforcement learn- ing policies. Additional material including the apen-source code is available at https: /tss::://fezrzzzaa wtwpraaazaaaccab..o.ae cc/bot_sttteeeeeeeeeeeeeeeeeeees site.
Keywords: Robot Learning, Graph Neural Networks, Imitation Learning, Rein- forcement Learning
Tokenizer
Body Transformer
Detokenizer
Turso
Foond Right His
Hight
Front Right Thigh
Bol
Cigh
Foone Celt
Loott
Roar Left Trigh
Banr at Cvi
actions per nodee
Figure 1: Body Transformer (BoT) is an architecture that considers physical agents graphs of sensors and actuators as nodes, and edges reflecting the structure of the robot body. BoT leverages masked attention simple but flexible mechanism provide body-induced bias the policy. The figure presents the overall schematic of our architecture, exemplified on a Unitree A robot.
1 Introduction
For most of their correcting and stabilizing actions, physical agents exhibit motor responses that are spatially correlated with the location the external stimuli they perceive [1]. This is the case of surfer, where the lower body, i.e. feet and ankles, is mostly responsible for counteracting the imbalance induced by the wave under the board [2]. In fact, humans present feedback loops the
level of the spinal cord's neural circuitry that specifically responsible for the response of single actuators [3].
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

Conrective localized actuation is a main factor for efficient locomotion [4]. This is particularly important for robots too, where however, learning architectures do not typically exploit spatial inter- relations between sensors and actuators. In fact, robot policies have mostly been exploiting the same architectures 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 se- quence dependencies and seamlessly absorb large amount of data. The transformer architecture [5] has been developed for unstructured natural language processing (NIP) tasks, e.g.. language trans- lations, 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 graph with sensors and actuators at its nodes. Then, it applies 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 an imitation learning setting, showing how the inductive bias provided by BoT leads better steady-state performance and generalization, well stronger scaling properties.
We show how BoT improves online reinforcement learning (RL), outperforming MLP and vanilla transformer baselines.
We analyze the computational advantages of BoT, by showing how reformulating the scaled dot product in the computation the attention operation leads to near-200% runtime and floating point operations (FLOPs) reduction.
2
Related Work
Transformers in robotics. Originally developed for NLP applications [5], transformers have been successfully applied across domains, for example in computer vision [6] and audio processing [7]. Several works have shown applications of transformers as means to represent robot policies [8, 9, 10], demonstrating its core advantages in this setting. i.e., variable context length, handling long sequences [11] and multiple modalities [12, 13]. 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 [14] are a class of learning architectures that can pro- cess inputs in the form graph [15]. While early versions of these architectures featured ex- plicit message-passing schemes along the graph [16, 17], more recent architectures mostly feature attention-based approaches. In fact, the vanilla transformer, with its variable context length, inher- ently supports fully connected graphs. However, fftte-of-the-ard 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 [18]. A contemporaneous work, Buterez et al. [19], following similar idea as in the work from Veličković et al. [20], utilizes masked atten- tion layers, where each node only attends to its neighbors, and interleaves such layers with unmasked attention layers. this work, we exploit masked attention in policy learning setting, by addition- ally proposing an architecture that only comprises layers where each can attend to itself and its direct neighbors, resulting in naturally growing context layers, i.e., the outputs of the first layers 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 RI. policies that are effective across different robot morphologies. Earlier works focused on message passing algorithms [21, 22], and were later out- performed by vanilla transformers [23, 24] and transformer-based GNNs that make use of learned encoding and attention biases [25]. All these approaches were only demonstrated in simulated benchmarking scenarios and not applied
real-world robotics setting. Compared to previous work, additionally show that introducing bottlenecks in the attention mask fully exploits the em- bodiment structure and benefits policy learning also for tasks achieved by single agent, leading to better performance and more favorable scaling.
Background
3.1 Attention Mechanisms in Transformers
Transformer, foundational architecture in modern machine learning applications well in work, is powered by the self-attention mechanism [5]. Self-attention weighs the values correspond- ing 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 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:
QKT
AttentionnQ, K, V) = softmax
V.
Vde
where Q. K, V (respectively query, key, and value matrices) learnable linear projections of the sequence elements' embedding vectors, and de is the dimension of the embedding space. As embedding pairs with higher correspondence will have higher score (from the computation of QK"), the corresponding value vector of an embedding's associated key will receive higher weight in the attention mechanism output.
Transformer-based GNNS
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 pessang-passing GNNs 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 [26].
More recently, self-attention was proposed as an alternative to message- passagg-passing [18, 23]. While the standard self-attention mechanism amounts modeling fully connected graph, popular transformer-based GNN, Graphormer [18], injects node-edges information through graph-based po- sitional encodings [27, 28, 18, 25] and by biasing the scaled dot- dot-product, i.c..
(OKT VTL +
+B)
Attention(Q K, V) softmax
(1)
where B is learnable matrix that depends on graph features, e.g., shortest path adjacency matrix, and can effectively encode the graph structure in flexible manner.
3.3 Masked Attention
The attention mechanism can be altered [5] with a binary mask M E (0, 1) (where n is the sequence length), which equivalent to replacing the elements of (1):
O
Mrj =1
B
M 0'
where i and 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.

---

4 Body Transformer
Robot learning policies that employ the vanilla transformer architecture 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, where at each layer in the resulting architecture, a node can only attend information from itself and its direct neighbors. As result, information flows according the graph structure, with the upstream layers reasoning according to local information and the downstream layers pooling more global information from the farther nodes.
Embodiment Graph
M
0
0
1
0
1
0
1
0
0
Adjacency Matrix
0
0
0
1
0
6
20
0
1
O
0
0
0
0
0
0
0
0
5
0
0
0 0
Figure 2: Formulation of Embodiment Mask. The mask M is constructed by adding diagonal of Is to the embodiment graph's 's adjacency matrices. Here, we visualize a simple example of a mask M for an arbitrary agent's 's embodiment where 10.
We present below the various components of the BoT architecture (see also Figure 1): (1) a tokenizer that projects the sensory inputs into the corresponding node embedding, (2) a transformer encoder that processes the input embeddings and generates output features of the same dimension, and (3) detokenizer that decodes the features to actions (or values, for RL critic's training).
Tokenizer. We map the observation vector to 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 [22, 23, 24, 25]. Then, a linear layer projects the local state vector into an embedding vector. Each node's state fed into its node-specific learnable linear projection, resulting in sequence of n embeddings, where n represents the number of nodes (or sequence length). This is in contrast the existing works [23, 24, 25] that 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 [5] 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 = In + A, where In 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 attend itself and its direct neighbors, and introduces considerable sparsity in the problem, which is particularly appealing from computational perspective as highlighted in Section 5.4.
BoT-Mix interleaves layers with masked attention (constructed as in BoT-Hard) with layers with unmasked attention. This similar to the concurrent work in Buterez al. [19], with the distinctions that, D) we find it more effective in our experimental setting to have masked attention layer the first layer, II) mask M is not equivalent to the adjacency matrix, allowing 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 the RL setting, as in the experiments presented in Section 5.2, the detokenizers output values rather than actions, which then averaged across body parts.
Experiments
We assess the performance of BoT across imitation learning and reinforcement learning settings. We keep the same structure as in Figure 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 emhedding vectors its input, (ii) 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 similar number of trainable parameters.
With the following experiments, we aim to answer the following questions:
Does masked attention benefit imitation learning terms of performance and generaliza- tion?
Does BoT exhibit positive scaling trend compared to vanilla transformer architecture?
Is BoT compatible with the RL framework, and what are sensible design choices maxi- mize performance?
Can BoT policies be applied to real-world robotics task?
What the computational advantages of masked attention?
5.1 Imitation Learning Experiments
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 tra- jectories with over 5M transitions, spanning 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 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 the unseen validation clips, demonstrating the generalization capabilities provided by the embodiment-awwan inductive bias. We also report the performance on the training set obtained by 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, strongly outperformed by our architecture. This is particularly remarkable result, as the comparison presents conditions more favorable to the baseline, which features more flexible stochastic policy, optimized in recurrent fashion tailored to the tracking task, and was trained on a larger set of rollouts.
As shown Figure 3b, we also find that BoT-Hard exhibits strong scaling capabilities, as its per- formance 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 overfit to the training data, which is induced by the embodiment bias. Additional comparisons reported in Appendix E, including experiments dexterous manipulation benchmark, Figure 4.
5

---

normalized episode return
normalized episode length
train
val
train
val
ML.P
0.623 0.572 + 0.022 0.568 / 0.534 + 0.025
0.808 /.766 + 0.018 0.777 / 007441 + 0.022
Transformer
0.713 0.664 + 0.024
0.656 0.576 0.038
0.875 0.836 + 0.022
0.834 0.779 0.026
BoT-Hard (ours)
0.751 0.691 = 0.024
0.698 0.650 = 0.035
0.908 / 0.865 = 0.018 0.879 / 0.835 = 0.025
Multi-Clip [29]
/0.654
/0.855
(a) Training and Validation Performance Across Architectures. Statisties of the various architecture- criterion combinations are shown with two values, the leftside being the maximum value recorded during training, and the rightside heing the mean evaluation scores with standard deviation. Results are averaged across five seeds.
Transformer
BoT-Hard (ours)
train
val
Nomaaized Return
0.70
Normalized Return
0.675
0.650
0.68
0.625
0.66
0.600
0.575
0.62
0.550
17.5M
Number of Parameters
Number of Parameters
(b) Training and Validation Performance on MoCapAet Across Number of Trainable Parameters, Each datapoint represents performance averaged across five seeds. We use the 17.5M models for the other imitation learning experiments in this paper.
Figure 3: BoT Performance on Imitation Learning.
Figure 4: Adroit Hand Door, Hammer, and Relocate Tasks (See Results in the Appendix).
5.2
Reinforcement Learning Experiments
We evaluate the RL performance of BoT and baselines using PPO [30] on 4 robotic control tasks in Isaac Gym [31]: Humanoid-Mod, HumanniddBoaart Humanoid-Hill, and A A1-Walk.
All humanoid environments build 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 Appendix A) and include contact forces at all limbs. Humanoid-Me classical running task on flat ground, while in Humanoid- we replaced the flat ground with an irregular hilly terrain. Humanoid Board is newly designed task, where the task for the humanoid keep balancing on board placed on top of cylinder. Finally, adapt the A A1-Walk environment, which is part of the Legged Gym repository [32], where the task is for Unitree Al quadruped robot to follow fixed velocity command.
features the
Figure presents the average episode return of evaluation rollouts during training for MLP, Trans- former, and BoT (Hard and Mix). The solid eurve corresponds to the mean, and the shaded area to the standard error over five 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.

---

MLP
Transformer
BaT-Hard (aurs)
BaT-Mix (ours)
A1-Walk
Humanoid-Mod
Humanoid-Boad --Board
Humanoid-Hii
10000
2000
episade return
6000
8000
4000
20
7000
1000
6000
500
2000
0.0
0.5
1.0
5000
0
6
2
cnvv steps
1e0
cnv
Steps
1e8
env steps
1e8
cnv steps
1e8
(a) Performance various tasks across evaluated architectures.
(b) Snapshots of successful rollouts BoT policies.
Figure 5: Reinforcement Learning Performance on Robotic Control Tasks.
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-Bood) and Humanoid-Hin Hill). Given that the masked attention bottlenecks information propagation from distant body parts, BoT-Hard's strong constraints information communication may hinder effi- cient RL exploration: In Humanoid-Bone 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 intermediate layers to ensure efficient exploration. In contrast, A A1-Walk Humanoid-Mod, the environment's 's state changes more reg- ularly, thus the strong body-induced bias can effectively reduce the search space, enabling faster learning with BoT-Hard.
and Humanoid-Hill, it may be useful for information
5.3 Real World Experiments
The Isaac Gym simulated locomotion environments widely popular for sim-to-real transfer of RL policies without requiring adaptation in the real-world [32]. 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 real-world Unitree A1 robot, adapting the codebases Zhuang et al. [33] and Wu et al. [34]. This showcased in the supplementary video, demonstrating feasibility architecture for real-world deployment. We note that for simplicity we did not make use of teacher- student training or memory mechanisms [35] as common in the locomotion literature, which are known to further improve the transfer by resulting in more natural gaits.
our
5.4 Computational Analysis
Connections between body parts physical agent are often sparse, and are the pre-computable masks M for embodiment graphs in BoT. Masked attention mechanisms 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 [36]), but do not leverage possible sparsity in the masked attention mechanism, ascribed by Buterez et [19] to missing use cases so far. For fair computational comparison, re-implement the scaled dot product in Equation (1) using CPU-based NumPy and evaluate single sample and attention head, being their batched and multi- multi-head versions further parallelizable on GPUs.

---

4,0
Vanilia Attention
Vanilla la Attention
3.5
Custom (ours)
20M
230
beration
2.5
Number Numberro of PLOO
155M
11.5
1.00
0.5
00
0M
48
16
32
48
64
96
128
48
16
32
48
64
128
Number of Nodes
Number Nodes
(a) Runtime
(b) FLOPs
Figure 6: Computational Analysis of the Custom Masked Attention Implementation. Across 10,000 randomly sampled masks, found that custom implementation provides 200% speed- up in runtime at sequence lengths up to 128 nodes and scales better in number of FL.OPs
Our custom implementation comprises two major changes in the computation of the masked atten- tion mechanism from its vanilla implementation in PyTorch. First all, we only perform the dot product computation for elements that will be masked, resulting in fewer matrix multiplication- induced FLOPs from the computation of vdz QK", Secondly, we only use unmasked values to compute the softmax Equation (I), also resulting 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 is and sparsity equal to that used in the MoCapAct experiments (B 0.908)
In Figure 6, we present scaling results of our implementation against vanilla attention across se- quence lengths (number of nodes). We observe potential speed-ups of up 206% for 128 nodes (e.g., in the order of humanoids with dexterous hands [37]). 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 paral- lelization, this approach may significantly reduce the training time of learning algorithms as shown above. Further details and derivations about these experiments can be found in Appendix H.
Conclusion
In this work, presented novel graph-based policy architecture, Body Transformer, which ploits the robot body structure an inductive bias. Our experiments show how masked attention, which is at the core of BoT, benefits both imitation and reinforcement learning algorithms. Addition- ally, the architecture exhibits favorable scaling and computational properties, making it appealing for applications on high-dimensional systems.
Here, used transformers process sequences of distributed sensory information from the same timestep. However, transformers have been shown to excel processing information across time too. We leave the extension of BoT to the temporal dimension 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 devel- opments in this direction. In addition, our architecture requires 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.
A mask with sparsity 3 has In² zero elements. When 3 = 1 "' the mask reduces to the identity In. In practice, the maximum degree of vertex (or node) in robots will be approximately constant, making the computational complexity of masked attention grow linearly with the number of nodes.

---

Acknowledgments
We would like thank Huang Huang and Antonio Loquercio for their help setting up the Al ex- periments. This work supported in part by the SNSF Postdoc Mobility Fellowship 21 211086, ONR MURI N00014-22-1-27755 BAIR Industrial Consortium, NSF A14OPT AI Centre, NSF AC- TION AI Centre Award IIS-2229876, an ONR DURIP grant. We also thank NVIDIA for providing compute resources through the NVIDIA Academic DGX Grant.
References
[1] H. Forssberg. Stumbling corrective reaction: a phase-dependent compensatory reaction during locomotion. Journal of neurophysiology, 42(():933- 1979.
[2] J. L.. Secomb, O. R. Farley, L. Lundgren, T. T. Tran, A. King, S. Nimphius, and J. M. Sheppard. Associations between the performance of scoring manoeuvres and lower-body strength and power in elite surfers. International Journal of Sports Science & Coaching, 10(5):911-918, 2015.
[3] L. Seminara, S. Dosen, Mastrogiovanni, M. Bianchi, S. Watt, P. Beckerle, Nanayakkara, K. Drewing, Moscatelli, R. L. Klatzky, al. hierarchical sensorimotor control framework for human-in-the-loop robotic hands. Science Robotics, 8(78)eadd5544, 2023.
[4] H. Collins and A. D. Kuo. Recycling energy to restore impaired ankle function during human walking. PLoS one, 5(2):e9307, 2010.
[5] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t. Kaiser, and I. Polo- sukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
[6] A. Dosovitskiy, Beyer, A. Kolesnikov, D. Weissenhorn, Zhai, Unterthiner, M. De- hghani, M. Minderer, G. Heigold, S. Gelly, al. An image is worth 16x16 words: Transform- ers for image recognition at scale. arXiv preprint arXiv:2010.11929/ 2020.
[7] L. Dong, S. Xu, and B. Xu. tprechetranssoormer: no-recurrence sequence-to-sequence model for speech recognition. In 2018 IEEE international conference acoustics, speech and signal processing (ICASSP), pages 5884-5883 IEEE, 2018
[8] T. Z. Zhao, V. Kumar, Levine, and C. Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304.
2023.
[9] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 2021.
[10] Q Zheng, A. Zhang, and A. Grover. Online decision transformer. In international conference on machine learning, pages 27042-27059, PMLR, 2022.
[11] I. Radosavovic, Zhang, B. Shi, J. Rajasegaran, S. Kamat, Darrell, K. Sreenath, and J. Ma- lik. Humanoid locomotion next token prediction. arXiv preprint 2xXv:2240.//9669 2024.
[12] H. Qi, Yi, Suresh, M. Lambeta, Ma, R. Calandra, and J. Malik. General in-hand object rotation with vision and touch. In Conference Robot Learning, pages 1549-2564 PMLR, 2023.
[13] C. Sferrazza, Seo, H. Liu, Y. Lee, and P. Abbeel. The power of the senses: Generalizable manipulation from vision and touch through masked multimodal learning, 2023.
[14] Z. Wu, S. Pan, F. Chen, G. Long, C. Zhang, and S. Y. Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 32(1): 4-24, 2020.

---

[15] F. Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner, and G. Monfardini. The graph neural network model. IEEE transactions on neural networks, 20 20(1):61-80, 2008
[16] W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gionzalez, Zambaldi, M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 012666 2018.
[17] J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pages 1263-1272. PMLR, 2017.
[18] C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877-28888,, 2021.
[19] D. Buterez, J. P. Janet, D. Oglic, and P. Lio. Masked attention is all you need for graphs, 2024.
[20] Veličković, G. Cucurull, A. Casanova, A. Romero, P. Lio, and Y. Bengio. Graph attention networks. arXiv preprint arXiv: 10903, 2017.
[21] T. Wang, R. Liao, J. Ba, and S. Fidler. Nervenet: Learning structured policy with graph neural networks. International conference on learning representations, 2018.
[22] W. Huang, I. Mordatch, and Pathak. One policy to control them all: Shared modular policies for agent-agnostic control. In International Conference on Machine Learning, pages 4455-4464. PMLR, 2020.
[23] V. Kurin, M. Igl, T. Rocktäschel, W. Boehmer, and S. Whiteson. My body is a cage: the role of morphology in graph-based incompatible control, 2021.
[24] A. Gupta, L. Fan, S. Ganguli, and L. Fei-Fei. Metamorph: Learning universal controllers with transformers. arXiv preprint 2rXiv:2203.119333 2022.
[25] S. Hong, D. Yoon, and K.-E. Kim. Structure-awar transformer policy for inhomogeneous multi-task reinforcement learning. In International Conference on Learning Representations, 2021.
[26] U. Alon and E. Yahav. On the bottleneck of graph neural networks and its practical implica- tions, 2021.
[27] W. Park, W. Chang, D. Lee, J. Kim, and S. won Hwang. Grpe: Relative positional encoding for graph transformer, 2022.
[28] Yeom, T. Kim, R. Chang, and K. Song. Structural and positional ensembled encoding for graph transformer. Pattern Recognition Letters, 2024.
[29] N. Wagener, A. Kolobov, F. V. Frujeri, R. Loynd, C.-A. Cheng, and M. Hausknecht. Mocapact: multi-task dataset for simulated humanoid control, 2023.
[30] Schulman, Wolski, P. Dhariwal, A. Radford, and O. Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv: M 1707.06347, 2017.
[31] V. Makoviychuk, L. Wawrzyniak, Y. Guo, M. Lu, K. Storey, M. Macklin, D. Hoeller, N. Rudin, A. Allshire, A. Handa, and G. State. Isaac gym: High performance gpu-based physics simula- tion for robot learning, 2021.
[32] N. Rudin, D. Hoeller, P. Reist, and M. Hutter. Learning walk in minutes using massively parallel deep reinforcement learning. In Conference on Robot Learning, pages 91-100. 9LM, 2022.
10

---

[33] Z. Zhuang, Z. Fu, J. Wang, C. Atkeson, S. Schwertfeger, C. Finn, and H. Zhao. Robot parkour learning. arXiv preprint arXiv:2309. 05666, 2023.
[34] Wu, A. Escontrela, D. Hafner, P. Abbeel, and K. Goldberg. Daydreamer: World models for physical robot learning. In Conference Robot Learning, pages 2226-2240. PMLR, 2023.
[35] Miki, J. Lee, Hwangbo, Wellhausen, V. Koltun, and M. Hutter. Learning robust per- ceptive locomotion for quadrupedal robots in the wild. Science Robotics, 7(62))ebb22822, 2022.
[36] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.
[37] C. Sferrazza, D.-M. Huang, X. Lin, Y. Lee, and P. Abbeel. Humanoidbench: Simu- lated humanoid benchmark for whole-body locomotion and manipulation. arXiv Preprint arxiv:2403. 10506, 2024.
[38] A. Rajeswaran, V. Kumar, A. Gupta, G. Vezzani, J. Schulman, E. Todorov, and S. Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstra- tions. arXiv preprint (rxxiv/779.1008/7 2017.
[39] Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.

---

A Details on RL environments
We adapt the IsaacGym humanoid environment for the three humanoid-related tasks, by modifying the observation space include the vertical position of the torso, root coordinates and angular velocity, joint positions and velocities, and per-limb contact forces. We leave the reward for the Humanoid-Mod and Humanoid-HH Hiill unchanged, while we adapt the reward for Humanoid- Hombnnoid-Bob by forcing the forward target velocity to zero, and appropriately adjusting the target and termination heights to take the balancing board into account. For the Al-Walk task, we adapt the codebase in Zhuang et al. [33] and train the policies using proprioception only for the actor, and additional simulation parameters for the critic. We define the task mantain target velocity of 0.5 m/s irregular terrain.
B World Deployment
We deployed the RL policy trained for Al-Walk task to a real-world Unitree A1 Robot. Three main components which standard for sim- to-real locomotion [33] were required for suc- cessful transfer, namely terrain randomiza- tion during training, similar to Zhuang et al. [33], II) higher stiffness in the joint controllers, and IIII a low-pass action filter. The experi- ments were run on flat ground, with offboard computation on CPU. Commands were sent via WiFi or Ethernet connection.
The policy was evaluated through 5 different rollouts, which were considered successful as long as the robot walked for seconds (based on the available experimental space) without falling. All five rollouts were successful.
We attach supplementary video that demon- strates the real-world deployment. A frame overlay representing the robot motion also shown in Figure 7.
Figure 7: Real- World Deployment. Frame over- lay demonstrating the deployment of the BoT walking policy Unitree Al quadruped robot.
C Positional Encodings
For the reinforcement learning experiments presented in Section 5.2, we found that the use of po- sitional encodings improves the performance of BoT architectures. Specifically, we compute po- sitional encodings through an embedding layer that maps indices - up to n to encoding vectors, which are then added the tokenizers' outputs. While this is beneficial for the reinforcement learn- ing setting, we did not report considerable improvement in the imitation learning setting, which present without the use of positional encodings. In fact, these not strictly necessary, as in the BoT architecture tokenizers do not share weights across body parts, and may in principle replace the role of positional encodings.
D Allocating Observations and Actions to Nodes
We follow simple rules allocate observations and actions the different nodes of the graph, each of which represents robot limb and is mapped to row (or column) of the mask M. We distinguish between two types of nodes:
12

---

A root node, to which we assign global observations (e.g., robot position and orientation) and environment observations (e.g., door handle angle, etc.).
Non-root nodes, to which assign local observations (c.g., joint angles) and local actions (e.g., joint commands).
giobal
No actuators
Environment observations
Root Node
All sersors at body part Non Non-Roat Node
Figure Rules for Allocating Quantities Nodes.
As rule of thumb, observations or actions that spanned multiple nodes were assigned to the closest node in the graph, or to either of the nodes involved.
We present example allocation for the A1-Walk task below:
base observations: orientation and angular velocity, actions: none.
front-left-hip observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
front-left-thigh observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
front-left-calf observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
front-right-hij observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
froht-right-thigh observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
front-right-calf observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
rear-left-hip observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
rear-left-thigh observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
rear-left-calf observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
rear-right-hip observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
rear-right-thigh observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
rear-right-calf observations: corresponding joint angle, joint velocity, and previous joint command, actions: corresponding joint command.
13

---

E
Additional Imitation Learning Experiments
normalized episode return
normalized episode length
train
val
train
val
BoT--Hard (ours)
0.75/ 0.69 0.02 0.69 0.66 0.04
0.91 /0.87 0.02 0.88/ 0.84 + 0.03
B-T- (ours)
0.72 / 0.68 + 0.03 0,66/ 0.63 = 0.01
0.89/ /0.86 + 0.03
0,85/082 + 0,01
Boft
0.71 0.68 0.01 0.64 0.58 0.03
0.88 0.87 0.01
0.83 0.03
BoT-Hard/Random
0.71 0.68 + 0.01 0,65 0.60 = 0.03
0.89 / 0,85 + 0.02
0,83/0800 = 0,02
SeeT-Hard-Stochastie (ours)
0.76 0.70 0.02 0.70/0.66 = 0.04
0.91 /0.88 + 0.01 0.88/084 = 0.02
(a) Additional Imitation Learning Ablations on MoCapAct Tracking Task. Statistics of the various architecture-criteni combinations are shown with two values, the leftside heing the maximum value recorded during training, and the rightside being the mean evaluation scores with standard deviation. Results are aver- aged across five seeds.
BoT-Hard (ours)
train
val
0.7
0.7
Normalized Retum
lize Retum
0.6
0.5
0.5
0.4
16
Number of Layers
Number of Layers
(b) BoT-Hard vs Number of Layers. Normalized return as function of the number of masked attention layers. Note that the diameter of the MoCapAct humanoid graph is 14. Results are averaged across seeds.
Transformer
BoT-Hard
Dear
Hammer
Relocate
0.8
Success Bate
20.3
V0.2
0.0
Number of Demorstrations
Number Demonstrations
Number of Demorstrations
(c) Adroit Tasks, We report success rates for varying number of training demonstrations. Statistics are aged across last ten training epochs and seeds, shaded areas indicating standard deviation.
Figure 9: Additional Imitation Learning Experiments.
In this section we provide several ablations on the MoCapAct dataset, in addition to those presented in Section 5.1, as well as results on different tasks in the Adroit Hand benchmark [38, 39].
MoCapAct. We compare (i) BoT-Hard, (ii) BoT-Mix, (iii) BoT-Soft, which similarly [25] leams the matrix B in (1) function of the graph's shortest path matrix, and (iv) BoT- Hard/Random with randomly sampled mask, i.e. having ones on its diagonal and the same sparsity as the mask M used for the correct implementation of BoT-Hard, The table in Figure 9a shows the result of this comparison, with BoT-Hard outperforming all baselines on most of of the metrics. The bottlenecks introduced by the masked attention result in better performance compared both to a mixed approach (BoT-Mix) and an approach that also accounts for structure but does not prevent long-range communication (BoT-Soft). As expected, simply sampling random mask without prop- erly accounting for the embodiment structure deteriorates performance. While this work focused on deterministic policies, we also implemented stochastic version of B-T-Hard, which uses zero-
14

---

mean Gaussian policy with fixed variance of 0.01. This variant shows that explicitly accounting for stochasticity can marginally improve performance further.
In Figure 9b, we show BoT-Hard performance as function of the number of masked attention layers. Despite requiring 14 layers (equal to the embodiment graph diameter) for full information propagation across the graph, the plots show how our architecture may already exhibit good perfor- mance for a smaller number of layers.
Adroit Hand. We compare our BoT-Hard architecture with baselines on three dexterous manipu- lation tasks, which are part of the Adroit Hand benchmark. Specifically, restrict the experiments to low-data regime, where transformers notably struggle and tend to overfit [6]. The results in Figure 9c show how BoT--Hard consistently outperforms the vanilla Transformer, while being com- petitive or outperforming the MLP baseline across all three tasks.
F
Additional Reinforcement Learning Ablations
F.1
Effect of Body-Induced Masking in BoT
BoT-Hard (ours)
BoT-Mix (ours)
BoT-Hard/Random
Bot-Mix/Random
Humanoid-Mod
Humanoid-Board
10000
episode retur
2000
8000
1000
6000
5000 epochs
10000
0
5000
10000
epochs
Figure 10: Additional RL Experimental Results on the Effect of Body-induced Masking.
BoT relies on masked attention with its mask determined by the embodiment structure. We conduct an additional experiment in the RL setting to further demonstrate the effect of the body-induced masking in this setting. We compare with BoT-Hard/Random and BoT-MMxxRandom,, where the attention mask M is given by a randomly sampled symmetric binary matrix with the same degree of sparsity (5 = 0.82 for the IsaacGym humanoid). The results are presented in Figure 10. Overall, BoT with random masking (dotted lines) underperforms BoT with body-induced masking (solid lines) in both simpler task (Humanoid-Me) (oodm) and bard-exploration task Hoamanoid-Boarde which highlights that the use of bood--inducee masking is crucial for the performance of BoT.
F.2 Effect of Per-Limb Tokenizer vs. Shared Tokenizer
Transformer
BoT-Hard (ours)
BoT-Mix (ours)
Transformer/Shared
BoT-HHrd/Shared
BoT-Mix/Sharee
Humanoid-Mod
Humanoid-Board
10000
2000
episode return
9000
1500
8000
1000
7000
6000
500
5000
2500
50D0 epochs
7500
10000
2500
5000
7500
10000
epachs
Figure 11: Additional RL. Experimental Results on the Effect of Per-Node (oe)oxzeszerss
15

---

The existing works using Transformer-bast policies [23, 24, 25] for multi-task RL adopt shared linear projections for tokenizers and detokenizers to deal with the varying numher of limbs, i.e., per-limb observation features are projected into embedding vectors by the single shared tokenizer network, and the per-limb hidden vectors are transformed to per-limb actions via the single shared detokenizer network. In contrast, our BoT is designed for tasks with a single morphology, thus we adopt per-node linear projections for tokenizer and detokenizer. We conduct an additional experi- ment to investigate the effect of this design choice, and the results are demonstrated in Figure 11.
In Figure 11, the solid lines denote the results of using per-node tokenizzen/ddeookeniier, and the dotted lines present the results of using a shared lokenizerrdetookcnizrr (which can understood as representatives of the existing methods [23, 24, 25]). Overall, Transformer/BoT with per-node (de)tokenizers significantly outperform their shared (de)tokenizer counterparts in both simpler task (Humanoid-Mod) and hard-exploration task (Humanoid-Boarde This shows that the of tokenizers shared across different limbs for Transformer-based policies hinders efficient learning.
G Training Details
The training parameters of the experiments detailed in Section 5.1 and Section 5.2 are as summarized in Tables 12a, 12b, and 12c.
H FLOP Derivation for Custom Masked Attention Implementation
Below, we comparatively analyze an asymptotic bound for the amount of floating-poie operations required in one scaled dot product (see Equation (I)) call between the vanilla and the masked ap- proach. From hereon, let denote the sequence length and dy the input and output dimension of our attention mechanism.
QKT Vdd." Considering Q E Raxxx (and similarly for K), the computation of QKT will generally require d& multiplications and dy - additions for all n² pairs. Division by vdg results n² divisions and one constant factor e₁ of FLOOPs for computing vdg. The total amount of flops 2nddj + C1.
Computing
Vda Exploiting sparsity, we ignore all inner product computations for zero entries in M, computing only 3n pairs of multiplications. This results in reduction to 23n dg + 01 FLOOSs
Masked computation of QKT
Computing Softmax(S). A softmax for one vector of dimension n requires n exponentiations, n additions, and n divisions, performed for n rows. Let exponentiations require 02 FL.OPs per element, then total of c2)n² n FLOPs is performed.
Masked computation of Softmax(S). As result of sparsity, there is instead total of Bne exponentiations 3n² divisions, and 3n² n additions to compute, reducing our demand to (2 + c2),n2 - n FL.OPs.
Computing the multiplication ( total of ndl pairs are multiplied, where each pair requires 2n operations complete. The total amount of FL FLOPs is 2n nd,, Following similar reasoning with previous writing, total of 2n?dy nıdy FLOPs are performed.
Assuming that our physical agent provides a graph-induced mask M E {0, 1 of sparsity B E [ 1] (such that there are 3n2 > n nonzero entries), then the amount of FLOPs required by a vanilla masked self-attention implementation is 4n2d, (2+ -nd n + C1 while that of custom masked implementation is (2 2)n (2+++2++))) (2+c2) ndy - n C1. Therefore, the performance gap between the vanilla and masked implementations is determined by the sparsity coefficient B, that is, the number of FLOPs that vanilla approach requires will be c(n) times the number of FLOPs
16

---

Parameter
Values
MLP
Transformers
Batch Size
256
256
# Epochs
100
100
# Encoder Layers
3
16
Embedding Input Size
320
320
Feedforward Size
2500
1024
Attention Heads
N/A
5
Learning Rate
le-4
le-4
Eval Rollouts
100
100
# Parameters
16,696,656
17,544, 120
Training Parameters Used for MoCapAct Experiments.
Parameter
Values
MLP
Transformers
Batch Size
256
256
# Epochs
200
200
# Encoder Layers
3
16
Embedding Input Size
320
320
Feedforward Size
2500
1024
# Attention Heads
N/A
5
Learning Rate
le-4
le-4
Eval Rollouts
20
20
# Parameters
14,305,524
17,131,292
(b) Training Parameters Used for Adroit Hand Experiments.
Parameter
Values
MLP
Transformers
Num Envs
2048
2048
Batch Size
8192
8192
# Encoder Layers
3
10
Attention Heads
N/A
2
Embedding Input Size
N/A
64
Feedforward Size
150
128
# Parameters
699,467
688,762
(c) Training Parameters Used the Humanoid Reinforcement Learning Experiments.
Figure 12: Training Parameters Used for Experiments in Section 5.
custom masked approach requires:
FLOPs in vanilla approach
lim R
FLOPs masked implementation
4n2dg (2 c2)n² ndy 0
lim
000 (2 + 2)n dg + (2 o2) - ndy - n n
4dumodel 2 +
21
(23 + 2)d 2,3 Bo2
Therefore, even though these implementations share the same asymptotic bound Onr)dy), the cus- tom masked implementation's amount of FLOPPs scales hetter than the vanilla implementation.
Note that it is possible further optimize implementation by sparsifying the multiplication Softmax( (SVV; this is left direction of future work, and requires the of sparse array libraries, which was not in the scope of this analysis.
17

---

I Computation Analysis vs Sparsity
Vanilla
Varilla Alterkion
Custom Mosked Attantion (ours)
Custom Masked Attention lours)
2.5
Custom Masked Attendion (ours)
(mma
50.30
Teng Rerdiin
0.00
1.5
0.00
2.00
66220
11.0
0.15
0.0 0.1 0.2 0.3 044 6.5 0.6 0.7 0.8 0.9
0.0 0.1 0.2 0.3 0.4 0.5 0.7 9.8 0.9
0.0 0.1 0.2
0.3 04 0.5 0.6 0.7 088
Spersiyy
Figure 13: Computational Analysis vs Mask Sparsity. Across 50,000 randomly sampled masks for different number of nodes (16, 32, 64, from left to right), runtime decreases with mask sparsity. BoT-HHrd remains more efficient than the vanilla Transformer for mask sparsities lower than s 0.6, which are much lower than common mask sparsities encountered in a robotics setting (including those in this paper).
Figure 13 shows comparison terms of runtime between masked and unmasked attention for varying degrees of sparsity. The plots highlight how BoT-Hard retains its computational advantages across wide range of sparsity degrees, covering the most common scenarios encountered in robotics setting.
18


Reviews:
[{"review_id": "QBZWMoj2d2", "reviewer": "Reviewer_f458", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "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.", "summary_of_recommendation": "** 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.", "technical_quality": 2, "clarity_of_presentation": 3, "potential_impact": 2, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "- 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?\n- Can you provide more details about the real-world deployment, such as success rate, failure scenarios, or any challenges encountered?\n- How does the epoch on the x-label in Fig. 4(a) correspond to environment steps?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x8dOg1bXDr", "reviewer": "Reviewer_W7Gq", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 4, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "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.", "summary_of_recommendation": "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.", "technical_quality": 4, "clarity_of_presentation": 5, "potential_impact": 4, "robotics_focus": 4, "quality_of_limitations": 2, "questions_for_rebuttal": "## Construction of Mask M:\nIn 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.\nAdditionally, 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.\n\n## Assumptions on the Sparsity of Mask M:\nIn 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.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EDLd9toC5L", "reviewer": "Reviewer_tLjG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "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.", "summary_of_recommendation": "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.", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 3, "robotics_focus": 3, "quality_of_limitations": 2, "questions_for_rebuttal": "How does BoT handle the potential computational burden introduced by the additional transformer layers, especially in resource-constrained environments?\n\nCan the authors provide more insights into the practical implications and challenges of deploying BoT in real-world robotic applications?\n\nCould the authors elaborate on the limitations related to the minimum number of transformer layers required to maintain representation power?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept