Model Training

This guide walks you through training imitation learning models for OMY, based on datasets collected via the Web UI.

Once preparing your dataset is done, the policy model can be trained using either the Web UI or the LeRobot CLI.

You can choose one of the following options:

Option 1
Web UI (Recommended) Option 2
LeRobot CLI (Optional)

Model Training With Web UI

1. Prepare Your Dataset

The dataset to be used for training should be located at

USER PC

<your_workspace>/physical_ai_tools/docker/huggingface/lerobot/${HF_USER}/

Datasets collected using Physical AI Tools are automatically saved to that path. However, if you downloaded the dataset from a hub or copied it from another PC, you need to move the dataset to that location.

INFO

• ${HF_USER} can be any folder name you prefer.
• <your_workspace> is the directory containing physical_ai_tools

Please refer to the folder structure tree below:

<your_workspace>/physical_ai_tools/docker/huggingface/lerobot/
  ├── USER_A/           # ← ${HF_USER} folder
  │   ├── dataset_1/    # ← Dataset
  │   │   ├── data/
  │   │   ├── meta/
  │   │   └── videos/
  │   └── dataset_2/
  └── USER_B/
      └── dataset_3/

2. Setup Physical AI Tools Docker Container

WARNING

If the Physical AI Tools is already set up, you can skip this step.

If you haven't set up the Physical AI Tools Docker container, please refer to the link below for setup instructions.

Setup Physical AI Tools Docker Container

3. Train the Policy

a. Launch Physical AI Server

WARNING

If the Physical AI Tools Docker container is already running, you can skip this step.

Go to physical_ai_tools/docker directory:

USER PC

cd physical_ai_tools/docker

Enter the Physical AI Tools Docker container:

./container.sh enter

Then, launch the Physical AI Server with the following command:

USER PC 🐋 PHYSICAL AI TOOLS

ai_server

b. Open the Web UI

Open your web browser and navigate to the Web UI (Physical AI Manager).

(Refer to the Dataset Preparation > Web UI > 1. Open the Web UI for more details.)

On the Home page, select the type of robot you are using.

c. Train the Policy

Go to the Training page and follow the steps below:

• Step 1: Select the Dataset, Policy Type and Device.
• Step 2: Enter the Output Folder Name.
• Step 3: (Optional) Modify Additional Options if needed.

For more information about these options, please refer to the descriptions below.

The datasets stored in the <your_workspace>/physical_ai_tools/docker/huggingface/ directory on the host (or /root/.cache/huggingface/ inside the Docker container) will be listed automatically.

Click Start Training to begin training the policy. The training results will be saved in the physical_ai_tools/lerobot/outputs/train/ directory.

You can monitor the training loss while training is in progress.

(Optional) Uploading Checkpoints to Hugging Face

Navigate to physical_ai_tools/docker directory and enter the Docker container:

USER PC

cd physical_ai_tools/docker

./container.sh enter

Navigate to the LeRobot directory:

USER PC 🐋 PHYSICAL AI TOOLS

cd /root/ros2_ws/src/physical_ai_tools/lerobot

To upload the latest trained checkpoint to the Hugging Face Hub:

huggingface-cli upload ${HF_USER}/act_omy_test \
  outputs/train/act_omy_test/checkpoints/last/pretrained_model

This makes your model accessible from anywhere and simplifies deployment.

Model Training With LeRobot CLI

1. Prepare Your Dataset

The dataset to be used for training should be located at <your_workspace>/physical_ai_tools/docker/huggingface/lerobot/${HF_USER}/. If your dataset is in a different location, please move it to this path.

INFO

• You can replace ${HF_USER} with the folder name you used when recording your dataset.
• <your_workspace> is the directory containing physical_ai_tools

2. Train the Policy

Go to physical_ai_tools/docker directory:

USER PC

cd physical_ai_tools/docker

Enter the Physical AI Tools Docker container:

container.sh enter

Navigate to the LeRobot directory:

USER PC 🐋 PHYSICAL AI TOOLS

cd /root/ros2_ws/src/physical_ai_tools/lerobot

Once the dataset has been transferred, you can train a policy using the following command:

python -m lerobot.scripts.train \
  --dataset.repo_id=${HF_USER}/omy_test \
  --policy.type=act \
  --output_dir=outputs/train/act_omy_test \
  --policy.device=cuda \
  --log_freq=100 \
  --save_freq=1000 \
  --policy.push_to_hub=false