Build Physical AI with Isaac Lab: Data Augmentation, Imitation Learning & Evaluation

Physical AI Overview

Physical AI aims to train large‑model architectures such as Transformers and Diffusion models so that they can guide embodied robots to accomplish tasks in the real world.

Toolchain Integration on PAI

Alibaba Cloud PAI tightly integrates the Isaac Sim, Isaac Lab, and Cosmos tool stacks, and provides ready‑to‑run notebooks that demonstrate the full Physical AI pipeline.

Best‑Practice 1: Action Data Augmentation & Imitation Learning with Isaac

The typical workflow consists of five stages: manual teleoperation demos, data augmentation, data enhancement, imitation‑learning training, and model evaluation.

Manual demonstration – record expert tele‑operated trajectories in simulation or on a real robot.

Data augmentation – randomize initial poses and trajectories to generate varied samples.

Data enhancement – use generative models to improve video texture, lighting, and realism.

Imitation learning – train a BC‑RNN policy on the augmented and enhanced dataset.

Model evaluation – run the policy in diverse simulated conditions and measure success rates.

Manual Demonstrations

In the notebook a few tele‑operated demos are collected with Isaac Lab. Example command to launch the tele‑operation interface:

# Use keyboard teleoperation to label dataset
cmd = f"PUBLIC_IP=$(curl -s ifconfig.me) /workspace/isaaclab/isaaclab.sh -p /workspace/isaaclab/scripts/environments/teleoperation/teleop_se3_agent.py --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --num_envs 1 --teleop_device keyboard --livestream 1"
print(f"Executing command: {cmd}")
!{cmd}

After recording, the annotate_demos.py script adds sub‑task labels required by Isaac‑Lab Mimic.

Data Augmentation with Isaac‑Lab Mimic

Using the Mimic pipeline, a small set of expert demos is expanded into a large synthetic dataset. The number of parallel environments ( --num_envs) and the total number of generated trials ( --generation_num_trials) can be tuned, e.g., 8 environments generating 1 000 trials.

# Generate dataset with Mimic
os.environ["ACCEPT_EULA"] = "Y"
cmd = f"PUBLIC_IP=$(curl -s ifconfig.me) /workspace/isaaclab/isaaclab.sh -p /workspace/isaaclab/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
    --enable_cameras --headless --num_envs 8 --generation_num_trials 1000 \
    --input_file {annotated_dataset_path} --output_file {mimic_dataset_path} \
    --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0 --livestream 1"
print(f"Executing command: {cmd}")
!{cmd}

For large‑scale jobs, DLC can launch distributed Mimic tasks with Ray support.

Data Enhancement with Cosmos‑Transfer

Generated MP4 videos are often not photorealistic. The Cosmos‑Transfer‑1‑7B model is deployed to enhance texture, lighting, and detail.

import json, os, requests, gradio_client
# simplified Cosmos request creation and upload logic omitted for brevity

Imitation Learning

A simple BC‑RNN model is trained on the augmented and enhanced dataset using the Robomimic framework.

# Train BC‑RNN model
cmd = f"cd {workspace_dir} && /workspace/isaaclab/isaaclab.sh -p /workspace/isaaclab/scripts/imitation_learning/robomimic/train.py \
    --task {task_name} --algo bc --dataset {dataset_path} --name {model_name}"
print(f"Training command: {cmd}")
!{cmd}

Distributed Training with DLC

DLC jobs can be created to run the training on multiple GPUs, specifying resource configurations for master and worker pods.

# Create DLC job (PyTorchJob example)
create_job_resp = dlc_client.create_job(CreateJobRequest().from_map({
    'WorkspaceId': workspace_id,
    'DisplayName': display_name,
    'JobType': 'PyTorchJob',
    'ResourceId': resource_quota_id,
    'JobSpecs': [{
        'Type': 'Master',
        'Image': image_uri,
        'PodCount': 1,
        'ResourceConfig': {'CPU': '48', 'Memory': '256Gi', 'GPU': '0'}
    }, {
        'Type': 'Worker',
        'Image': image_uri,
        'PodCount': 2,
        'ResourceConfig': {'CPU': '56', 'Memory': '448Gi', 'GPU': '4'}
    }],
    'UserCommand': f"cd {workspace_dir} && /workspace/isaaclab/isaaclab.sh -p /workspace/isaaclab/scripts/imitation_learning/robomimic/train.py --task {task_name} --algo bc --dataset {dataset_path} --name {model_name} && sleep 30"
}))
job_id = create_job_resp.body.job_id
wait_for_job_to_terminate(dlc_client, job_id)

Model Evaluation

Trained policies are evaluated under varied lighting, texture, and environmental conditions to compute success rates.

# Evaluate model robustness
cmd = f"cd {workspace_dir} && /workspace/isaaclab/isaaclab.sh -p /workspace/isaaclab/scripts/imitation_learning/robomimic/robust_eval.py \
    --task {task_name} --input_dir {model_dir_path}/models --log_dir {log_dir} \
    --enable_cameras --livestream 1 --seeds 0 --num_rollouts 15 --headless"
print(f"Evaluating command: {cmd}")
!{cmd}

Summary

The end‑to‑end workflow demonstrates how to collect a few tele‑operated demos, expand them with Isaac‑Lab Mimic, enhance visual fidelity with Cosmos‑Transfer, train a BC‑RNN imitation‑learning policy, and evaluate it across diverse simulated conditions, achieving higher success rates especially under lighting variations.