{ "cells": [ { "cell_type": "markdown", "id": "udlwdep9gm", "metadata": {}, "source": [ "# LeRobot Data: Download, Convert, Train, and Evaluate\n", "\n", "This notebook walks through the full workflow for using [LeRobot](https://github.com/huggingface/lerobot) datasets with VLA Foundry: downloading a dataset from HuggingFace Hub, converting it to WebDataset format, training a diffusion policy model, and running inference. We set the training scale to a small number to make the tutorial but this is not enough to get a meaningful policy. The user is welcome to modify the tutorial as they see fit.\n", "\n", "We use `lerobot/pusht` as the running example throughout this tutorial. PushT is a simple 2D pushing task that makes a great starting point." ] }, { "cell_type": "markdown", "id": "qo9x6aplrsj", "metadata": {}, "source": [ "## Prerequisites\n", "\n", "- Python 3.12+ with VLA Foundry installed via `uv` (see [Installation](../README.md#installation))\n", "- The **`Python (vla_foundry)`** Jupyter kernel \u2014 install it once from the repo root:\n", " ```bash\n", " bash tutorials/install_kernel.sh\n", " ```\n", " This syncs dependencies (preprocessing, gym-pusht, etc.) and registers the kernel.\n", "- Restart your notebook kernel after installing the vla_foundry kernel and select it for execution.\n", "- For training, a machine with one or more GPUs\n", "- (Optional) AWS credentials configured if you want to read/write data on S3\n", "\n", "**Storage:** This tutorial downloads ~12MB of LeRobot data from HuggingFace and generates ~2GB of preprocessed WebDataset shards. Ensure you have at least 3GB of free disk space before starting.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3ezsf6xwp6i", "metadata": {}, "outputs": [], "source": [ "import os\n", "import subprocess\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n", "\n", "# Change to the repo root (works regardless of where the notebook is opened from)\n", "repo_root = subprocess.check_output([\"git\", \"rev-parse\", \"--show-toplevel\"], text=True).strip()\n", "os.chdir(repo_root)\n", "print(f\"Working directory: {os.getcwd()}\")" ] }, { "cell_type": "markdown", "id": "uwpnxdkjm3d", "metadata": {}, "source": [ "**Stop here.** Switch to the **\"Python (vla_foundry)\"** kernel in your notebook UI (Kernel \u2192 Change Kernel) and restart it before running the cells below." ] }, { "cell_type": "markdown", "id": "169zui94qoi", "metadata": {}, "source": [ "## 1. Download a LeRobot Dataset\n", "\n", "### Browsing Available Datasets\n", "\n", "LeRobot datasets are hosted on HuggingFace Hub. You can browse them at:\n", "- [https://huggingface.co/lerobot](https://huggingface.co/lerobot)\n", "\n", "Some well-known datasets include:\n", "- `lerobot/pusht` -- Push-T task (2D)\n", "- `lerobot/aloha_sim_transfer_cube_human` -- ALOHA bimanual cube transfer (simulation)\n", "- `lerobot/aloha_sim_insertion_human` -- ALOHA peg insertion (simulation)\n", "- `lerobot/xarm_lift_medium_replay` -- xArm lift task\n", "\n", "### Downloading with huggingface-cli\n", "\n", "> **Important -- Use `--revision v2.0`.** The default version on HuggingFace is now v3.0, which uses a different directory layout that the converter does not support yet. Always specify `--revision v2.0` when downloading datasets for use with VLA Foundry." ] }, { "cell_type": "markdown", "id": "dac249cd", "metadata": {}, "source": [ "> \u26a0\ufe0f **Heavy download & preprocessing ahead.** The next cells download `lerobot/pusht` (~1 GB) from HuggingFace and produce ~2 GB of preprocessed WebDataset shards. Both steps are idempotent (safe to re-run; skips files already on disk). Under `jupyter nbconvert`, pass `--ExecutePreprocessor.timeout=1800` or run the download from a terminal first.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "xcwzdxp84oh", "metadata": {}, "outputs": [], "source": [ "from huggingface_hub import snapshot_download\n", "\n", "snapshot_download(\n", " \"lerobot/pusht\",\n", " repo_type=\"dataset\",\n", " revision=\"v2.0\",\n", " local_dir=\"./tutorials/data/lerobot/pusht\",\n", ")" ] }, { "cell_type": "markdown", "id": "ymjfu6o8uto", "metadata": {}, "source": [ "### LeRobot Data Format (v2.0)\n", "\n", "After downloading, the dataset directory has this structure (v2.0 format):\n", "\n", "```\n", "pusht/\n", " meta/\n", " info.json # Dataset metadata (fps, features, shapes, etc.)\n", " episodes.jsonl # Per-episode metadata (length, task, index)\n", " tasks.jsonl # Task descriptions (language instructions)\n", " data/\n", " chunk-000/\n", " episode_000000.parquet # Episode data (states, actions)\n", " episode_000001.parquet\n", " ...\n", " videos/ # (Optional) Video files for camera observations\n", " chunk-000/\n", " observation.image/\n", " episode_000000.mp4\n", " ...\n", "```\n", "\n", "Key files:\n", "- **`meta/info.json`**: Contains the dataset FPS, feature names, data shapes, and number of episodes\n", "- **`meta/episodes.jsonl`**: One JSON object per line, each describing an episode\n", "- **`data/chunk-XXXX/episode_XXXXXX.parquet`**: Parquet files containing observation and action data\n", "- **`videos/`**: (Optional) Video files organized by camera name and chunk" ] }, { "cell_type": "markdown", "id": "pilpmtfemkf", "metadata": {}, "source": [ "## 2. Convert to WebDataset Format\n", "\n", "VLA Foundry uses WebDataset tar shards for training. The conversion script handles LeRobot datasets with `--type lerobot`.\n", "\n", "### Inspecting Your Dataset\n", "\n", "Before running preprocessing, inspect the dataset to discover the correct column names for observations, actions, and images. Column names vary between datasets, so always check before proceeding." ] }, { "cell_type": "code", "execution_count": null, "id": "wi3xkjpr52s", "metadata": {}, "outputs": [], "source": [ "import pandas as pd\n", "\n", "df = pd.read_parquet(\"./tutorials/data/lerobot/pusht/data/chunk-000/episode_000000.parquet\")\n", "print(\"Columns:\", df.columns.tolist())\n", "print(\"\\nFirst few rows:\")\n", "df.head()" ] }, { "cell_type": "markdown", "id": "pv643o9srw", "metadata": {}, "source": [ "**For video-based datasets** (like `pusht`), image/camera columns will **not** appear in the parquet file. Instead, camera names come from the subdirectories under `videos/`:" ] }, { "cell_type": "code", "execution_count": null, "id": "pq2zelgogy9", "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "# For video-based datasets, camera names come from the videos/ directory\n", "video_dirs = os.listdir(\"./tutorials/data/lerobot/pusht/videos/chunk-000/\")\n", "print(\"Camera names:\", video_dirs)" ] }, { "cell_type": "markdown", "id": "30vn7cdtuzs", "metadata": {}, "source": [ "### Basic Conversion (Local)\n", "\n", "The example below shows the preprocessing command for converting a LeRobot dataset to webdataset tar shards.\n", "\n", "Arguments that accept lists use Python list literal syntax as a string (e.g., `\"['observation.state']\"`).\n", "\n", "### Key Arguments\n", "\n", "| Argument | Description |\n", "|---|---|\n", "| `--type \"lerobot\"` | Selects the LeRobot converter |\n", "| `--source_episodes` | List of paths to LeRobot dataset root directories (local or S3) |\n", "| `--output_dir` | Where to write the converted WebDataset shards (local or S3) |\n", "| `--camera_names` | List of camera/image names (e.g., `observation.image`) |\n", "| `--observation_keys` | List of observation column names to include (e.g., `observation.state`) |\n", "| `--action_keys` | List of action column names (e.g., `action` or `actions`) |\n", "| `--config_path` | Path to a YAML file with windowing/preprocessing parameters |\n", "| `--samples_per_shard` | Number of training samples per tar shard (default: 128) |\n", "\n", "### Understanding the Preprocessing Config\n", "\n", "The preprocessing config at `vla_foundry/config_presets/data/robotics_preprocessing_params_1past_14future.yaml` defines how raw episodes are windowed into training samples. Key parameters:\n", "\n", "- **`past_lowdim_steps: 1`, `future_lowdim_steps: 14`** -- Each training sample contains a window of 16 timesteps: 1 past frame + the current frame + 14 future frames. The name \"1past_14future\" reflects this. Actions and proprioception are stored for all 16 timesteps.\n", "- **`image_indices: [-1, 0]`** -- Two images are extracted per sample: one at timestep -1 (the past frame) and one at timestep 0 (the current/anchor frame). This gives the model a sense of motion.\n", "- **`resize_images_size: [384, 384]`** -- Images are resized to 384x384 pixels using center cropping (`image_resizing_method: center_crop`).\n", "- **`stride: 1`** -- Sliding window moves by 1 timestep, so consecutive frames generate overlapping samples, maximizing data utilization.\n", "- **`padding_strategy: copy`** -- At episode boundaries, missing frames are filled by copying the nearest available frame (up to `max_padding_left: 3` and `max_padding_right: 15`).\n", "\n", "**Timing:** On a single GPU machine, data conversion takes approximately 2-3 minutes for the PushT dataset (~25k frames across 206 episodes)." ] }, { "cell_type": "code", "execution_count": null, "id": "b056v2jnkba", "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "output_dir = os.path.abspath(\"./tutorials/data/preprocessed/pusht/\")\n", "source_dir = os.path.abspath(\"./tutorials/data/lerobot/pusht/\")\n", "\n", "# Point Ray workers to the real venv so uv doesn't create a new one (avoids build race condition)\n", "os.environ[\"UV_PROJECT_ENVIRONMENT\"] = os.path.abspath(\".venv\")\n", "\n", "!MPLBACKEND= uv run --group preprocessing python vla_foundry/data/preprocessing/preprocess_robotics_to_tar.py \\\n", " --type \"lerobot\" \\\n", " --source_episodes \"['{source_dir}']\" \\\n", " --output_dir {output_dir} \\\n", " --camera_names \"['observation.image']\" \\\n", " --observation_keys \"['observation.state']\" \\\n", " --action_keys \"['action']\" \\\n", " --samples_per_shard 100 \\\n", " --config_path \"vla_foundry/config_presets/data/robotics_preprocessing_params_1past_14future.yaml\"" ] }, { "cell_type": "markdown", "id": "44ggmymx13", "metadata": {}, "source": [ "### Output Structure\n", "\n", "After conversion, the output directory contains:\n", "\n", "```\n", "tutorials/data/preprocessed/pusht/\n", " shards/\n", " manifest.jsonl # Manifest listing all shards and their sample counts\n", " stats.json # Dataset statistics (mean, std, min, max per field)\n", " shard_000000.tar # WebDataset tar shards\n", " shard_000001.tar\n", " ...\n", " episodes/\n", " manifest.jsonl # Episode-based manifest\n", " stats.json # Same statistics (copy)\n", " episode_*.tar # Episode-based tar shards\n", " frames/\n", " manifest.jsonl # Frame-level manifest\n", " frame_*.tar # Per-timestep frame metadata\n", "```\n", "\n", "What each directory contains:\n", "- **`shards/`** -- The main training data. Each tar file is a WebDataset shard containing images (JPEG) and low-dimensional data (NPZ files with actions, proprioception, and masks) for a batch of training samples. The `manifest.jsonl` and `stats.json` here are referenced directly by the training config.\n", "- **`episodes/`** -- Per-episode tar archives containing episode-level metadata such as language instructions and episode boundaries. Used for episode-aware sampling and evaluation.\n", "- **`frames/`** -- Per-timestep frame metadata used for windowed sampling. Contains the mapping from frame indices back to episodes, enabling the sliding-window sample construction.\n", "\n", "The key files for training are:\n", "- **`shards/manifest.jsonl`** -- referenced by `dataset_manifest` in the training config\n", "- **`shards/stats.json`** -- referenced by `dataset_statistics` in the training config\n", "\n", "### Additional Options\n", "\n", "**Ray Parallelization** -- For large datasets, add `--ray_address local --ray_num_cpus 16` to parallelize episode processing across multiple CPU cores.\n", "\n", "**Writing to S3** -- Replace `--output_dir ./tutorials/data/preprocessed/pusht/` with an S3 path like `--output_dir s3://my-bucket/vla_foundry_datasets/lerobot/pusht/`.\n", "\n", "**Multi-Camera** -- For datasets with multiple cameras, list all camera names: `--camera_names \"['cam_left', 'cam_right', 'cam_wrist']\"`." ] }, { "cell_type": "markdown", "id": "cq91i9thw64", "metadata": {}, "source": [ "## 3. Train a Model\n", "\n", "After preprocessing, you can train a model using the converted data. VLA Foundry uses a YAML config system powered by draccus.\n", "\n", "### Creating a Training Config\n", "\n", "Create a YAML config file that references your converted data. Here is an example for training a diffusion policy on the PushT dataset.\n", "\n", "**Note:** The `<<: !include` syntax is a draccus feature that merges in fields from another YAML file. This lets you reuse common presets for models, data pipelines, and hyperparameters.\n", "\n", "> **Important -- Temporal config must match between preprocessing and training.**\n", "> The `past_lowdim_steps` and `future_lowdim_steps` values in your preprocessing config must match the `lowdim_past_timesteps` and `lowdim_future_timesteps` in the training data config.\n", "\n", "### Available Model Presets\n", "\n", "Model presets are located in `vla_foundry/config_presets/models/`. Some commonly used ones:\n", "\n", "| Preset | Description |\n", "|---|---|\n", "| `diffusion_policy.yaml` | Diffusion/flow matching policy with CLIP vision backbone |\n", "| `transformer_tiny.yaml` | Tiny transformer (for testing/debugging) |\n", "| `transformer_100m.yaml` | 100M parameter transformer |\n", "| `vlm_3b.yaml` | 3B vision-language model (PaliGemma) |\n", "| `vla_diffusion_paligemma2.yaml` | VLA diffusion policy with PaliGemma2 backbone |" ] }, { "cell_type": "code", "execution_count": null, "id": "5r84eq1a9g", "metadata": {}, "outputs": [], "source": [ "%%writefile pusht_training_config.yaml\n", "# Training config for PushT diffusion policy with a small ViT backbone (no pretrained weights).\n", "save_path: tutorials/checkpoints/pusht\n", "# Uses 96x96 images directly (native PushT resolution), no augmentation.\n", "\n", "model:\n", " <<: !include vla_foundry/config_presets/models/diffusion_policy.yaml\n", " vision_language_backbone:\n", " type: vit_backbone\n", " img_size: 96\n", " hidden_dim: 64\n", " inter_dim: 256\n", " n_heads: 4\n", " n_layers: 4\n", " patch_size: 8\n", " # Minimal denoising head (this IS the diffusion head, not an LLM)\n", " transformer:\n", " type: transformer\n", " hidden_dim: 64\n", " n_heads: 4\n", " n_layers: 2\n", " is_causal: True\n", "\n", "data:\n", " <<: !include vla_foundry/config_presets/data/diffusion_policy.yaml\n", " processor: \"none\"\n", " image_size: 96\n", " dataset_manifest:\n", " - ./tutorials/data/preprocessed/pusht/shards/manifest.jsonl\n", " dataset_statistics:\n", " - ./tutorials/data/preprocessed/pusht/shards/stats.json\n", " dataset_modality:\n", " - robotics\n", " dataset_weighting:\n", " - 1.0\n", " camera_names:\n", " - observation.image\n", " image_indices:\n", " - -1\n", " - 0\n", " action_fields:\n", " - action\n", " proprioception_fields:\n", " - observation.state\n", " num_workers: 4\n", "\n", "distributed:\n", " fsdp: True\n", "\n", "hparams:\n", " <<: !include vla_foundry/config_presets/hparams/diffusion_policy.yaml\n", " per_gpu_batch_size: 32\n", " global_batch_size: 512" ] }, { "cell_type": "markdown", "id": "ebtvunyyj26", "metadata": {}, "source": [ "### Running Training Locally\n", "\n", "Launch training with `torchrun`. Key training arguments:\n", "\n", "| Argument | Description |\n", "|---|---|\n", "| `--config_path` | Path to training config YAML |\n", "| `--total_train_samples` | Total number of training samples to process |\n", "| `--num_checkpoints` | Number of checkpoints to save during training |\n", "| `--remote_sync` | S3 path to sync checkpoints to (optional) |\n", "| `--wandb False` | Disable Weights & Biases logging |\n", "| `--name` | Experiment name (used for logging) |\n", "\n", "For this tutorial, we run a small training job (5000 samples, 1 checkpoint) with wandb disabled.\n", "\n", "**Timing:** On a single GPU, training 5000 samples with this small model takes approximately 10-15 seconds. A production-scale run (100k+ samples with a larger model) will take significantly longer." ] }, { "cell_type": "code", "execution_count": null, "id": "mkl29rh9ilp", "metadata": {}, "outputs": [], "source": [ "!uv run torchrun --nproc_per_node=1 --nnodes=1 vla_foundry/main.py \\\n", " --config_path pusht_training_config.yaml \\\n", " --total_train_samples 5000 \\\n", " --num_checkpoints 1 \\\n", " --wandb False" ] }, { "cell_type": "markdown", "id": "njt032hu3d", "metadata": {}, "source": [ "### Additional Training Options\n", "\n", "**Previewing the resolved config** -- Before training, preview the fully resolved configuration with all `!include` directives expanded:\n", "```bash\n", "uv run torchrun --nproc_per_node=1 --nnodes=1 vla_foundry/main.py \\\n", " --config_path pusht_training_config.yaml \\\n", " --total_train_samples 100000 \\\n", " --num_checkpoints 5 \\\n", " --resolve_configs True \\\n", " --resolve_configs_path ./\n", "```\n", "\n", "**Syncing checkpoints to S3** -- Add `--remote_sync s3://my-bucket/model_checkpoints/pusht_diffusion_policy` to save checkpoints to S3 during training.\n", "\n", "**Finetuning from pretrained weights** -- Use `--model.resume_from_checkpoint ` and `--model.resume_weights_only True` to load only model weights without resuming optimizer state." ] }, { "cell_type": "markdown", "id": "c2jg1tq1lm", "metadata": {}, "source": [ "## 4. Verify Inference\n", "\n", "Now let's verify the trained model works by loading the checkpoint, feeding it a sample from the preprocessed data, and running a forward pass to generate action predictions." ] }, { "cell_type": "code", "execution_count": null, "id": "70rcndmbrs3", "metadata": {}, "outputs": [], "source": [ "import glob\n", "import os\n", "\n", "# Find the experiment directory (most recent one)\n", "experiment_dirs = sorted(glob.glob(\"tutorials/checkpoints/pusht/*/\"), key=os.path.getmtime)\n", "assert len(experiment_dirs) > 0, \"No experiment directories found. Did training complete?\"\n", "experiment_dir = experiment_dirs[-1].rstrip(\"/\")\n", "print(f\"Using experiment directory: {experiment_dir}\")\n", "\n", "# Find the checkpoint\n", "checkpoint_files = sorted(glob.glob(os.path.join(experiment_dir, \"checkpoints\", \"checkpoint_*.pt\")))\n", "assert len(checkpoint_files) > 0, \"No checkpoints found.\"\n", "checkpoint_path = checkpoint_files[-1]\n", "print(f\"Using checkpoint: {checkpoint_path}\")\n", "\n", "# Config path\n", "config_path = os.path.join(experiment_dir, \"config.yaml\")\n", "print(f\"Using config: {config_path}\")\n", "\n", "# Data shard\n", "shard_path = \"tutorials/data/preprocessed/pusht/shards/shard_000000.tar\"\n", "print(f\"Using data shard: {shard_path}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "i8i7uoqffd", "metadata": {}, "outputs": [], "source": [ "import io\n", "import json\n", "import tarfile\n", "\n", "import numpy as np\n", "import torch\n", "from PIL import Image\n", "\n", "from vla_foundry.data.processor.robotics_processor import RoboticsProcessor\n", "from vla_foundry.file_utils import load_model_checkpoint\n", "from vla_foundry.models import create_model\n", "from vla_foundry.params.train_experiment_params import load_experiment_params_from_yaml\n", "\n", "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "print(f\"Device: {DEVICE}\")\n", "\n", "# Load experiment config\n", "cfg = load_experiment_params_from_yaml(config_path, localize_params=True)\n", "print(f\"Model type: {cfg.model.type}\")\n", "print(f\"Action dim: {cfg.model.action_dim}\")\n", "print(f\"Proprioception dim: {cfg.model.proprioception_dim}\")\n", "\n", "# Create model and load checkpoint\n", "model = create_model(cfg.model, load_pretrained=False)\n", "start_ckpt, global_step, _ = load_model_checkpoint(model, checkpoint_path)\n", "print(f\"Checkpoint loaded: checkpoint_num={start_ckpt}, global_step={global_step}\")\n", "\n", "model.to(DEVICE)\n", "model.eval()\n", "print(f\"Model on {DEVICE}, eval mode\")" ] }, { "cell_type": "code", "execution_count": null, "id": "e6y1umeeu8u", "metadata": {}, "outputs": [], "source": [ "# Load RoboticsProcessor for normalization\n", "robotics_processor = RoboticsProcessor.from_pretrained(experiment_dir)\n", "print(\"RoboticsProcessor loaded\")\n", "print(f\" Normalizer enabled: {robotics_processor.normalizer is not None}\")\n", "\n", "# Extract one sample from the tar shard\n", "sample_files = {}\n", "first_sample_key = None\n", "\n", "with tarfile.open(shard_path, \"r\") as tar:\n", " for member in tar.getmembers():\n", " name = member.name\n", " sample_key = name.split(\".\", 1)[0]\n", " if first_sample_key is None:\n", " first_sample_key = sample_key\n", " if sample_key != first_sample_key:\n", " if len(sample_files) >= 4:\n", " break\n", " continue\n", " f = tar.extractfile(member)\n", " if f is not None:\n", " sample_files[name] = f.read()\n", "\n", "print(f\"\\nSample key: {first_sample_key}\")\n", "print(f\"Files in sample: {list(sample_files.keys())}\")\n", "\n", "# Parse the sample data\n", "images = {}\n", "lowdim_data = None\n", "metadata = None\n", "language_instructions = None\n", "\n", "for filename, data in sample_files.items():\n", " if filename.endswith(\".jpg\") or filename.endswith(\".png\"):\n", " img = Image.open(io.BytesIO(data)).convert(\"RGB\")\n", " images[filename] = img\n", " print(f\" Image: {filename} -> size={img.size}\")\n", " elif filename.endswith(\".npz\"):\n", " lowdim_data = dict(np.load(io.BytesIO(data)))\n", " print(f\" Lowdim fields: {list(lowdim_data.keys())}\")\n", " for k, v in lowdim_data.items():\n", " print(f\" {k}: shape={v.shape}, dtype={v.dtype}\")\n", " elif filename.endswith(\"metadata.json\"):\n", " metadata = json.loads(data.decode())\n", " print(f\" Metadata: {metadata}\")\n", " elif filename.endswith(\"language_instructions.json\"):\n", " language_instructions = json.loads(data.decode())\n", " print(f\" Language instructions: {language_instructions}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "w0pdynqmpp", "metadata": {}, "outputs": [], "source": [ "import torchvision.transforms as T\n", "\n", "# Build model input and run inference\n", "lowdim_past_timesteps = cfg.data.lowdim_past_timesteps\n", "lowdim_future_timesteps = cfg.data.lowdim_future_timesteps\n", "total_timesteps = lowdim_past_timesteps + 1 + lowdim_future_timesteps\n", "action_dim = cfg.data.action_dim\n", "image_names = cfg.data.image_names\n", "\n", "print(\"Data config:\")\n", "print(f\" past_timesteps: {lowdim_past_timesteps}\")\n", "print(f\" future_timesteps: {lowdim_future_timesteps}\")\n", "print(f\" total_timesteps: {total_timesteps}\")\n", "print(f\" action_dim: {action_dim}\")\n", "print(f\" image_names: {image_names}\")\n", "\n", "# Collect images in order and convert to tensors\n", "img_to_tensor = T.ToTensor()\n", "\n", "ordered_image_tensors = []\n", "for img_name in image_names:\n", " for filename, img in images.items():\n", " if img_name in filename:\n", " # Resize to model's expected size (96x96)\n", " img_resized = img.resize((96, 96))\n", " ordered_image_tensors.append(img_to_tensor(img_resized))\n", " break\n", "\n", "print(f\"\\nNumber of images for model: {len(ordered_image_tensors)}\")\n", "\n", "# Stack into [1, N, C, H, W]\n", "pixel_values = torch.stack(ordered_image_tensors).unsqueeze(0)\n", "# Dummy text inputs (ViT backbone ignores them)\n", "input_ids = torch.zeros(1, 1, dtype=torch.long)\n", "attention_mask = torch.ones(1, 1, dtype=torch.long)\n", "\n", "print(f\" pixel_values: {pixel_values.shape}\")\n", "\n", "# Build action and proprioception tensors\n", "raw_actions = torch.tensor(lowdim_data[\"action\"], dtype=torch.float32)\n", "if robotics_processor.normalizer is not None:\n", " actions_normalized = robotics_processor.normalizer.normalize_tensor(raw_actions.unsqueeze(0), \"action\")\n", "else:\n", " actions_normalized = raw_actions.unsqueeze(0)\n", "\n", "raw_proprio = torch.tensor(lowdim_data[\"observation.state\"], dtype=torch.float32)\n", "if robotics_processor.normalizer is not None:\n", " proprio_normalized = robotics_processor.normalizer.normalize_tensor(raw_proprio.unsqueeze(0), \"observation.state\")\n", "else:\n", " proprio_normalized = raw_proprio.unsqueeze(0)\n", "proprio_normalized = proprio_normalized[:, : lowdim_past_timesteps + 1]\n", "\n", "# Build past_mask\n", "past_mask = torch.zeros(1, total_timesteps, dtype=torch.bool)\n", "past_mask[:, : lowdim_past_timesteps + 1] = True\n", "\n", "# Move to device\n", "input_ids = input_ids.to(DEVICE)\n", "attention_mask = attention_mask.to(DEVICE)\n", "pixel_values = pixel_values.to(DEVICE, dtype=torch.float32)\n", "actions_normalized = actions_normalized.to(DEVICE, dtype=torch.float32)\n", "past_mask = past_mask.to(DEVICE)\n", "proprio_normalized = proprio_normalized.to(DEVICE, dtype=torch.float32)\n", "\n", "# Run inference\n", "print(\"\\nRunning generate_actions...\")\n", "NUM_FLOW_STEPS = 10\n", "\n", "with torch.no_grad():\n", " predicted_actions = model.generate_actions(\n", " input_ids=input_ids,\n", " pixel_values=pixel_values,\n", " actions=actions_normalized,\n", " attention_mask=attention_mask,\n", " num_inference_steps=NUM_FLOW_STEPS,\n", " past_mask=past_mask,\n", " proprioception=proprio_normalized,\n", " )\n", "\n", "print(f\"\\nPredicted actions shape: {predicted_actions.shape}\")\n", "print(f\"Predicted actions dtype: {predicted_actions.dtype}\")\n", "\n", "pred_cpu = predicted_actions.cpu().float().numpy()\n", "print(\"\\nPredicted actions stats (normalized):\")\n", "print(f\" Min: {pred_cpu.min():.6f}\")\n", "print(f\" Max: {pred_cpu.max():.6f}\")\n", "print(f\" Mean: {pred_cpu.mean():.6f}\")\n", "print(f\" Std: {pred_cpu.std():.6f}\")\n", "\n", "print(\"\\nFirst 5 predicted action timesteps (normalized):\")\n", "for t in range(min(5, pred_cpu.shape[1])):\n", " print(f\" t={t}: {pred_cpu[0, t]}\")\n", "\n", "# Denormalize to get actual action values\n", "if robotics_processor.normalizer is not None:\n", " denorm_actions = robotics_processor.normalizer.denormalize_tensor(predicted_actions.cpu(), \"action\").numpy()\n", " print(\"\\nDenormalized actions (first 5 timesteps):\")\n", " for t in range(min(5, denorm_actions.shape[1])):\n", " print(f\" t={t}: {denorm_actions[0, t]}\")\n", " print(f\"\\n Denormalized range: [{denorm_actions.min():.2f}, {denorm_actions.max():.2f}]\")\n", "\n", "print(\"\\nInference verification complete!\")" ] }, { "cell_type": "markdown", "id": "zf4zjrteaj9", "metadata": {}, "source": [ "## 5. Interactive PushT Evaluation\n", "\n", "Now let's close the loop: run the trained model inside the actual PushT simulation environment. This requires the `gym-pusht` package (a Gymnasium environment for the 2D Push-T task).\n", "\n", "The eval loop:\n", "1. Resets the environment and gets the initial observation (96\u00d796 image + 2D agent position)\n", "2. At each step, builds the model input from the live observation (two images, proprioception, text instruction)\n", "3. Runs `generate_actions` to get an action chunk (16 timesteps)\n", "4. Executes the **first predicted future action** in the environment\n", "5. Records frames for visualization\n", "\n", "> **Note:** The model was only trained for 5000 samples in this tutorial, so don't expect great performance. For a well-performing policy, train for 100k+ samples." ] }, { "cell_type": "code", "execution_count": null, "id": "2okicqm1hht", "metadata": {}, "outputs": [], "source": [ "import gym_pusht # noqa: F401 -- registers the gymnasium environment\n", "import gymnasium as gym\n", "import numpy as np\n", "import torch\n", "from PIL import Image\n", "from torchvision import transforms\n", "\n", "# --- Configuration ---\n", "MAX_STEPS = 300 # Maximum environment steps per episode\n", "NUM_FLOW_STEPS = 10 # Denoising steps for action generation\n", "ACTION_EXEC_HORIZON = 8 # How many actions to execute before re-planning\n", "IMG_SIZE = 96 # Native PushT resolution, matches training\n", "\n", "# Data config from training\n", "lowdim_past_timesteps = cfg.data.lowdim_past_timesteps # 1\n", "lowdim_future_timesteps = cfg.data.lowdim_future_timesteps # 14\n", "total_timesteps = lowdim_past_timesteps + 1 + lowdim_future_timesteps # 16\n", "anchor_idx = lowdim_past_timesteps # index of \"current\" timestep in the window\n", "\n", "# Pre-compute static tensors\n", "past_mask = torch.zeros(1, total_timesteps, dtype=torch.bool, device=DEVICE)\n", "past_mask[:, : lowdim_past_timesteps + 1] = True\n", "dummy_actions = torch.zeros(1, total_timesteps, cfg.data.action_dim, device=DEVICE)\n", "\n", "# Image transform: uint8 HWC -> float32 CHW [0, 1]\n", "img_transform = transforms.ToTensor() # handles both conversion and normalization to [0,1]\n", "\n", "\n", "def obs_to_pixel_values(prev_pixels: np.ndarray, curr_pixels: np.ndarray) -> torch.Tensor:\n", " \"\"\"Convert two observation images to model input [1, 2, C, H, W].\"\"\"\n", " t_prev = img_transform(prev_pixels) # [C, H, W]\n", " t_curr = img_transform(curr_pixels) # [C, H, W]\n", " return torch.stack([t_prev, t_curr]).unsqueeze(0) # [1, 2, C, H, W]\n", "\n", "\n", "# --- Create environment ---\n", "env = gym.make(\"gym_pusht/PushT-v0\", obs_type=\"pixels_agent_pos\", render_mode=\"rgb_array\")\n", "obs, info = env.reset(seed=42)\n", "\n", "prev_pixels = obs[\"pixels\"]\n", "curr_pixels = obs[\"pixels\"]\n", "prev_agent_pos = obs[\"agent_pos\"].copy()\n", "curr_agent_pos = obs[\"agent_pos\"].copy()\n", "\n", "frames = [env.render()]\n", "rewards = []\n", "action_queue = []\n", "\n", "print(f\"Environment created. Running eval for up to {MAX_STEPS} steps...\")\n", "print(f\" Image size: {obs['pixels'].shape} (native, no resizing)\")\n", "print(f\" Agent pos: {obs['agent_pos']}\")\n", "print()\n", "\n", "for step in range(MAX_STEPS):\n", " if len(action_queue) == 0:\n", " # Build proprioception: [t-1, t0] normalized\n", " proprio_raw = torch.tensor(\n", " np.stack([prev_agent_pos, curr_agent_pos]),\n", " dtype=torch.float32,\n", " ).unsqueeze(0) # [1, 2, 2]\n", " if robotics_processor.normalizer is not None:\n", " proprio_norm = robotics_processor.normalizer.normalize_tensor(proprio_raw, \"observation.state\")\n", " else:\n", " proprio_norm = proprio_raw\n", "\n", " # Build pixel values directly (no CLIP processor needed)\n", " pixel_values = obs_to_pixel_values(prev_pixels, curr_pixels).to(DEVICE, dtype=torch.float32)\n", "\n", " # Generate actions (no input_ids/attention_mask \u2014 ViT backbone ignores them)\n", " with torch.no_grad():\n", " predicted_actions = model.generate_actions(\n", " input_ids=torch.zeros(1, 1, dtype=torch.long, device=DEVICE),\n", " pixel_values=pixel_values,\n", " actions=dummy_actions,\n", " attention_mask=torch.ones(1, 1, dtype=torch.long, device=DEVICE),\n", " num_inference_steps=NUM_FLOW_STEPS,\n", " past_mask=past_mask,\n", " proprioception=proprio_norm.to(DEVICE, dtype=torch.float32),\n", " )\n", "\n", " # Denormalize and extract future actions\n", " if robotics_processor.normalizer is not None:\n", " denorm = robotics_processor.normalizer.denormalize_tensor(predicted_actions.cpu(), \"action\").numpy()[0]\n", " else:\n", " denorm = predicted_actions.cpu().numpy()[0]\n", "\n", " future_actions = denorm[anchor_idx : anchor_idx + ACTION_EXEC_HORIZON]\n", " action_queue = list(future_actions)\n", "\n", " # Execute next action\n", " action = action_queue.pop(0)\n", " action = np.clip(action, 0.0, 512.0).astype(np.float32)\n", "\n", " obs, reward, terminated, truncated, info = env.step(action)\n", " rewards.append(reward)\n", " frames.append(env.render())\n", "\n", " # Update observation history\n", " prev_pixels = curr_pixels\n", " prev_agent_pos = curr_agent_pos.copy()\n", " curr_pixels = obs[\"pixels\"]\n", " curr_agent_pos = obs[\"agent_pos\"].copy()\n", "\n", " if terminated or truncated:\n", " print(f\"Episode ended at step {step + 1}\")\n", " break\n", "\n", " if (step + 1) % 50 == 0:\n", " print(f\" Step {step + 1}/{MAX_STEPS} reward={reward:.4f} coverage={info.get('coverage', 0):.4f}\")\n", "\n", "env.close()\n", "\n", "print(f\"\\nEval complete: {len(rewards)} steps\")\n", "print(f\" Final reward: {rewards[-1]:.4f}\")\n", "print(f\" Max reward: {max(rewards):.4f}\")\n", "print(f\" Final coverage: {info.get('coverage', 0):.4f}\")\n", "print(f\" Success: {info.get('is_success', False)}\")\n", "print(f\" Frames captured: {len(frames)}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "0wc18ir4n7me", "metadata": {}, "outputs": [], "source": [ "# Render the rollout as an inline video in the notebook\n", "import base64\n", "\n", "import imageio.v3 as iio\n", "from IPython.display import HTML, display\n", "\n", "video_path = \"tutorials/data/pusht_eval_rollout.mp4\"\n", "iio.imwrite(video_path, np.stack(frames), fps=10, codec=\"h264\", plugin=\"pyav\")\n", "print(f\"Video saved to {video_path}\")\n", "\n", "# Display inline\n", "with open(video_path, \"rb\") as f:\n", " video_data = f.read()\n", "b64 = base64.b64encode(video_data).decode()\n", "display(\n", " HTML(f\"\"\"\n", "\n", "\"\"\")\n", ")" ] }, { "cell_type": "markdown", "id": "5u7a12vgsig", "metadata": {}, "source": [ "## 6. Deployment\n", "\n", "For simulation evaluation of a trained VLA checkpoint, use [`lbm_eval`](https://github.com/ToyotaResearchInstitute/lbm_eval). See `vla_foundry/inference/scripts/README.md` for instructions on launching the policy server.\n", "\n", "For deploying a trained model as a policy server (e.g., for real robot control), see `vla_foundry/inference/robotics/inference_policy.py`.\n", "\n", "## Tips and Common Issues\n", "\n", "- **Camera name mismatch**: The `--camera_names` passed during preprocessing must match the column names in the LeRobot parquet files\n", "- **Action key names vary**: Some LeRobot datasets use `action` (singular) while others use `actions` (plural)\n", "- **Image format**: The converter handles both inline image bytes in parquet files and video-based storage (MP4 files). It auto-detects which format is used\n", "- **Memory usage**: For very large datasets, use Ray parallelization to distribute processing across cores\n", "- **Statistics file**: The `stats.json` file is critical for normalization during training. If preprocessing is interrupted before statistics are saved, re-run\n", "\n", "**See also**:\n", "- `tutorials/diffusion_policy.md` -- Diffusion policy (Spartan) data workflow\n", "- `tutorials/adding_new_models.md` -- Adding custom models\n", "- `vla_foundry/data/preprocessing/README.md` -- Ray cluster setup for preprocessing\n", "- `vla_foundry/inference/scripts/README.md` -- inference policy server scripts" ] } ], "metadata": { "kernelspec": { "display_name": "Python (vla_foundry)", "language": "python", "name": "vla_foundry" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.12" } }, "nbformat": 4, "nbformat_minor": 5 }