{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a79f1dae",
   "metadata": {},
   "source": [
    "# 3 - Record, train, and deploy in one loop\n",
    "\n",
    "The full data loop in one notebook: **record** a LeRobotDataset, **train** an [ACT](https://tonyzhaozh.github.io/aloha/) policy on it, **export** the checkpoint, and **load** it back as a runnable policy. It runs on CPU - small dataset, two training steps - so you can see the whole loop close on a laptop with no GPU.\n",
    "\n",
    "**Requirements:** the recording and loading steps need only `strands-robots[sim-mujoco,lerobot]`, but **step 2 also needs `lerobot[training]`, on CPU as well as GPU**:\n",
    "\n",
    "```bash\n",
    "pip install \"lerobot[training]\"\n",
    "```\n",
    "\n",
    "LeRobot's `train()` calls `require_package(\"accelerate\", extra=\"training\")` before it branches on device, so \"no GPU\" does not mean \"no extra\". Step 2 checks `result.status` and re-raises `result.message`, which carries LeRobot's own install remedy - `train()` reports failure in its result rather than raising, so an unchecked call would hand step 3 a `checkpoint_dir` of `None`.\n",
    "\n",
    "For a real policy you would raise the step count and run on a GPU, but the code path is identical. Part of the [examples/notebooks](./README.md) series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65844717",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-29T14:28:42.601260Z",
     "iopub.status.busy": "2026-06-29T14:28:42.601148Z",
     "iopub.status.idle": "2026-06-29T14:28:42.636012Z",
     "shell.execute_reply": "2026-06-29T14:28:42.635371Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# macOS uses the 'cgl' offscreen GL backend; Linux headless uses 'egl'.\n",
    "os.environ.setdefault(\"MUJOCO_GL\", \"cgl\" if sys.platform == \"darwin\" else \"egl\")\n",
    "os.environ.setdefault(\"STRANDS_TRUST_REMOTE_CODE\", \"1\")  # ACT loads through LeRobot\n",
    "import shutil\n",
    "\n",
    "ROOT, OUT = \"/tmp/nb3_dataset\", \"/tmp/nb3_ft\"\n",
    "# LeRobot's trainer refuses to write into an existing output_dir, so clear it\n",
    "# first to keep this notebook re-runnable.\n",
    "shutil.rmtree(OUT, ignore_errors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6777c453",
   "metadata": {},
   "source": [
    "## 1. Record a small dataset\n",
    "\n",
    "Same recording flow as [notebook 2](./02_record_and_stream.ipynb), kept short for a fast CPU run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c87459",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-29T14:28:42.638007Z",
     "iopub.status.busy": "2026-06-29T14:28:42.637894Z",
     "iopub.status.idle": "2026-06-29T14:28:51.869477Z",
     "shell.execute_reply": "2026-06-29T14:28:51.869072Z"
    }
   },
   "outputs": [],
   "source": [
    "from strands_robots import MockPolicy, Robot, create_policy\n",
    "from strands_robots.training import TrainSpec, create_trainer\n",
    "\n",
    "sim = Robot(\"so100\", mesh=False)\n",
    "sim.add_camera(name=\"front\", position=[0.5, 0.0, 0.4], target=[0.2, 0.0, 0.05])\n",
    "sim.start_recording(repo_id=\"local/nb3_demo\", root=ROOT, fps=30, task=\"pick up the red cube\", overwrite=True)\n",
    "sim.run_policy(robot_name=\"so100\", policy_object=MockPolicy(), instruction=\"pick up the red cube\", n_steps=60)\n",
    "sim.stop_recording()\n",
    "sim.destroy()\n",
    "print(\"recorded ->\", ROOT)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f261c07e",
   "metadata": {},
   "source": [
    "## 2. Train\n",
    "\n",
    "`create_trainer(provider)` returns a `Trainer` - the peer of `create_policy()`. The trainer is selected by the **same provider name** as the inference policy, so the same string that runs a policy also trains one. `TrainSpec` describes the run; `validate()` is a pure preflight that launches nothing if the spec is wrong. Here ACT trains from scratch for two steps on CPU.\n",
    "\n",
    "> Swap the provider to `\"groot\"` or `\"cosmos3\"` and only the provider string changes - the LeRobot, GR00T, and Cosmos training pipelines all sit behind the same `TrainSpec` + `Trainer` lifecycle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "234c5fcb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-29T14:28:51.871123Z",
     "iopub.status.busy": "2026-06-29T14:28:51.870958Z",
     "iopub.status.idle": "2026-06-29T14:28:56.078829Z",
     "shell.execute_reply": "2026-06-29T14:28:56.078287Z"
    }
   },
   "outputs": [],
   "source": [
    "trainer = create_trainer(\"lerobot_local\", device=\"cpu\")\n",
    "spec = TrainSpec(\n",
    "    dataset_root=ROOT,\n",
    "    base_model=\"\",  # ACT from scratch (smallest CPU path)\n",
    "    output_dir=OUT,\n",
    "    steps=2,\n",
    "    save_freq=2,\n",
    "    global_batch_size=2,\n",
    "    extra={\"policy_type\": \"act\", \"num_workers\": 0},\n",
    ")\n",
    "\n",
    "problems = trainer.validate(spec)\n",
    "assert not problems, problems\n",
    "\n",
    "result = trainer.train(spec)\n",
    "if result.status != \"success\":\n",
    "    # train() converts any failure into a TrainResult rather than raising, and\n",
    "    # result.message carries the underlying cause - including lerobot's own\n",
    "    # \"'accelerate' is required but not installed\" remedy. Without this check the\n",
    "    # cell prints status=error, and the next cell fails on a None checkpoint_dir\n",
    "    # instead of on the thing that actually went wrong.\n",
    "    raise RuntimeError(f\"training failed: {result.message}\")\n",
    "print(\"train status:\", result.status)\n",
    "print(\"checkpoint:\", result.checkpoint_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb9487c6",
   "metadata": {},
   "source": [
    "## 3. Load the trained checkpoint back\n",
    "\n",
    "The exported checkpoint loads through the same `create_policy()` entry point from [notebook 1](./01_getting_started.ipynb) - the loop closes: the thing you trained is a policy you can run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "910ff14c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-29T14:28:56.080351Z",
     "iopub.status.busy": "2026-06-29T14:28:56.080266Z",
     "iopub.status.idle": "2026-06-29T14:28:56.595008Z",
     "shell.execute_reply": "2026-06-29T14:28:56.594530Z"
    }
   },
   "outputs": [],
   "source": [
    "policy = create_policy(result.checkpoint_dir)\n",
    "print(\"loaded:\", type(policy).__name__)\n",
    "print(\"record -> train -> export -> load closed on CPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d946b26b",
   "metadata": {},
   "source": [
    "## Where to go from here\n",
    "\n",
    "- Raise `steps`, point `base_model` at a pretrained checkpoint, and run on a GPU for a real policy.\n",
    "- `stop_recording(bucket=...)` plus `stream_dataset()` let you train straight from a Hugging Face Storage Bucket without a full download.\n",
    "- The [`07_post_tune_any_policy.py`](../07_post_tune_any_policy.py) script is the same loop in one file."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "07f210a62e8d45149b3c8042c037a685": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "3cfc996f83844d62af349ce92c632043": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "44c03ac9d6a140f7909e31c6f5f095da": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "990f486a45ce4ca7875a17d6de88cbea": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_3cfc996f83844d62af349ce92c632043",
       "placeholder": "​",
       "style": "IPY_MODEL_a7c51642b1394393bc369bab5ad41984",
       "tabbable": null,
       "tooltip": null,
       "value": " 60/60 [00:00&lt;00:00, 2581.19 examples/s]"
      }
     },
     "a07e93b294a346f290630b3de8c0a929": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "a7c51642b1394393bc369bab5ad41984": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "d1a1a80378df4cb288c9e82889fd0764": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "e8f42cbe8afa4fbe8e77065a7d11d540": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_44c03ac9d6a140f7909e31c6f5f095da",
       "placeholder": "​",
       "style": "IPY_MODEL_d1a1a80378df4cb288c9e82889fd0764",
       "tabbable": null,
       "tooltip": null,
       "value": "Map: 100%"
      }
     },
     "f169e0a2128d48f58d8f7790a8723c16": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "f4cc24ebe2e34c7a998887a913845fdb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_e8f42cbe8afa4fbe8e77065a7d11d540",
        "IPY_MODEL_f7c696c177c44757ad683554c52f8500",
        "IPY_MODEL_990f486a45ce4ca7875a17d6de88cbea"
       ],
       "layout": "IPY_MODEL_a07e93b294a346f290630b3de8c0a929",
       "tabbable": null,
       "tooltip": null
      }
     },
     "f7c696c177c44757ad683554c52f8500": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_07f210a62e8d45149b3c8042c037a685",
       "max": 60,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_f169e0a2128d48f58d8f7790a8723c16",
       "tabbable": null,
       "tooltip": null,
       "value": 60
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
