EmbodiedData

BEGINNER PROJECT · PUSHT

Train your first robot-learning model

Learn the complete workflow with a small visual pushing dataset—without needing a physical robot.

Beginner31.6 MB datasetSimulationDiffusion Policy

STEP 1

What are we teaching the model?

A circular agent must push a T-shaped block into a matching target. The model observes an image and the agent position, then predicts the next 2D movement. This is a compact way to learn the same observation-to-action idea used by robot policies.

Camera imageWhat the scene looks likeAgent positionWhere the controller is2D actionWhere it should move next

STEP 2

Create a clean environment

Use a virtual environment so the project does not interfere with your other Python packages. The installation extra follows the current official LeRobot packaging.

python -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install "lerobot[diffusion,pusht]"
CPU or GPU?

You can inspect the dataset on CPU. Training a useful Diffusion Policy is much faster on an NVIDIA GPU. Do not start with the full training run if installation or inspection fails.

STEP 3

Load one sample before training

This checks that LeRobot can find the dataset, download its metadata and expose the fields needed by the policy.

from lerobot.datasets.lerobot_dataset import LeRobotDataset

dataset = LeRobotDataset("lerobot/pusht_image")
print("Frames:", len(dataset))
print("Available fields:", dataset[0].keys())
print("First action:", dataset[0]["action"])

This inspection snippet follows the documented LeRobotDataset interface. Dataset revisions can change, so use the linked source card if a field name differs.

Open the PushT source card ↗

STEP 4

Understand what one training example means

Observation

The image and current agent position available to the policy.

Action

The demonstrated 2D movement the policy should learn to predict.

Episode

One full attempt to move the T block into the target.

Policy

The learned function that converts observations into future actions.

STEP 5

Launch a Diffusion Policy training run

Run this only after the inspection step succeeds. The command uses the current LeRobot training CLI structure.

lerobot-train \
  --dataset.repo_id=lerobot/pusht_image \
  --policy.type=diffusion \
  --output_dir=outputs/train/diffusion_pusht \
  --job_name=diffusion_pusht \
  --policy.device=cuda \
  --wandb.enable=false
Important

This starts training; it does not guarantee a successful policy. Results depend on LeRobot version, hardware, training steps and hyperparameters. Check the generated logs and upstream documentation before increasing compute.

STEP 6

Common beginner problems

The dataset download fails

Check your connection and available disk space, then open the Hugging Face source card to confirm the repository is accessible.

CUDA is not available

Inspect the dataset on CPU first. For training, install the PyTorch build matching your NVIDIA driver or use a hosted GPU environment.

A field name is different

The dataset or LeRobot format may have changed. Print dataset[0].keys() and compare it with the current source card and LeRobot documentation.

Training loss falls but the policy performs badly

A low training loss is not the same as successful control. Evaluate rollouts and consult the current upstream PushT guidance before changing hyperparameters.

WHAT NEXT?

Move from PushT to real robot data

After you can inspect and train on PushT, try ALOHA simulation for bimanual actions or SO-100 pick-and-place for a compact real-robot VLA dataset.