Skip to main content
Version: 2.0

Training

Dataset format conversion to LeRobot dataset format​

After collecting the data, we need to convert the data for training the policy.

Here, we will show how to convert the collected data to the LeRobot dataset format.

git clone https://github.com/enactic/openarm_dataset.git
cd openarm_dataset
uv sync
uv run openarm-dataset-convert path/to/collected_dataset_path path/to/output_path --format lerobot_v2.1

LeRobot Dataset v2.1 format file structure:

output_path/
├── data/
│   ├── chunk-000/
│   │   ├── episode_000000.parquet
│   │   ├── episode_000001.parquet
│   │   └── ...
├── meta/
│   ├── info.json
│   ├── episodes.jsonl
│   ├── episodes_stats.jsonl
│   ├── tasks.jsonl
│   ├── stats.jsonl
├── videos/
│   ├── chunk-000/
│   │   ├── observation.images.wrist_right/
│   │   ├── observation.images.wrist_left/
│   │   ├── observation.images.ceiling/
│   │   ├── observation.images.head_left/
│   │   ├── observation.images.head_right/
│   │   |  ├── episode_000000.mp4
│   │   |  ├── episode_000001.mp4
│   │   |  └── ...

You can also convert to the LeRobot Dataset v3 by '--format lerobot_v3.0'. Please see the README.md in the repository for further details.

Training the policy​

After you collected the data and convert it, you can use it to train a policy.

You can train any model you want, but in this tutorial, we will show how to train the LeRobot ACT policy.

Dataset​

We assume you already have the dataset in the LeRobot dataset format.

However, in order to make it easier for you to follow the inference tutorial later, we will use the dataset that we collected and converted in the data collection tutorial.

You can use this dataset to train the policy. In order to make training faster, we filtered out ceiling and head_right.

Training setup​

Pre-requisites​

  • ffmpeg
  • GPU machine (ACT is not large but still benefits from GPU for training)
  • Hugging Face Hub account (if you want to push the trained model to the Hub)

Software​

In the previous section we converted the data to the LeRobot Dataset v2.1 format, so we will use LeRobot 0.3.3 for training the ACT policy.

uv venv
source .venv/bin/activate
uv pip install lerobot==0.3.3

Training script​

After setting up the environment, you can use the following script to train the ACT policy on the dataset.

lerobot-train \
  --dataset.repo_id=enactic/openarm-2-cell-pick_up_cube_mujoco-lerobot \
  --policy.type=act \
  --output_dir=outputs/train \
  --policy.device=cuda \
  --wandb.enable=false \
  --policy.repo_id=$HF_USER/act-openarm-2-cell-pick_up_cube_mujoco \
  --policy.push_to_hub True

For more details on training the ACT policy, please check the LeRobot Documentation.

Last updated on