Body Transformer: Leveraging Robot Embodiment for Policy 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. Therefore, we propose Body Transformer (BoT), an architecture that leverages 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 throughout the architecture. The resulting architecture outperforms the vanilla transformer, as well as the classical multilayer perceptron, in terms of task completion, scaling properties, and computational efficiency when representing either imitation or reinforcement learning policies. Additional material including the open-source code is available at https://sferrazza.cc/bot_site.

• The paper introduces the BoT architecture that embeds robot morphology using a graph-based masked attention mechanism to enhance policy learning.
• It achieves a near-200% reduction in runtime and outperforms standard transformers and MLPs in both imitation and reinforcement learning tasks.
• The study highlights the practical and theoretical benefits of incorporating physical structure into deep learning for more robust robotics applications.

Body Transformer: Leveraging Robot Embodiment for Policy Learning

The paper "Body Transformer: Leveraging Robot Embodiment for Policy Learning" presents a novel architectural approach aimed at improving policy learning in robotics by utilizing the embodiment structure inherent in physical agents. The authors, Carmelo Sferrazza et al., propose the Body Transformer (BoT) that explicitly accounts for the spatial arrangement of sensors and actuators on the robot, a departure from the traditional vanilla transformers typically used for NLP and computer vision tasks.

Summary of Proposed Architecture

The core innovation of BoT lies in its representation of the robot body as a graph, where nodes represent sensors and actuators, and edges capture physical connections. Crucially, BoT employs masked attention mechanisms within the transformer architecture, allowing each node to attend only to its immediate neighbors. This embedding of the robot's morphology within the learning architecture provides a more relevant inductive bias, which the authors posit improves performance in both imitation learning and reinforcement learning (RL).

Contributions and Numerical Results

The authors enumerate their key contributions as follows:

1. The introduction of the BoT architecture, embedding a novel masking mechanism reflective of the robot's structure.
2. Demonstrated improvements in steady-state performance, generalization, and scaling properties in imitation learning settings.
3. Enhanced RL performance, surpassing both MLP and vanilla transformer baselines.
4. A computational efficiency analysis showing a near-200% reduction in runtime and floating point operations (FLOPs) with the BoT architecture.

These contributions are substantiated with empirical results from a variety of simulated and real-world environments. Notably, BoT outperformed baselines in imitation learning tasks with substantial improvements in return and episode length metrics on both training and unseen validation clips. For instance, the BoT architecture demonstrated a 6-point gain in normalized episode return on the MoCapAct dataset's validation clips over standard transformers.

In reinforcement learning tasks, BoT consistently provided superior performance across several tests, including challenging environments like Humanoid-Board and Humanoid-Hill. The results were evident in metrics such as average episode return during training, where BoT policies learned more effectively and attained higher asymptotic performance compared to MLPs and standard transformers.

Theoretical and Practical Implications

Theoretically, BoT illustrates the significance of incorporating domain-specific structures into general-purpose architectures like transformers. This tailored approach not only enhances learning efficiency and effectiveness but also underlines the potential of structural biases in deep learning.

Practically, BoT's architecture is particularly compelling for real-world robotics applications. By integrating the physical structure directly into the policy learning mechanism, robots can perform more robustly and with greater efficiency. The successful deployment of a BoT-trained policy on a Unitree A1 robot underscores the architecture's feasibility and efficacy outside controlled simulation environments.

Future Directions

While the current study focuses on leveraging spatial information within single time steps, extending BoT to process information across temporal dimensions presents an exciting future direction. Such an extension could further enhance the architecture's applicability to dynamic tasks requiring temporal coherence, potentially making BoT even more robust in real-world scenarios.

Additionally, the authors highlight a notable challenge: the existing deep learning libraries do not fully exploit the computational advantages of the sparse masked attention mechanisms utilized in BoT. Addressing this gap could lead to significant improvements in training efficiency, further broadening BoT's applicability.

Conclusion

The Body Transformer architecture offers a significant advancement in the domain of robot learning by embedding physical structure directly into the learning process. The demonstrated improvements in both imitation and reinforcement learning setups, coupled with practical deployment, make BoT a promising approach for future robotics applications. Further adaptations and optimizations could extend its efficacy, making it a pivotal tool in the evolution of intelligent robotic systems.