Skip to main content

Your First LeRobot Training Job

Submit a LeRobot behavioral cloning training job to OSMO using a HuggingFace dataset and verify that the trained policy appears in Azure ML. By the end of this recipe, you will have a trained ACT policy for the ALOHA sim insertion task.

[!NOTE] This recipe requires deployed infrastructure with OSMO running. Complete the Quickstart first.

📋 Prerequisites

RequirementDetails
InfrastructureAzure resources deployed via Terraform
OSMOControl plane and backend running
VPNConnected to private cluster (if using private AKS)
Azure CLIAuthenticated (az login)

🚀 Steps

Step 1: Preview the training configuration

Preview what will be submitted:

cd training/il/scripts
./submit-osmo-lerobot-training.sh \
-d lerobot/aloha_sim_insertion_human \
--config-preview

Review the dataset source, policy type, training hyperparameters, and Azure ML context. No job is submitted.

Step 2: Submit a training job

Submit a training run with the default ACT policy:

./submit-osmo-lerobot-training.sh \
-d lerobot/aloha_sim_insertion_human

The script submits an OSMO workflow that pulls the dataset from HuggingFace, trains an ACT policy, and logs metrics to MLflow.

Customize training hyperparameters:

./submit-osmo-lerobot-training.sh \
-d lerobot/aloha_sim_insertion_human \
--policy-type act \
--training-steps 50000 \
--batch-size 32 \
--save-freq 5000

Step 3: Train with data from Azure Blob Storage

Use --from-blob when your dataset is in Azure Storage instead of HuggingFace:

./submit-osmo-lerobot-training.sh \
-d my-org/my-dataset \
--from-blob \
--storage-account <your-storage-account> \
--storage-container datasets \
--blob-prefix my-dataset/v1

The script downloads the dataset from Blob Storage using managed identity credentials before training starts.

Step 4: Monitor training progress

OSMO UI: Open the OSMO dashboard to view workflow status, pod logs, and real-time metrics. See Accessing OSMO for connection instructions (VPN or port-forward).

Azure ML Studio: Navigate to your workspace at ml.azure.com, open the Jobs section, and select the MLflow experiment. Training metrics (loss, gradient norm, learning rate) stream in real time as the job progresses.

To view results in detail:

  1. Open ml.azure.com and select your workspace
  2. Navigate to Jobs in the left sidebar
  3. Find your experiment (named after the task, e.g., lerobot-act-training) and select it
  4. Click the latest run to open the run detail page
  5. Select the Metrics tab to view training curves (loss, gradient norm, learning rate) — use the chart controls to overlay multiple metrics or smooth noisy curves
  6. Select the Outputs + logs tab to view stdout/stderr logs from the training container
  7. If --register-checkpoint was used, navigate to Models in the left sidebar to confirm the registered model and its version

CLI (optional): Check workflow status via the OSMO CLI:

osmo workflow list

Step 5: Register the trained model (optional)

Register the trained checkpoint to Azure ML for versioned model management:

./submit-osmo-lerobot-training.sh \
-d lerobot/aloha_sim_insertion_human \
--register-checkpoint my-act-policy

The --register-checkpoint flag triggers automatic model registration after training completes.

✅ Verify

The recipe succeeded when:

  • OSMO training pod reached Completed status
  • MLflow experiment shows training loss decreasing over steps
  • Model artifacts exist in Azure ML (if --register-checkpoint was used)

⚙️ Configuration Reference

ParameterDefaultDescription
-d, --dataset-repo-id(required)HuggingFace dataset repository
--policy-typeactPolicy architecture (act or diffusion)
--training-steps100000Total training iterations
--batch-size32Training batch size
--learning-rate1e-4Optimizer learning rate
--save-freq5000Checkpoint save frequency
--val-split0.1Validation split ratio
--from-blob(disabled)Use Azure Blob Storage as data source
--register-checkpoint(none)Model name for Azure ML registration

See Scripts Reference for the full parameter table.

🤖 Crafted with precision by ✨Copilot following brilliant human instruction, then carefully refined by our team of discerning human reviewers.