Model Training

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

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

Model Training

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/.cache/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 ai_server 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 ~/.cache/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.

(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_omx_test \
  outputs/train/act_omx_test/checkpoints/last/pretrained_model

This makes your model accessible from anywhere and simplifies deployment.