Physical AI Platform (MVP)

A modular, scalable platform for Robot Foundation Models, featuring a closed-loop workflow: Data Collection → World Model Training → Planning (MPC) → Evaluation

This repository implements the core infrastructure for a "Platform-level Narrative," allowing for rapid iteration on world models and policies using a structured observation schema.

🏗 Architecture

The platform operates on a continuous feedback loop:

1. Sim Rollout: Collect structured episodes using Isaac Lab (or mock envs). Use standard schema (Proprio, EE, Objects).
2. World Model: Train a dynamics model ($S_{t+1} | S_t, A_t$) on collected data.
3. Offline Eval: Assess model quality via rollout error (MSE) and success proxies.
4. Planning (MPC): Use the trained world model to verify performance in-loop (Model Predictive Control).
5. Failure Mining: Automatically cluster failure modes (Collision, Drop, Timeout) to target data collection.

---

🚀 Workflow

1. Data Collection (Rollout)

Generate synthetic episodes. Currently uses a Mock Environment (replace with real Isaac Lab later).

# Generate 32 parallel envs for 100 steps each
pai sim-rollout --override sim.num_envs=32 sim.horizon=100

Output: runs/sim-rollout-YYYY-MM-DD_HH-MM-SS/episodes/*.jsonl

2. Train World Model

Train a transition model to predict the next state given current state and action.

# Train on the data generated above
pai train-world-model --override data.episode_glob="runs/**/episodes/*.jsonl" train.epochs=10

Output: Checkpoints in runs/train-world-model-.../ckpts/

3. Offline Evaluation

Evaluate the trained model on rollout error and success rates without running simulation.

# Evaluate a specific checkpoint
pai eval-offline --override eval.model_ckpt="runs/train-world-model-.../ckpts/ckpt_final.pt"

4. Failure Mining 🔍

Analyze failed episodes to understand why the robot failed (Collision vs. Object Drop vs. Timeout).

pai failure-mining

Output: Generates reports/failure_report.html helping you visualize hard cases.

5. Planning with World Model (MPC)

Close the loop: Use the trained World Model to control the robot.

# Run rollout using MPC policy instead of random
pai sim-rollout --override sim.policy.type=mpc sim.policy.model_ckpt="runs/train-world-model-.../ckpts/ckpt_final.pt"

---

📂 Project Structure

physical-ai-platform/
├── configs/              # Hydra configurations
│   ├── sim/              # Simulation & Policy configs
│   ├── train/            # Training hyperparameters
│   └── eval/             # Evaluation settings
├── src/pai/
│   ├── schema/           # Universal Episode/Step Schema (Pydantic)
│   ├── sim/              # Environment adapters & Rollout loops
│   ├── models/           # World Models (RSSM/Transformer)
│   ├── planning/         # MPC & Search Policies
│   ├── training/         # PyTorch training loops
│   ├── eval/             # Offline Eval & Failure Mining
│   └── data/             # Dataset loaders
└── runs/                 # Experiment outputs (Auto-generated)

🛠 Integration Guide

Switching to Real Isaac Lab

1. Modify src/pai/sim/isaaclab_env.py:
   • Replace the mock reset() and step() with calls to the actual Isaac Lab vectorized environment.
   • Ensure the returned obs dictionary matches the structure expected by obs_adapter.yaml.
2. Update Config:
   • Edit configs/sim/isaaclab_cabinet.yaml to match your task interface.

Adding a New Model

1. Add your model class in src/pai/models/.
2. Update src/pai/training/train_world_model.py to instantiate your model based on config.

---

📊 Episode Schema

All data flows satisfy the strict Episode schema defined in src/pai/schema/episode.py.

• Proprioception: Joint positions (q), velocities (qd).
• End-Effector: Pose (p, q), velocities.
• Objects: State of relevant objects in the scene.
• Events: Sparse booleans for semantic events (grasped, dropped, collision).

This schema ensures that the World Model is agnostic to the underlying simulator implementation.