Appendix D: PyTorch and JAX for Embodied AI

Why This Appendix Exists

PyTorch and JAX for Embodied AI supports the self-contained promise of the book. It gives readers enough reference material to continue without opening a second textbook.

Reference Notes

This appendix should be used on demand. The main chapters link here when a mathematical tool, software package, benchmark, compute estimate, or citation convention needs a compact refresher.

Practical Checklist

  1. Identify the concept or tool needed by the chapter.
  2. Review the minimal definition and notation.
  3. Run the smallest example.
  4. Return to the chapter with the missing prerequisite restored.
Operational Shortcut

For tools and libraries, prefer official documentation and pinned environment files. For math and notation, prefer the definitions used in this book so symbols stay consistent.

Bibliography & Further Reading

Foundational Papers, Tools, and References

Sutton, R. S., and Barto, A. G.. "Reinforcement Learning: An Introduction." (2018). http://incompleteideas.net/book/the-book-2nd.html

A foundation for value functions, policy gradients, exploration, and the RL framing used throughout the book.

Todorov, E., Erez, T., and Tassa, Y.. "MuJoCo: A physics engine for model-based control." (2012). https://mujoco.org/

The simulator lineage behind much modern robot learning, now extended through MJX and Warp workflows.

Brohan, A. et al.. "RT-1: Robotics Transformer for real-world control at scale." (2022). https://arxiv.org/abs/2212.06817

A landmark in large-scale robot policy learning with transformer policies.

Brohan, A. et al.. "RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control." (2023). https://arxiv.org/abs/2307.15818

A central reference for connecting web-scale VLM knowledge to robot actions.

Open X-Embodiment Collaboration. "Open X-Embodiment: Robotic Learning Datasets and RT-X Models." (2023). https://arxiv.org/abs/2310.08864

The cross-embodiment data and transfer reference used by the data chapters.

Chi, C. et al.. "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion." (2023). https://arxiv.org/abs/2303.04137

The practical diffusion policy reference for imitation learning and continuous action generation.

Hafner, D. et al.. "Mastering Diverse Domains through World Models." (2023). https://arxiv.org/abs/2301.04104

DreamerV3, a modern reference for latent world models and imagination-based control.

Hugging Face. "LeRobot." (2024). https://github.com/huggingface/lerobot

The open robot-learning stack used for datasets, policies, demos, and low-cost embodied AI workflows.

Official documentation and source repositories for PyTorch and JAX for Embodied AI.

Use official docs for install commands, current APIs, and version caveats before running the chapter lab.