PyTorch

Today’s dominant deep-learning framework, used to train most embodied AI models.

PyTorch is an open-source deep-learning framework originally developed by Meta’s (then Facebook’s) AI research team, now hosted by the PyTorch Foundation under the Linux Foundation. It’s built around tensor (multi-dimensional array) operations and automatic differentiation, with code that reads close to ordinary Python and is easy to debug, which is why it dominates in academia. Most embodied AI models — Diffusion Policy, ACT, OpenVLA, the PyTorch version of π0, LeRobot, and more — are built on PyTorch, which calls into an NVIDIA GPU through CUDA for acceleration. The core things a newcomer needs to learn are tensor operations, building a network with nn.Module, loading data with DataLoader, writing an optimizer training loop, and saving and loading checkpoints.

ExampleAfter import torch, model.to('cuda') moves the network onto the GPU for training, and torch.save writes out a checkpoint once training finishes.

Also called

torch

Related

CUDA · JAX · LeRobot · Checkpoint · Python and C++ · Hugging Face Transformers

Sources

PyTorch official site
PyTorch Foundation