microsoft/osmo-lerobot-training

OSMO LeRobot Training

Submit, monitor, analyze, and evaluate LeRobot behavioral cloning training workflows on the OSMO platform. Covers the full lifecycle: job submission, log streaming, Azure ML metric retrieval, training summary generation, and post-training inference evaluation.

Read the skill file .github/skills/osmo-lerobot-training/SKILL.md for parameter defaults, GPU configuration, and training duration estimates. Read references/DEFAULTS.md for known datasets, GPU profiles, and Azure environment auto-resolution.

Prerequisites

| Requirement | Purpose | |-------------|---------| | osmo CLI | Workflow submission and monitoring | | az CLI | Azure authentication and model registry | | terraform | Infrastructure output resolution | | zip, base64 | Training payload packaging | | Python 3.11+ with azure-ai-ml, mlflow | Metric retrieval from Azure ML |

Authentication must be configured before any OSMO or Azure ML operations:

az login
osmo login <service-url> --method dev --username guest

Quick Start

Train from Azure Blob Storage (typical production flow)

scripts/submit-osmo-lerobot-training.sh \
  -d my-robot-dataset \
  --from-blob \
  --storage-account mystorageaccount \
  --blob-prefix my-robot-dataset \
  --no-val-split \
  --steps 100000 \
  --batch-size 32 \
  --learning-rate 1e-4 \
  --save-freq 10000 \
  -j my-robot-act-train \
  --experiment-name my-robot-training \
  -r my-robot-act-model

Train from HuggingFace Hub

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

Run Continuous Eval During Training (preferred)

Start the background poller immediately after submitting training. It watches AzureML for new checkpoint versions and submits an inference job per version automatically, stopping when training reaches a terminal state.

# Launch in the background — runs until training completes
nohup scripts/poll-and-eval-checkpoints.sh \
  --model-name my-robot-act-model \
  --training-workflow-id lerobot-training-32 \
  --blob-prefix my-robot-dataset \
  --job-prefix my-robot-eval \
  --experiment-name my-robot-inference \
  --poll-interval 60 \
  --max-concurrent 2 \
  > /tmp/my-robot-eval.log 2>&1 & disown

# Monitor the poller
tail -f /tmp/my-robot-eval.log

The poller caps concurrent inference workflows at --max-concurrent (default 2) to avoid cluster saturation. Submitted versions are tracked in /tmp/<model-name>-submitted-versions.txt.

Run a Single Inference Job

# OSMO inference (GPU, evaluates against the same dataset)
scripts/submit-osmo-lerobot-inference.sh \
  --from-aml-model \
  --model-name my-robot-act-model \
  --model-version 3 \
  --from-blob-dataset \
  --storage-account mystorageaccount \
  --blob-prefix my-robot-dataset \
  --mlflow-enable \
  --eval-episodes 10 \
  -j my-robot-eval \
  --experiment-name my-robot-inference

# Local inference (CPU/MPS, for quick validation)
python scripts/run-local-lerobot-inference.py \
  --model-name my-robot-act-model \
  --model-version 3 \
  --dataset-dir /path/to/local/dataset \
  --episodes 5 \
  --output-dir outputs/local-eval \
  --device cpu

Post-Submission Browser Monitoring

After every successful training or inference submission, open the OSMO workflow page in VS Code's SimpleBrowser so the user can track progress and access logs directly.

Steps:

1. Capture the workflow ID from the submission output (the line Workflow ID - <id>).
2. Construct the URL: http://10.0.5.7/workflows/<workflow-id>.
3. Open it with the open_browser_page tool (VS Code SimpleBrowser).
4. Tell the user that the Logs tab on that page streams live output per task (e.g., lerobot-train, lerobot-infer).

Example — after training submission output:

Workflow ID - lerobot-training-31
Workflow Overview - http://10.0.5.7/workflows/lerobot-training-31

Open: http://10.0.5.7/workflows/lerobot-training-31

Example — after inference submission output:

Workflow ID - lerobot-inference-20
Workflow Overview - http://10.0.5.7/workflows/lerobot-inference-20

Open: http://10.0.5.7/workflows/lerobot-inference-20

The page has a Logs tab with per-task log streams. For training, select the lerobot-train task. For inference, select the lerobot-infer task. Use the OSMO CLI (osmo workflow logs <id> -t <task> -n 100) as a fallback when the browser is not reachable.

Azure ML Portal Monitoring (Playwright)

After submitting a training job, and whenever the background eval poller reports a new inference job, open the Azure ML portal with Playwright to view live metrics and trajectory plots. Use mcp_playwright_browser_navigate, mcp_playwright_browser_snapshot, mcp_playwright_browser_click, and mcp_playwright_browser_take_screenshot.

Training Metrics — Open Immediately After Submission

After the training job is submitted, navigate to the training experiment page and open the Metrics tab:

1. Construct the experiment URL from Azure environment variables in scripts/.env:

   https://ml.azure.com/experiments/{experiment_name}?wsid=/subscriptions/{AZURE_SUBSCRIPTION_ID}/resourceGroups/{AZURE_RESOURCE_GROUP}/providers/Microsoft.MachineLearningServices/workspaces/{AZUREML_WORKSPACE_NAME}
   

2. Call mcp_playwright_browser_navigate with that URL.

3. Call mcp_playwright_browser_snapshot to confirm the page loaded and identify the latest run row in the table.

4. Click the first (most recent) run link.

5. On the run detail page, call mcp_playwright_browser_snapshot to locate the Metrics tab.

6. Click Metrics.

7. Call mcp_playwright_browser_take_screenshot and show the live training curves to the user.

Key metrics to surface: train/loss, train/learning_rate (confirm 1e-04, not 1e-05), train/grad_norm, gpu_percent.

Refresh by calling mcp_playwright_browser_navigate again on the same URL at any time.

See references/REFERENCE.md for exact click paths, tab selectors, and screenshot guidance.

Inference / Eval Plots — Open When Poller Submits a Job

While the background eval poller is running, monitor the poller log and navigate to Azure ML to view trajectory plots as each inference job completes:

1. Tail the poller log to detect a new inference submission:

   tail -n 30 /tmp/<model-name>-eval.log | grep -E "Submitting|Workflow ID"
   

2. Construct the inference experiment URL using the --experiment-name passed to the poller:

   https://ml.azure.com/experiments/{inference_experiment_name}?wsid=/subscriptions/{AZURE_SUBSCRIPTION_ID}/resourceGroups/{AZURE_RESOURCE_GROUP}/providers/Microsoft.MachineLearningServices/workspaces/{AZUREML_WORKSPACE_NAME}
   

3. Call mcp_playwright_browser_navigate with that URL.

4. Call mcp_playwright_browser_snapshot to identify the latest run row (most recently submitted checkpoint eval).

5. Click that run.

6. On the run detail page, click the Images tab.

7. Call mcp_playwright_browser_take_screenshot and show the trajectory plots to the user.

The Images tab contains per-episode trajectory plots logged by the inference job (episode_NNN_trajectory.png and eval_summary.png). They appear after the OSMO inference workflow reaches completed status. If images are not yet present, check osmo workflow query <inference-workflow-id> and wait for completed.

Parameters Reference

Training Submission Parameters

| Parameter | Flag | Default | Description | |-----------|------|---------|-------------| | Dataset repo ID | -d, --dataset | (required) | HuggingFace dataset or blob dataset name | | Policy type | -p, --policy | act | act or diffusion | | Job name | -j, --job-name | lerobot-act-training | Unique job identifier | | Training steps | --steps | 100000 | Total training iterations | | Batch size | --batch-size | 32 | Training batch size (64 for 48GB GPUs) | | Learning rate | --learning-rate | 1e-4 | Maps to --policy.optimizer_lr internally | | Save frequency | --save-freq | 5000 | Checkpoint interval (model registered at each) | | Validation split | --val-split | 0.1 | Ratio for train/val split | | No val split | --no-val-split | — | Disable validation splitting | | Register checkpoint | -r | (none) | Model name for Azure ML registration | | From blob | --from-blob | false | Use Azure Blob Storage as data source | | Storage account | --storage-account | (terraform) | Azure Storage account name | | Blob prefix | --blob-prefix | (none) | Blob path prefix for dataset |

Inference Submission Parameters

| Parameter | Flag | Default | Description | |-----------|------|---------|-------------| | Policy repo ID | --policy-repo-id | (required) | HuggingFace repo, or use --from-aml-model | | From AML model | --from-aml-model | false | Load from AzureML model registry | | Model name | --model-name | (none) | AzureML model registry name | | Model version | --model-version | (none) | AzureML model version | | Dataset repo ID | -d, --dataset-repo-id | (none) | HuggingFace dataset | | From blob dataset | --from-blob-dataset | false | Download dataset from Azure Blob | | Eval episodes | --eval-episodes | 10 | Number of episodes to evaluate | | MLflow enable | --mlflow-enable | false | Log trajectory plots to AzureML |

Continuous Evaluation Parameters (poll-and-eval-checkpoints.sh)

| Parameter | Flag | Default | Description | |-----------|------|---------|-------------| | Model name | --model-name | (required) | AzureML model registry name to watch | | Training workflow | --training-workflow-id | (required) | OSMO workflow ID of the training job | | Blob prefix | --blob-prefix | (required) | Blob path prefix for the evaluation dataset | | Storage account | --storage-account | (from .env) | Azure Storage account | | Eval episodes | --eval-episodes | 10 | Episodes per inference run | | Job prefix | --job-prefix | (from model name) | Prefix for inference job names | | Experiment name | --experiment-name | (from model name) | MLflow experiment for inference runs | | Poll interval | --poll-interval | 60 | Seconds between AzureML registry polls | | Max concurrent | --max-concurrent | 2 | Max simultaneous inference workflows |

GPU Configuration Guidelines

| GPU | VRAM | Recommended Batch Size | Notes | |-----|------|----------------------|-------| | A10 | 24GB | 32 | Standard configuration | | RTX PRO 6000 | 48GB | 64 | Requires mig.strategy: single | | H100 | 80GB | 128 | Standard MIG disabled |

Azure ML Context

Resolved from CLI flags > environment variables > Terraform outputs:

| Variable | Flag | Env Var | |----------|------|---------| | Subscription ID | --azure-subscription-id | AZURE_SUBSCRIPTION_ID | | Resource group | --azure-resource-group | AZURE_RESOURCE_GROUP | | Workspace name | --azure-workspace-name | AZUREML_WORKSPACE_NAME |

Training Completion Estimation

Estimate training duration based on dataset and configuration:

| Dataset Size | Steps | GPU | Approximate Duration | |-------------|-------|-----|---------------------| | 20K frames / 64 episodes | 10,000 | A10 | ~30 minutes | | 20K frames / 64 episodes | 100,000 | A10 | ~5 hours | | 80K frames / 174 episodes | 100,000 | A10 | ~8 hours | | 20K frames / 64 episodes | 100,000 | RTX PRO 6000 | ~3 hours |

Checkpoints are registered to AzureML at every --save-freq interval. Jobs may be evicted on spot GPU instances — checkpoints already registered remain available for inference even if the job is interrupted.

OSMO CLI Reference

See references/REFERENCE.md for full CLI and SDK documentation.

osmo workflow query <workflow-id>
osmo workflow logs <workflow-id> -n 100
osmo workflow logs <workflow-id> --error
osmo workflow list
osmo workflow cancel <workflow-id>

Checkpoint Poller Commands

# Start continuous eval loop in background
nohup scripts/poll-and-eval-checkpoints.sh \
  --model-name <model-name> \
  --training-workflow-id <workflow-id> \
  --blob-prefix <dataset-blob-prefix> \
  > /tmp/<model-name>-eval.log 2>&1 & disown

# Monitor poller
tail -f /tmp/<model-name>-eval.log

# Check which versions have been submitted
cat /tmp/<model-name>-submitted-versions.txt

# Stop the poller early
pkill -f poll-and-eval-checkpoints

Key Metrics Logged

| Metric | Description | |--------|-------------| | train/loss | Training loss per step | | train/grad_norm | Gradient norm | | train/learning_rate | Current learning rate (verify 1e-4 not 1e-5) | | val/loss | Validation loss (when val split enabled) | | gpu_percent | GPU utilization (when system metrics enabled) |

Troubleshooting

| Symptom | Likely Cause | Resolution | |---------|-------------|------------| | lr: 1e-05 in logs | LEARNING_RATE not mapped | Verify train.py maps to --policy.optimizer_lr | | KeyError: chunk_index | v3.0 dataset not converted | Verify download_dataset.py has patch_info_paths() | | codebase_version warning | Dataset still marked v3.0 | Verify patch_info_paths() sets codebase_version = "v2.1" | | CUDA_ERROR_NO_DEVICE | MIG strategy misconfigured | Set mig.strategy: single for vGPU nodes | | VM eviction mid-training | Spot GPU preempted | Checkpoints already registered to AML survive eviction | | ImportError: patch_info_paths | Payload missing training fixes | Ensure training/il/ includes download_dataset.py with patch_info_paths | | OOM during training | Batch size too large | Reduce --batch-size (32 for 24GB, 64 for 48GB) | | Poller exits immediately | Training workflow already terminal | Check osmo workflow query <id>; rerun poller or submit inference manually | | Poller stalls at max-concurrent | Inference jobs not finishing | Check inference workflow status; increase --max-concurrent or cancel stuck jobs | | Many pending inference jobs after stopping poller | Poller submitted jobs faster than cluster could drain | osmo workflow list only returns the last 12 — iterate over expected ID range to cancel all: for id in $(seq <first> <last>); do osmo workflow cancel lerobot-inference-$id; done | | info: command not found in poller | common.sh not sourced | Verify scripts/lib/common.sh exists and is readable |

See references/REFERENCE.md for detailed debugging commands.

Brought to you by microsoft/physical-ai-toolchain