{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f9a4f8a5",
   "metadata": {},
   "source": [
    "# VLA Foundry: Training an LLM, VLM, and VLA from Scratch\n",
    "\n",
    "This notebook walks through the full three-stage pipeline for training a Vision-Language-Action (VLA) model from scratch using **small model configs** (100M backbone, tiny ViT) and local test data:\n",
    "\n",
    "1. **LLM** -- Train a language model on text data\n",
    "2. **VLM** -- Add vision capabilities with image-caption training\n",
    "3. **VLA** -- Add action prediction with robotics data\n",
    "\n",
    "Each stage builds on the previous one: language understanding transfers to vision, and vision-language understanding transfers to action prediction. The same 100M transformer backbone is used throughout. In the default settings that we set, none of the stages are trained enought to get good results, the user is welcome to change the settings to train longer.\n",
    "\n",
    "| Stage | Model Config | Backbone | Additional Components |\n",
    "|-------|-------------|----------|----------------------|\n",
    "| LLM | `transformer_100m.yaml` | 12 layers, 8 heads, hidden_dim=512 | -- |\n",
    "| VLM | `vlm_100m.yaml` | 12 layers, 8 heads, hidden_dim=512 | + tiny ViT (4 layers, hidden_dim=64) |\n",
    "| VLA | `vla_diffusion_100m.yaml` | 12 layers, 8 heads, hidden_dim=512 | + tiny ViT + diffusion head |\n",
    "\n",
    "**Requirements:** 1 GPU, `uv` package manager. Run this notebook from the repo root."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "added_cell",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "- GPU with ≥16 GB VRAM\n",
    "- The **`Python (vla_foundry)`** Jupyter kernel — install it once from the repo root:\n",
    "  ```bash\n",
    "  bash tutorials/install_kernel.sh\n",
    "  ```\n",
    "  Then restart your notebook kernel and select **Python (vla_foundry)**.\n",
    "- Run from the repository root (the first code cell handles this automatically)\n",
    "\n",
    "This notebook is **standalone** — no other tutorial needs to be completed first. It downloads all required data automatically.\n",
    "\n",
    "**Storage:** ~650 MB downloaded from HuggingFace, ~550 MB of preprocessed output in `tutorials/data/`. Ensure you have at least 1.5 GB free.\n",
    "\n",
    "| Stage | HF download | Output |\n",
    "|---|---|---|\n",
    "| Stage 1: DCLM text | ~135 MB | ~360 MB |\n",
    "| Stage 2: PixelProse images | ~25 MB + URL downloads | ~4 MB |\n",
    "| Stage 3: DROID robotics | ~443 MB | ~187 MB |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7323de58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Notebook must run from the repo root, not from tutorials/\n",
    "if os.path.basename(os.getcwd()) == \"tutorials\":\n",
    "    os.chdir(\"..\")\n",
    "print(f\"Working directory: {os.getcwd()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "519f9f9d",
   "metadata": {},
   "source": [
    "---\n",
    "## Stage 1: LLM Training\n",
    "\n",
    "We train a 100M parameter transformer on a minimal 1k-document sample from [DCLM-baseline](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0). The sample is converted into a local untokenized WebDataset so the training command matches the rest of the stack. We override `vocab_size` to 49280 to match the SmolVLM2 tokenizer used in the VLM stage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d0ab1b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tutorials.tutorial_utils import download_dclm_text\n",
    "\n",
    "dclm_manifest = download_dclm_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05424c68",
   "metadata": {},
   "outputs": [],
   "source": [
    "!uv run torchrun --nproc_per_node=1 --nnodes=1 --master_port=0 vla_foundry/main.py \\\n",
    "  --model.type transformer \\\n",
    "  --model \"include vla_foundry/config_presets/models/transformer_100m.yaml\" \\\n",
    "  --model.vocab_size 49280 \\\n",
    "  --model.cast_output_to_float32 True \\\n",
    "  --distributed.fsdp False \\\n",
    "  --data.type text_untokenized \\\n",
    "  --data.dataset_manifest '[\"tutorials/data/dclm_minimal/manifest.jsonl\"]' \\\n",
    "  --data.dataset_modality '[\"text_untokenized\"]' \\\n",
    "  --data.dataset_weighting '[1.0]' \\\n",
    "  --data.tokenizer HuggingFaceTB/SmolVLM2-256M-Video-Instruct \\\n",
    "  --data.seq_len 256 \\\n",
    "  --data.num_workers 2 \\\n",
    "  --data.allow_multiple_epochs True \\\n",
    "  --total_train_samples 10000 \\\n",
    "  --num_checkpoints 2 \\\n",
    "  --hparams.per_gpu_batch_size 8 \\\n",
    "  --hparams.global_batch_size 16 \\\n",
    "  --hparams.lr 0.0001 \\\n",
    "  --hparams.warmup 2 \\\n",
    "  --hparams.torchcompile True \\\n",
    "  --log_every_n_steps 1 \\\n",
    "  --save_path ./tutorials/checkpoints \\\n",
    "  --wandb False \\\n",
    "  --db_logging False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "533701c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "\n",
    "from tutorials.tutorial_utils import plot_training_loss\n",
    "\n",
    "llm_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*transformer*\"))[-1]\n",
    "plot_training_loss(llm_exp_dir, \"Stage 1: LLM Training Loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef73cdf5",
   "metadata": {},
   "source": [
    "### Vibe check: Text generation\n",
    "\n",
    "Load the trained LLM and generate a short text completion. With only 100M parameters and a tiny dataset, the output will not be coherent -- the goal is to confirm the model produces valid tokens and that loading works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e448db60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "\n",
    "import torch\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "from vla_foundry.file_utils import load_model_checkpoint\n",
    "from vla_foundry.models import create_model\n",
    "from vla_foundry.params.model_params import ModelParams\n",
    "from vla_foundry.params.train_experiment_params import load_params_from_yaml\n",
    "\n",
    "# Find the latest LLM experiment directory\n",
    "llm_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*transformer*\"))[0]\n",
    "print(f\"Loading from: {llm_exp_dir}\")\n",
    "\n",
    "# Load model config and create model\n",
    "model_params = load_params_from_yaml(ModelParams, f\"{llm_exp_dir}/config_model.yaml\")\n",
    "model = create_model(model_params)\n",
    "\n",
    "# Load the latest checkpoint\n",
    "ckpt_files = sorted(glob.glob(f\"{llm_exp_dir}/checkpoints/checkpoint_*.pt\"))\n",
    "print(f\"Loading checkpoint: {ckpt_files[-1]}\")\n",
    "load_model_checkpoint(model, ckpt_files[-1])\n",
    "model.eval()\n",
    "\n",
    "# Tokenize and generate\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"HuggingFaceTB/SmolVLM2-256M-Video-Instruct\")\n",
    "inputs = tokenizer([\"The future of robotics is\"], return_tensors=\"pt\", padding=True)\n",
    "with torch.no_grad():\n",
    "    output = model.generate(inputs[\"input_ids\"], inputs[\"attention_mask\"], max_new_tokens=30)\n",
    "print(\"Generated text:\")\n",
    "print(tokenizer.batch_decode(output, skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14979e58",
   "metadata": {},
   "source": [
    "---\n",
    "## Stage 2: VLM Training\n",
    "\n",
    "We add vision capabilities by combining the 100M transformer backbone with a tiny ViT encoder using a minimal sample from [PixelProse](https://huggingface.co/datasets/tomg-group-umd/pixelprose). PixelProse provides dense, high-quality captions generated by Gemini 1.0 Pro Vision for images drawn from CC12M and LAION — a step up in caption quality over shorter crowd-sourced alternatives. We download one parquet shard, fetch 64 images from their URLs, write a local manifest, and run a small training step. The `vlm_100m.yaml` config includes `transformer_100m.yaml` and `vit_tiny.yaml`; `vocab_size` is resolved from the `simple_vlm` processor's tokenizer (SmolVLM2, 49280 tokens) at runtime.\n",
    "\n",
    "**Note:** LLM→VLM weight transfer requires prefix remapping (not yet automated), so this stage trains from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a37810",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tutorials.tutorial_utils import download_pixelprose_images\n",
    "\n",
    "pixelprose_manifest = download_pixelprose_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7297fdca",
   "metadata": {},
   "outputs": [],
   "source": [
    "!uv run torchrun --nproc_per_node=1 --nnodes=1 --master_port=0 vla_foundry/main.py \\\n",
    "  --model \"include vla_foundry/config_presets/models/vlm_100m.yaml\" \\\n",
    "  --distributed.fsdp False \\\n",
    "  --data.type image_caption \\\n",
    "  --data.processor simple_vlm \\\n",
    "  --data.dataset_manifest '[\"tutorials/data/pixelprose_minimal/manifest.jsonl\"]' \\\n",
    "  --data.dataset_modality '[\"image_caption\"]' \\\n",
    "  --data.dataset_weighting '[1.0]' \\\n",
    "  --data.seq_len 512 \\\n",
    "  --data.img_num_tokens 256 \\\n",
    "  --data.num_workers 1 \\\n",
    "  --data.allow_multiple_epochs True \\\n",
    "  --total_train_samples 3000 \\\n",
    "  --num_checkpoints 1 \\\n",
    "  --hparams.per_gpu_batch_size 16 \\\n",
    "  --hparams.global_batch_size 32 \\\n",
    "  --hparams.lr 0.001 \\\n",
    "  --hparams.warmup 2 \\\n",
    "  --hparams.torchcompile True \\\n",
    "  --log_every_n_steps 1 \\\n",
    "  --save_path ./tutorials/checkpoints \\\n",
    "  --wandb False \\\n",
    "  --db_logging False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2dabe3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "\n",
    "from tutorials.tutorial_utils import plot_training_loss\n",
    "\n",
    "vlm_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*vlm*\"))[-1]\n",
    "plot_training_loss(vlm_exp_dir, \"Stage 2: VLM Training Loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b162e6",
   "metadata": {},
   "source": [
    "### Vibe check: Image captioning\n",
    "\n",
    "Load the trained VLM and generate a caption for an image from the training data. With a tiny ViT and 100M backbone trained on minimal data, the caption will not be meaningful -- the goal is to confirm the model loads, processes images, and generates tokens.\n",
    "\n",
    "> **Expected output:** You will likely see random or repeated letters and symbols (e.g. `aaa...`, `###`, partial words). This is normal at this scale. A valid run means the model completes without error and produces some token output — coherence is not expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d42af685",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "import io\n",
    "import tarfile\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "from vla_foundry.file_utils import load_model_checkpoint\n",
    "from vla_foundry.models import create_model\n",
    "from vla_foundry.models.vision_language_backbones.simple_vlm_processor import SimpleVLMProcessor\n",
    "from vla_foundry.params.model_params import ModelParams\n",
    "from vla_foundry.params.train_experiment_params import load_params_from_yaml\n",
    "\n",
    "# Find the latest VLM experiment directory\n",
    "vlm_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*vlm*\"))[-1]\n",
    "print(f\"Loading from: {vlm_exp_dir}\")\n",
    "\n",
    "# Load model config and create model\n",
    "model_params = load_params_from_yaml(ModelParams, f\"{vlm_exp_dir}/config_model.yaml\")\n",
    "model = create_model(model_params)\n",
    "\n",
    "# Load the latest checkpoint\n",
    "ckpt_files = sorted(glob.glob(f\"{vlm_exp_dir}/checkpoints/checkpoint_*.pt\"))\n",
    "print(f\"Loading checkpoint: {ckpt_files[-1]}\")\n",
    "load_model_checkpoint(model, ckpt_files[-1])\n",
    "model.eval()\n",
    "\n",
    "# Load a real image from the PixelProse training data\n",
    "img_idx = 3\n",
    "tar_path = glob.glob(\"tutorials/data/pixelprose_minimal/shard-*.tar\")[0]\n",
    "with tarfile.open(tar_path) as tar:\n",
    "    jpgs = [m for m in tar.getmembers() if m.name.endswith(\".jpg\")]\n",
    "    test_image = Image.open(io.BytesIO(tar.extractfile(jpgs[img_idx]).read())).convert(\"RGB\")\n",
    "    gt_caption = tar.extractfile(jpgs[img_idx].name.replace(\".jpg\", \".txt\")).read().decode(\"utf-8\").strip()\n",
    "# Generate caption\n",
    "processor = SimpleVLMProcessor()\n",
    "processor.image_seq_length = 256  # must match data.img_num_tokens\n",
    "prompt = \"<image>\"\n",
    "inputs = processor([test_image], [prompt], return_tensors=\"pt\", max_length=1024)\n",
    "\n",
    "with torch.no_grad():\n",
    "    output = model.generate(\n",
    "        input_ids=inputs[\"input_ids\"],\n",
    "        pixel_values=inputs[\"pixel_values\"],\n",
    "        attention_mask=inputs[\"attention_mask\"],\n",
    "        max_new_tokens=20,\n",
    "    )\n",
    "generated = processor.tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "\n",
    "# Display image with captions\n",
    "fig, ax = plt.subplots(figsize=(4, 4))\n",
    "ax.imshow(test_image)\n",
    "ax.axis(\"off\")\n",
    "ax.set_title(\"Input image\", fontsize=11)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"Ground truth: {gt_caption}\")\n",
    "print(f\"Generated:    {generated}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "685079e5",
   "metadata": {},
   "source": [
    "---\n",
    "## Stage 3: VLA Training\n",
    "\n",
    "The final stage adds action prediction on a tiny local subset of [lerobot/droid_100](https://huggingface.co/datasets/lerobot/droid_100). We download a few metadata files plus the three DROID camera videos, build a compatibility copy for the repo's LeRobot preprocessor, convert that into local tar shards, and train a short VLA run. The VLM checkpoint from Stage 2 is loaded via `vision_language_backbone.resume_from_checkpoint`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d6164b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tutorials.tutorial_utils import download_droid_robotics\n",
    "\n",
    "droid_compat_root = download_droid_robotics()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dce3eda5",
   "metadata": {},
   "source": [
    "**Data preprocessing:** The cell below converts the downloaded DROID episodes into WebDataset tar shards that VLA Foundry's data pipeline expects. Preprocessed data is written to `tutorials/data/droid_100_minimal/preprocessed/`. You can inspect the output manifest at `tutorials/data/droid_100_minimal/preprocessed/shards/manifest.jsonl`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462262bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "DROID_CAMERAS = [\n",
    "    \"observation.images.exterior_image_1_left\",\n",
    "    \"observation.images.exterior_image_2_left\",\n",
    "    \"observation.images.wrist_image_left\",\n",
    "]\n",
    "COMPAT_ROOT = Path(\"tutorials/data/droid_100_minimal/lerobot_compat\").resolve()\n",
    "PREPROC_ROOT = Path(\"tutorials/data/droid_100_minimal/preprocessed\").resolve()\n",
    "PREPROC_MANIFEST = PREPROC_ROOT / \"shards/manifest.jsonl\"\n",
    "PREPROC_STATS = PREPROC_ROOT / \"shards/stats.json\"\n",
    "\n",
    "vlm_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*vlm*\"))[-1]\n",
    "vlm_ckpt = sorted(glob.glob(f\"{vlm_exp_dir}/checkpoints/checkpoint_*.pt\"))[-1]\n",
    "\n",
    "if PREPROC_MANIFEST.exists() and PREPROC_STATS.exists():\n",
    "    print(f\"Using existing preprocessed data at {PREPROC_MANIFEST}.\")\n",
    "else:\n",
    "    preprocess_cmd = [\n",
    "        \"uv\",\n",
    "        \"run\",\n",
    "        \"--group\",\n",
    "        \"preprocessing\",\n",
    "        \"python\",\n",
    "        \"vla_foundry/data/preprocessing/preprocess_robotics_to_tar.py\",\n",
    "        \"--type\",\n",
    "        \"lerobot\",\n",
    "        \"--source_episodes\",\n",
    "        f\"['{COMPAT_ROOT}']\",\n",
    "        \"--output_dir\",\n",
    "        str(PREPROC_ROOT),\n",
    "        \"--camera_names\",\n",
    "        str(DROID_CAMERAS),\n",
    "        \"--observation_keys\",\n",
    "        \"['observation.state']\",\n",
    "        \"--action_keys\",\n",
    "        \"['action']\",\n",
    "        \"--samples_per_shard\",\n",
    "        \"32\",\n",
    "        \"--config_path\",\n",
    "        \"tutorials/tutorial_utils/robotics_preprocessing_params_tutorial.yaml\",\n",
    "        \"--max_episodes_to_process\",\n",
    "        \"2\",\n",
    "        \"--db_logging\",\n",
    "        \"False\",\n",
    "    ]\n",
    "    print(\"Running preprocessing:\")\n",
    "    print(\" \".join(preprocess_cmd))\n",
    "    subprocess.run(preprocess_cmd, check=True)\n",
    "\n",
    "print(f\"VLM checkpoint for backbone: {vlm_ckpt}\")\n",
    "print(f\"DROID manifest: {PREPROC_MANIFEST}\")\n",
    "print(f\"DROID stats:    {PREPROC_STATS}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ccc55ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "vlm_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*vlm*\"))[-1]\n",
    "vlm_ckpt = sorted(glob.glob(f\"{vlm_exp_dir}/checkpoints/checkpoint_*.pt\"))[-1]\n",
    "droid_manifest = str(Path(\"tutorials/data/droid_100_minimal/preprocessed/shards/manifest.jsonl\").resolve())\n",
    "droid_stats = str(Path(\"tutorials/data/droid_100_minimal/preprocessed/shards/stats.json\").resolve())\n",
    "\n",
    "# Camera names use the short form (without \"observation.images.\" prefix)\n",
    "# because the training pipeline strips prefixes via key.split(\".\")[-2]\n",
    "cameras = '[\"exterior_image_1_left\",\"exterior_image_2_left\",\"wrist_image_left\"]'\n",
    "\n",
    "# Stage 3 is fully self-contained: the VLA backbone loads the Stage-2 VLM checkpoint,\n",
    "# so data shape (image_size=224, img_num_tokens=256, processor=simple_vlm) must match\n",
    "# what Stage 2 was trained with. The hparams include provides MSE loss for diffusion.\n",
    "cmd = [\n",
    "    \"uv\",\n",
    "    \"run\",\n",
    "    \"torchrun\",\n",
    "    \"--nproc_per_node=1\",\n",
    "    \"--nnodes=1\",\n",
    "    \"--master_port=0\",\n",
    "    \"vla_foundry/main.py\",\n",
    "    \"--model\",\n",
    "    \"include vla_foundry/config_presets/models/vla_diffusion_100m.yaml\",\n",
    "    \"--model.vision_language_backbone.resume_from_checkpoint\",\n",
    "    vlm_ckpt,\n",
    "    \"--distributed.fsdp\",\n",
    "    \"False\",\n",
    "    \"--data.type\",\n",
    "    \"robotics\",\n",
    "    \"--data.processor\",\n",
    "    \"simple_vlm\",\n",
    "    \"--data.image_size\",\n",
    "    \"224\",\n",
    "    \"--data.img_num_tokens\",\n",
    "    \"256\",\n",
    "    \"--data.seq_len\",\n",
    "    \"2048\",\n",
    "    \"--data.dataset_manifest\",\n",
    "    f'[\"{droid_manifest}\"]',\n",
    "    \"--data.dataset_statistics\",\n",
    "    f'[\"{droid_stats}\"]',\n",
    "    \"--data.dataset_modality\",\n",
    "    '[\"robotics\"]',\n",
    "    \"--data.dataset_weighting\",\n",
    "    \"[1.0]\",\n",
    "    \"--data.camera_names\",\n",
    "    cameras,\n",
    "    \"--data.action_fields\",\n",
    "    '[\"action\"]',\n",
    "    \"--data.proprioception_fields\",\n",
    "    '[\"observation.state\"]',\n",
    "    \"--data.language_instruction_types\",\n",
    "    '[\"original\"]',\n",
    "    \"--data.pose_groups\",\n",
    "    \"[]\",\n",
    "    \"--data.intrinsics_fields\",\n",
    "    \"[]\",\n",
    "    \"--data.extrinsics_fields\",\n",
    "    \"[]\",\n",
    "    \"--data.lowdim_past_timesteps\",\n",
    "    \"2\",\n",
    "    \"--data.lowdim_future_timesteps\",\n",
    "    \"14\",\n",
    "    \"--data.allow_multiple_epochs\",\n",
    "    \"True\",\n",
    "    \"--data.num_workers\",\n",
    "    \"0\",\n",
    "    \"--hparams\",\n",
    "    \"include vla_foundry/config_presets/hparams/diffusion_policy.yaml\",\n",
    "    \"--hparams.per_gpu_batch_size\",\n",
    "    \"1\",\n",
    "    \"--hparams.global_batch_size\",\n",
    "    \"1\",\n",
    "    \"--hparams.warmup\",\n",
    "    \"2\",\n",
    "    \"--total_train_samples\",\n",
    "    \"128\",\n",
    "    \"--num_checkpoints\",\n",
    "    \"1\",\n",
    "    \"--max_checkpoint_limit\",\n",
    "    \"2\",\n",
    "    \"--save_path\",\n",
    "    \"./tutorials/checkpoints\",\n",
    "    \"--wandb\",\n",
    "    \"False\",\n",
    "    \"--db_logging\",\n",
    "    \"False\",\n",
    "]\n",
    "print(\"Running VLA training...\\n\")\n",
    "subprocess.run(cmd, check=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f24011a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "\n",
    "from tutorials.tutorial_utils import plot_training_loss\n",
    "\n",
    "vla_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*diffusion*\"))[-1]\n",
    "plot_training_loss(vla_exp_dir, \"Stage 3: VLA Training Loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f7d60bc",
   "metadata": {},
   "source": [
    "### Vibe check: VLA action prediction\n",
    "\n",
    "Load the trained VLA, feed it a real robotics sample from the training data, run `generate_actions()` to produce predicted action trajectories, and compare them against the ground truth. We also display camera images from the sample to show what the model sees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0569b23f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "import io\n",
    "import json\n",
    "import tarfile\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "from vla_foundry.data.pipelines.robotics import extract_robotics_fields\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.model_params import ModelParams\n",
    "from vla_foundry.params.train_experiment_params import load_params_from_yaml\n",
    "\n",
    "# ---------- 1. Load VLA model ----------\n",
    "vla_exp_dir = sorted(glob.glob(\"tutorials/checkpoints/*diffusion*\"))[-1]\n",
    "model_params = load_params_from_yaml(ModelParams, f\"{vla_exp_dir}/config_model.yaml\")\n",
    "model = create_model(model_params)\n",
    "ckpt = sorted(glob.glob(f\"{vla_exp_dir}/checkpoints/checkpoint_*.pt\"))[-1]\n",
    "load_model_checkpoint(model, ckpt)\n",
    "model.eval().cuda()\n",
    "print(f\"Loaded VLA from {vla_exp_dir}\")\n",
    "\n",
    "# ---------- 2. Load processor (with normalizer) ----------\n",
    "processor = RoboticsProcessor.from_pretrained(vla_exp_dir)\n",
    "data_params = processor.data_params\n",
    "\n",
    "# ---------- 3. Read one sample from a preprocessed tar ----------\n",
    "tar_path = glob.glob(\"tutorials/data/droid_100_minimal/preprocessed/shards/shard_*.tar\")[0]\n",
    "raw_sample = {}\n",
    "with tarfile.open(tar_path, \"r\") as tar:\n",
    "    members = tar.getmembers()\n",
    "    first_id = members[0].name.split(\".\")[0]\n",
    "    for m in members:\n",
    "        if not m.name.startswith(first_id):\n",
    "            break\n",
    "        suffix = m.name[len(first_id) + 1 :]\n",
    "        f = tar.extractfile(m)\n",
    "        if suffix.endswith(\".jpg\"):\n",
    "            raw_sample[suffix] = torch.from_numpy(np.array(Image.open(io.BytesIO(f.read())).convert(\"RGB\"))).permute(\n",
    "                2, 0, 1\n",
    "            )\n",
    "        elif suffix.endswith(\".npz\"):\n",
    "            raw_sample[suffix] = dict(np.load(io.BytesIO(f.read())))\n",
    "        elif suffix.endswith(\".json\"):\n",
    "            raw_sample[suffix] = json.load(f)\n",
    "\n",
    "# ---------- 4. Build a batch ----------\n",
    "fields = extract_robotics_fields(\n",
    "    raw_sample,\n",
    "    language_instruction_types=list(data_params.language_instruction_types),\n",
    "    action_fields=list(data_params.action_fields),\n",
    "    proprioception_fields=list(data_params.proprioception_fields),\n",
    "    lowdim_past_timesteps=data_params.lowdim_past_timesteps,\n",
    "    lowdim_future_timesteps=data_params.lowdim_future_timesteps,\n",
    ")\n",
    "batch_of_one = {k: [v] for k, v in fields.items()}\n",
    "batch = processor.process_inputs(batch_of_one, image_names=list(data_params.image_names))\n",
    "batch = processor.add_action_and_proprioception_fields(\n",
    "    batch,\n",
    "    action_fields=list(data_params.action_fields),\n",
    "    proprioception_fields=list(data_params.proprioception_fields),\n",
    ")\n",
    "\n",
    "# ---------- 5. Run generate_actions ----------\n",
    "actions = batch[\"actions\"].cuda()\n",
    "past_mask = torch.as_tensor(np.stack(fields[\"past_mask\"][None]), dtype=torch.bool).cuda()\n",
    "\n",
    "with torch.no_grad():\n",
    "    predicted = model.generate_actions(\n",
    "        input_ids=batch[\"input_ids\"].cuda(),\n",
    "        pixel_values=batch[\"pixel_values\"].cuda(),\n",
    "        actions=actions,\n",
    "        attention_mask=batch[\"attention_mask\"].cuda().bool(),\n",
    "        attention_mask_images=batch[\"attention_mask_images\"].cuda().bool(),\n",
    "        past_mask=past_mask,\n",
    "        proprioception=batch[\"proprioception\"].cuda() if \"proprioception\" in batch else None,\n",
    "        num_inference_steps=10,\n",
    "    )\n",
    "\n",
    "gt = actions.cpu().numpy()[0]\n",
    "pred = predicted.cpu().numpy()[0]\n",
    "past = past_mask.cpu().numpy()[0]\n",
    "T, D = gt.shape\n",
    "t = np.arange(T)\n",
    "future_start = int(past.sum())\n",
    "\n",
    "# ---------- 6. Plot: cameras (top) + x/y/z position (bottom, 3 plots) ----------\n",
    "cam_keys = sorted([k for k in raw_sample if k.endswith(\"_t0.jpg\")])[:3]\n",
    "axis_names = [\"x\", \"y\", \"z\"]\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, figsize=(15, 8))\n",
    "\n",
    "# Top row: camera images\n",
    "for i in range(3):\n",
    "    if i < len(cam_keys):\n",
    "        img = raw_sample[cam_keys[i]]\n",
    "        if isinstance(img, torch.Tensor):\n",
    "            img = img.permute(1, 2, 0).numpy()\n",
    "        axes[0, i].imshow(img)\n",
    "        cam_name = cam_keys[i].replace(\"observation.images.\", \"\").replace(\"_t0.jpg\", \"\")\n",
    "        axes[0, i].set_title(cam_name, fontsize=10)\n",
    "    axes[0, i].axis(\"off\")\n",
    "\n",
    "# Bottom row: one plot per xyz dimension\n",
    "for i in range(3):\n",
    "    ax = axes[1, i]\n",
    "    if i < D:\n",
    "        ax.plot(t, gt[:, i], \"-\", color=\"tab:blue\", linewidth=2, label=\"Ground truth\")\n",
    "        ax.plot(t, pred[:, i], \"--\", color=\"tab:red\", linewidth=2, label=\"Predicted\")\n",
    "        ax.axvline(future_start - 1, color=\"gray\", ls=\":\", alpha=0.5, label=\"Present timestep\")\n",
    "        ax.fill_betweenx(\n",
    "            ax.get_ylim() if ax.get_ylim()[0] != ax.get_ylim()[1] else [-1, 1],\n",
    "            0,\n",
    "            future_start - 1,\n",
    "            alpha=0.05,\n",
    "            color=\"gray\",\n",
    "        )\n",
    "        ax.set_title(f\"Position {axis_names[i]}\", fontsize=11)\n",
    "        ax.set_xlabel(\"Timestep\")\n",
    "        if i == 0:\n",
    "            ax.set_ylabel(\"Normalized value\")\n",
    "        ax.legend(fontsize=8)\n",
    "    else:\n",
    "        ax.axis(\"off\")\n",
    "\n",
    "plt.suptitle(\"VLA Vibe Check: GT vs Predicted action trajectories\", fontsize=13)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "print(\n",
    "    f\"Actions: {D} dims, {T} timesteps ({future_start} past (past actions + past observations) \"\n",
    "    \"+ 1 present (current observation, no action) \"\n",
    "    f\"+ {T - future_start - 1} future (no observations, no actions))\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61473236",
   "metadata": {},
   "source": [
    "---\n",
    "## Evaluation References\n",
    "\n",
    "For rigorous benchmarking beyond these quick vibe checks:\n",
    "\n",
    "- **LLM evaluation:** [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) -- integrates with VLA Foundry checkpoints for standard LLM benchmarks\n",
    "- **VLM evaluation:** [VLMEvalKit](https://github.com/open-compass/VLMEvalKit) -- supports VLA Foundry model loading for vision-language benchmarks\n",
    "- **VLA evaluation:** [lbm_eval](https://github.com/ToyotaResearchInstitute/lbm_eval) -- simulation evaluation framework. See `vla_foundry/inference/scripts/README.md` for instructions on launching the policy server."
   ]
  }
 ],
 "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
}