{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {},
   "source": [
    "# The Streaming Data Loop\n",
    "\n",
    "The full physical-AI data loop in one notebook: **record** a dataset into a\n",
    "Hugging Face Storage Bucket, **stream** it back with no download, **train** a\n",
    "policy on the streamed data, and **load** the checkpoint ready to deploy.\n",
    "\n",
    "This is the notebook version of\n",
    "[`examples/06_agent_collect_and_stream.py`](https://github.com/strands-labs/robots/blob/main/examples/06_agent_collect_and_stream.py), the agent-driven data-loop example.\n",
    "\n",
    "**Requirements:** the bucket path needs `strands-robots >= 0.5.1`, which is the\n",
    "first release whose `[lerobot]` extra floors LeRobot at the `>= 0.6.1` that\n",
    "serves `stream_dataset(..., repo_type=\"bucket\")`.\n",
    "\n",
    "```bash\n",
    "pip install -U \"strands-robots[sim-mujoco,lerobot]>=0.5.1\"\n",
    "```\n",
    "\n",
    "Keep the `-U` and the version floor on any install line you adapt. Extras alone\n",
    "do not make a requirement unsatisfied, so a plain install into an environment\n",
    "that already has an older release reports `Requirement already satisfied` and\n",
    "upgrades nothing. On 0.4.1 the bucket read below then raises `TypeError: open()\n",
    "got an unexpected keyword argument 'repo_type'`, naming the keyword rather than\n",
    "the stale install behind it.\n",
    "\n",
    "Bucket steps additionally need `hf auth login`. Step 5 needs\n",
    "`lerobot[training]` on CPU as well as GPU - the trainer imports `accelerate`\n",
    "either way. The sim-only path (record + stream from a local root) runs on any\n",
    "laptop, no GPU, no credentials.\n",
    "\n",
    "The optional Step 7 (Isaac backend swap) additionally needs an RTX GPU +\n",
    "Isaac Sim 6.0+ and the `sim-isaac` extra; it self-skips otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acae54e37e7d407bbb7b55eff062a284",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "\n",
    "# macOS uses \"cgl\" for offscreen GL; 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\")\n",
    "\n",
    "ROOT = \"/tmp/nb5_dataset\"\n",
    "OUT = \"/tmp/nb5_ft\"\n",
    "BUCKET = None  # set to \"your-org/your-bucket\" to enable the bucket path\n",
    "RUN_ID = \"nb5_demo\"  # folder inside the bucket; one run per collection session\n",
    "\n",
    "shutil.rmtree(ROOT, ignore_errors=True)\n",
    "shutil.rmtree(OUT, ignore_errors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a63283cbaf04dbcab1f6479b197f3a8",
   "metadata": {},
   "source": [
    "## 1. Record a demonstration\n",
    "\n",
    "`Robot(\"so100\")` builds a MuJoCo simulation. We add a camera (so the dataset\n",
    "has video), run the Mock policy for 60 steps (structurally complete data), and\n",
    "stop recording. The result is a LeRobotDataset on disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd0d8092fe74a7c96281538738b07e2",
   "metadata": {},
   "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(\n",
    "    repo_id=\"local/nb5_demo\",\n",
    "    root=ROOT,\n",
    "    fps=30,\n",
    "    task=\"pick up the red cube\",\n",
    "    cameras=[\"front\"],  # record only the declared sensor, not the implicit 'default' view\n",
    "    overwrite=True,\n",
    ")\n",
    "# control_frequency must match the recording fps: the recorder writes one frame\n",
    "# per control step with no decimation, so a 50 Hz default rollout against a 30 fps\n",
    "# recording is refused rather than written at a distorted timestamp rate.\n",
    "result = sim.run_policy(\n",
    "    robot_name=\"so100\",\n",
    "    policy_object=MockPolicy(),\n",
    "    instruction=\"pick up the red cube\",\n",
    "    n_steps=60,\n",
    "    control_frequency=30.0,\n",
    ")\n",
    "if result.get(\"status\") != \"success\":\n",
    "    raise RuntimeError(f\"rollout failed: {result.get('content')}\")\n",
    "\n",
    "result = sim.stop_recording()\n",
    "if result.get(\"status\") != \"success\":\n",
    "    raise RuntimeError(f\"recording failed: {result.get('content')}\")\n",
    "print(\"recorded ->\", ROOT)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72eea5119410473aa328ad9291626812",
   "metadata": {},
   "source": [
    "## 2. See what the robot recorded\n",
    "\n",
    "Render the front camera so you can see the scene the dataset captured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8edb47106e1a46a883d545849b8ab81b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "# Render the sim before destroying it.\n",
    "frame = sim.render(camera_name=\"front\")\n",
    "for item in frame.get(\"content\", []):\n",
    "    if isinstance(item, dict) and \"image\" in item:\n",
    "        display(Image(data=item[\"image\"][\"source\"][\"bytes\"], width=480))\n",
    "        break\n",
    "sim.destroy()\n",
    "print(\"SO-100 arm with front camera - this is what the dataset contains.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10185d26023b46108eb7d9f57d49d2b3",
   "metadata": {},
   "source": [
    "## 3. Sync to a Hugging Face Storage Bucket (optional)\n",
    "\n",
    "If you have a bucket (`hf buckets create your-org/name --private` after\n",
    "`hf auth login`), set `BUCKET` in cell 1. The sync uploads only the bytes that\n",
    "changed (Xet deduplication), so daily re-syncs are fast.\n",
    "\n",
    "**Skip this cell** if running without credentials (the local path still works\n",
    "for Steps 4 and 5)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8763a12b2bbd4a93a75aff182afb95dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "if BUCKET:\n",
    "    from strands_robots.dataset_recorder import DatasetRecorder\n",
    "\n",
    "    # Reopen the dataset recorded in Step 1 and sync it. sync_to_bucket()\n",
    "    # returns a status dict - check it instead of assuming success.\n",
    "    rec = DatasetRecorder.resume(\"local/nb5_demo\", root=ROOT)\n",
    "    result = rec.sync_to_bucket(BUCKET, run_id=RUN_ID)\n",
    "    if result.get(\"status\") == \"success\":\n",
    "        print(f\"synced to bucket: {result['bucket_uri']}\")\n",
    "    else:\n",
    "        raise RuntimeError(f\"bucket sync failed: {result.get('message')}\")\n",
    "else:\n",
    "    print(\"BUCKET not set - skipping sync. Local path works for training below.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7623eae2785240b9bd12b16a66d81610",
   "metadata": {},
   "source": [
    "## 4. Stream the dataset back\n",
    "\n",
    "`stream_dataset()` is the read counterpart to `start_recording()`. It reads\n",
    "frames lazily with no full download. If you synced to a bucket (Step 3), pass\n",
    "`repo_type=\"bucket\"` to stream directly from it (requires LeRobot >= 0.6.1).\n",
    "\n",
    "Reading frames in capture order takes three arguments, not one. `shuffle=False`\n",
    "only fixes the random seed: the reader re-shards the dataset into\n",
    "`max_num_shards` shards and yields from a random reservoir of `buffer_size`\n",
    "frames, so `max_num_shards=1, buffer_size=1` are what collapse it to a single\n",
    "sequential pass. Use all three to inspect or replay an episode; drop all three\n",
    "to train, where the internal shuffle is what you want and video decoding\n",
    "parallelizes across worker processes.\n",
    "\n",
    "`episodes=[n]` does not filter: LeRobot accepts the argument and never applies\n",
    "it, so the whole dataset streams back with no error. Filter on\n",
    "`frame[\"episode_index\"]` while iterating instead.\n",
    "\n",
    "`sync_to_bucket` writes each run to its own `run_id` folder inside the bucket\n",
    "(`hf://buckets/{BUCKET}/{RUN_ID}`), so the streaming repo id is\n",
    "`f\"{BUCKET}/{RUN_ID}\"`, not `BUCKET` on its own. The bucket namespace is the\n",
    "first two segments of that id; everything after it is the path within the\n",
    "bucket. Passing `BUCKET` alone looks for `meta/` at the bucket root and fails\n",
    "with `FileNotFoundError`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cdc8c89c7104fffa095e18ddfef8986",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = Robot(\"so100\", mesh=False)\n",
    "\n",
    "if BUCKET:\n",
    "    # Read from the same path Step 3 wrote to: bucket namespace + run_id.\n",
    "    bucket_repo_id = f\"{BUCKET}/{RUN_ID}\"\n",
    "    reader = sim.stream_dataset(\n",
    "        bucket_repo_id, repo_type=\"bucket\", shuffle=False, max_num_shards=1, buffer_size=1\n",
    "    )\n",
    "    print(f\"streaming from bucket: hf://buckets/{bucket_repo_id}\")\n",
    "else:\n",
    "    reader = sim.stream_dataset(\n",
    "        \"local/nb5_demo\", root=ROOT, shuffle=False, max_num_shards=1, buffer_size=1\n",
    "    )\n",
    "    print(f\"streaming from local root: {ROOT}\")\n",
    "\n",
    "print(f\"episodes: {reader.num_episodes} | frames: {reader.num_frames} | fps: {reader.fps}\")\n",
    "\n",
    "for n, frame in enumerate(reader):\n",
    "    if n == 0:\n",
    "        img = frame.get(\"observation.images.front\")\n",
    "        state = frame.get(\"observation.state\")\n",
    "        print(\n",
    "            f\"frame 0 - image: {tuple(img.shape) if img is not None else None}, state: {tuple(state.shape) if state is not None else None}\"\n",
    "        )\n",
    "    if n >= 4:\n",
    "        break\n",
    "print(f\"streamed {n + 1} frames - camera decoded on the fly.\")\n",
    "sim.destroy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b118ea5561624da68c537baed56e602f",
   "metadata": {},
   "source": [
    "## 5. Train a policy on the dataset\n",
    "\n",
    "`create_trainer(\"lerobot_local\")` returns a Trainer (the peer of\n",
    "`create_policy()`). A `TrainSpec` describes the run. Here we train ACT for 2\n",
    "steps on CPU (fast for the demo). On a GPU with 500 steps this takes about 133\n",
    "seconds on an NVIDIA L4 (`g6.4xlarge`).\n",
    "\n",
    "This step needs `lerobot[training]` whichever device you are on: the trainer\n",
    "imports `accelerate` before it looks at the device. The cell checks\n",
    "`result.status` and re-raises `result.message`, which carries LeRobot's own\n",
    "install remedy - `train()` reports failure in its result rather than raising, so\n",
    "an unchecked call would hand Step 6 a `checkpoint_dir` of `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "938c804e27f84196a10c8828c723f798",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = create_trainer(\"lerobot_local\", device=\"cpu\")\n",
    "spec = TrainSpec(\n",
    "    dataset_root=ROOT,\n",
    "    base_model=\"\",  # ACT from scratch\n",
    "    output_dir=OUT,\n",
    "    steps=2,  # raise to 500+ on a GPU for a real checkpoint\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(f\"train status: {result.status}\")\n",
    "print(f\"checkpoint: {result.checkpoint_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504fb2a444614c0babb325280ed9130a",
   "metadata": {},
   "source": [
    "## 6. Load the trained checkpoint\n",
    "\n",
    "The same `create_policy()` entry point loads the checkpoint we just produced.\n",
    "On hardware you would pass `mode=\"real\"` to deploy against a physical arm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59bbdb311c014d738909a11f9e486628",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = create_policy(result.checkpoint_dir)\n",
    "print(f\"loaded: {type(policy).__name__}\")\n",
    "print(\"record -> stream -> train -> load: the data loop closes.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Same loop, different backend: Isaac (optional - RTX GPU + Isaac Sim)\n",
    "\n",
    "The record -> stream loop above ran on the default MuJoCo backend - no GPU, no\n",
    "credentials. Backend parity for recording landed in #1552 (`IsaacRecordingMixin`),\n",
    "so the *same* loop runs on NVIDIA Isaac Sim with a single `backend=\"isaac\"` swap.\n",
    "\n",
    "This step is optional and self-skips unless you have:\n",
    "\n",
    "- an RTX GPU and Isaac Sim 6.0+ installed out-of-band, and\n",
    "- `pip install -U 'strands-robots[sim-isaac,lerobot]'`.\n",
    "\n",
    "One pacing difference from the MuJoCo cells above: Isaac renders at\n",
    "`rendering_dt = 1/30` by default, so keep `control_frequency <= 1 / rendering_dt`.\n",
    "We use `fps=10` / `control_frequency=10.0` here (not the MuJoCo cells' `fps=30`)\n",
    "so every recorded frame reflects a freshly rendered product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from strands_robots.simulation.isaac import IsaacSimulation\n",
    "\n",
    "available, reason = IsaacSimulation.is_available()\n",
    "if not available:\n",
    "    print(f\"Isaac Sim not available - skipping the backend-swap demo: {reason}\")\n",
    "else:\n",
    "    ROOT_ISAAC = \"/tmp/nb5_isaac_dataset\"\n",
    "    shutil.rmtree(ROOT_ISAAC, ignore_errors=True)\n",
    "    isaac = Robot(\"so100\", backend=\"isaac\", mesh=False)  # same loop, one kwarg changed\n",
    "    isaac.add_camera(name=\"front\", position=[0.5, 0.0, 0.4], target=[0.2, 0.0, 0.05])\n",
    "    isaac.start_recording(\n",
    "        repo_id=\"local/nb5_isaac\",\n",
    "        root=ROOT_ISAAC,\n",
    "        fps=10,\n",
    "        task=\"pick up the red cube\",\n",
    "        cameras=[\"front\"],\n",
    "        overwrite=True,\n",
    "    )\n",
    "    isaac.run_policy(\n",
    "        robot_name=\"so100\",\n",
    "        policy_object=MockPolicy(),\n",
    "        n_steps=20,\n",
    "        control_frequency=10.0,\n",
    "        fast_mode=True,\n",
    "        instruction=\"pick up the red cube\",\n",
    "    )\n",
    "    isaac.stop_recording()\n",
    "    reader = isaac.stream_dataset(\n",
    "        \"local/nb5_isaac\", root=ROOT_ISAAC, shuffle=False, max_num_shards=1, buffer_size=1\n",
    "    )\n",
    "    print(f\"isaac dataset: episodes={reader.num_episodes} frames={reader.num_frames}\")\n",
    "    isaac.destroy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b43b363d81ae4b689946ece5c682cd59",
   "metadata": {},
   "source": [
    "## Where to go from here\n",
    "\n",
    "- Raise `steps` to 500 and run on a GPU to get a checkpoint that loads and runs.\n",
    "  It will not be a good policy: this notebook records one 60-frame episode from\n",
    "  the Mock policy, so what the run proves is the record-train-load path, not the\n",
    "  behaviour. Collect real demonstrations before you judge a checkpoint.\n",
    "- Set `BUCKET` and run with `hf auth login` to exercise the full bucket path.\n",
    "- Swap `mode=\"real\"` on `Robot(\"so100\")` to deploy to a physical SO-101.\n",
    "- Read the [first post in the series](https://huggingface.co/blog/amazon/strands-lerobot-hub-to-hardware) for the full sim-to-hardware walkthrough.\n",
    "- See [examples/06_agent_collect_and_stream.py](https://github.com/strands-labs/robots/blob/main/examples/06_agent_collect_and_stream.py) for the agent-driven version."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}