{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c7e7edd",
   "metadata": {},
   "source": [
    "# Agentic Fleet Orchestration with Strands Agents\n",
    "\n",
    "Give **one agent one goal in plain English**, and watch it plan across a fleet of\n",
    "robots, assign each sub-task to a robot that can do it, and dispatch over a safe\n",
    "command mesh, from simulation to hardware.\n",
    "\n",
    "This is the companion notebook to the blog post *Agentic AI for autonomous fleet\n",
    "operations*. It is the notebook form of\n",
    "[`fleet_orchestration_demo.py`](https://github.com/strands-labs/robots), and every\n",
    "cell runs the real Strands Robots API.\n",
    "\n",
    "**What you will do:**\n",
    "\n",
    "1. Bring up a heterogeneous fleet in one simulated world and read each robot's capabilities.\n",
    "2. Give a natural-language goal and watch a planner decompose it into sub-tasks.\n",
    "3. Dispatch the assignments in a synchronized multi-robot control loop.\n",
    "4. Take a robot offline mid-mission and see the plan re-form safely.\n",
    "5. (Optional) Dispatch over a Zenoh peer mesh with a human-in-the-loop approval on every action.\n",
    "\n",
    "**The whole default path runs on a laptop, CPU only. No GPU, no cloud credentials.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed4c7c6a",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Install Strands Robots with the simulation extra. The `mesh` extra is only needed\n",
    "for the optional Step 5.\n",
    "\n",
    "```bash\n",
    "uv pip install -U \"strands-robots[sim-mujoco,mesh]\"\n",
    "```\n",
    "\n",
    "On macOS set the MuJoCo GL backend to `cgl`; on headless Linux use `egl`. This must\n",
    "be set before `strands_robots` triggers the first `import mujoco`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b119ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# MuJoCo locks its GL backend at first import, so set this before importing the SDK.\n",
    "# cgl is macOS-only; egl is the headless-Linux backend. An exported MUJOCO_GL wins.\n",
    "os.environ.setdefault(\"MUJOCO_GL\", \"cgl\" if sys.platform == \"darwin\" else \"egl\")\n",
    "\n",
    "import logging\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(message)s\")\n",
    "\n",
    "from strands_robots.simulation import Simulation\n",
    "from strands_robots.policies.mock import MockPolicy\n",
    "print(\"Strands Robots ready.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5abd8b67",
   "metadata": {},
   "source": [
    "## Step 1 - Bring up a fleet and read its capabilities\n",
    "\n",
    "A fleet is a set of robots with different physical abilities. What the allocator\n",
    "needs to know is **what each robot can do**, not what it is: an arm offers\n",
    "`manipulate`, a mobile base offers `navigate` and `transport`.\n",
    "\n",
    "The registry records a `category` for every robot it knows (`arm`, `mobile`,\n",
    "`mobile_manip`, `humanoid`, ...) but no capability field, so the mapping from\n",
    "category to capabilities is orchestration policy and is defined below. Deriving it\n",
    "from `category` rather than listing robots by name means a new vendor's robot is\n",
    "classified as soon as the registry knows it, and an unmapped category raises\n",
    "instead of silently defaulting to something the robot cannot do. The allocator\n",
    "then matches a sub-task's required capability against that set, so adding a\n",
    "capability class is one row in a map, not new planner code.\n",
    "\n",
    "**`grasp` is the exception, and it is the interesting one.** A category cannot\n",
    "answer it: 23 registry robots have category `arm` and 3 of them declare a `gripper`\n",
    "block. So `grasp` is claimed from that block alone - the field that names the\n",
    "gripper's actuators and which end of their range is closed, i.e. the only field\n",
    "that says the mechanism exists. Over-advertising a capability is the one\n",
    "allocation error an allocator cannot detect afterwards: the robot accepts the\n",
    "task, the plan looks complete, and the failure happens in the world.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Capabilities are derived from the registry's own fields rather than hard-coded per\n",
    "# robot, so a robot the registry knows about is classified without editing this\n",
    "# notebook. The registry has no capability field of its own; the mapping below is\n",
    "# the orchestration policy, and it lives here.\n",
    "#\n",
    "# `grasp` is deliberately NOT claimed from the category. It is claimed only when the\n",
    "# registry entry carries a `gripper` block, because that block is the only thing that\n",
    "# says the mechanism exists - it names the actuators and which end of their range is\n",
    "# closed. Measured on the shipped registry, 23 robots have category `arm` and 3 of\n",
    "# them declare a gripper, so inferring `grasp` from the category would advertise it\n",
    "# for 20 arms that never declared one. That is the one allocation error an allocator\n",
    "# cannot detect afterwards: the robot accepts the task and fails in the world. The\n",
    "# library resolves a gripper the same way and for the same reason - see\n",
    "# `MuJoCoMotionPrimitives._registry_gripper_metadata`, which prefers this block over\n",
    "# its name heuristic rather than the reverse.\n",
    "from strands_robots.registry import get_robot, resolve_name\n",
    "\n",
    "CATEGORY_CAPABILITIES = {\n",
    "    \"arm\":          {\"manipulate\"},\n",
    "    \"bimanual\":     {\"manipulate\"},\n",
    "    \"hand\":         {\"manipulate\"},\n",
    "    \"mobile\":       {\"navigate\", \"transport\"},\n",
    "    \"mobile_manip\": {\"navigate\", \"transport\", \"manipulate\"},\n",
    "    \"humanoid\":     {\"navigate\", \"manipulate\"},\n",
    "    \"aerial\":       {\"navigate\", \"inspect\"},\n",
    "    # Mapped to nothing on purpose: an expressive robot (`reachy_mini`) has no\n",
    "    # capability this allocator dispatches, so every sub-task it is offered is left\n",
    "    # UNASSIGNED. Empty and explicit, rather than unmapped and raising, because\n",
    "    # \"advertises nothing\" is a real answer and a missing category is a gap.\n",
    "    \"expressive\":   set(),\n",
    "}\n",
    "\n",
    "\n",
    "def capabilities_of(robot_type):\n",
    "    \"\"\"Capability tags for a registry name or alias.\n",
    "\n",
    "    Raises `ValueError` for a name the registry does not know and for a category\n",
    "    this notebook has no row for. Neither case may return a default: a robot that\n",
    "    quietly advertises `manipulate` gets handed work it may be unable to do.\n",
    "    \"\"\"\n",
    "    entry = get_robot(robot_type)\n",
    "    if entry is None:\n",
    "        raise ValueError(\n",
    "            f\"{robot_type!r} is not in the robot registry (aliases resolve too, so \"\n",
    "            \"check the spelling against `list_robots()`). Capabilities cannot be \"\n",
    "            \"guessed for an unknown robot.\"\n",
    "        )\n",
    "    category = entry.get(\"category\")\n",
    "    caps = CATEGORY_CAPABILITIES.get(category)\n",
    "    if caps is None:\n",
    "        raise ValueError(\n",
    "            f\"{robot_type!r} (canonical {resolve_name(robot_type)!r}) has category \"\n",
    "            f\"{category!r}, which this notebook maps to no capabilities. Add it to \"\n",
    "            \"CATEGORY_CAPABILITIES before dispatching to it.\"\n",
    "        )\n",
    "    caps = set(caps)\n",
    "    if entry.get(\"gripper\"):\n",
    "        caps.add(\"grasp\")\n",
    "    return caps\n",
    "\n",
    "\n",
    "for name in (\"so101\", \"ur5e\", \"go2\", \"reachy_mini\"):\n",
    "    print(f\"{name:12s} {sorted(capabilities_of(name))}\")\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The fleet itself\n",
    "\n",
    "With the capability model in place, bring up the world. `arm_a` and `arm_b` declare\n",
    "grippers, so they advertise `grasp`; `go2` resolves through its alias to\n",
    "`unitree_go2` (category `mobile`) and advertises `navigate` and `transport`.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Stand up a heterogeneous fleet in one world: two arms + one mobile base.\n",
    "sim = Simulation()\n",
    "sim.create_world()\n",
    "\n",
    "FLEET_SPEC = [\n",
    "    (\"arm_a\", \"so101\", [0.0, 0.0, 0.0]),\n",
    "    (\"arm_b\", \"so100\", [0.6, 0.0, 0.0]),\n",
    "    (\"base\",  \"go2\",   [0.0, 0.8, 0.0]),\n",
    "]\n",
    "\n",
    "fleet = {}  # instance name -> {type, capabilities, online}\n",
    "for name, rtype, pos in FLEET_SPEC:\n",
    "    res = sim.add_robot(name=name, data_config=rtype, position=pos)\n",
    "    assert res.get(\"status\") == \"success\", res\n",
    "    fleet[name] = {\"type\": rtype, \"capabilities\": capabilities_of(rtype), \"online\": True}\n",
    "\n",
    "for name, m in fleet.items():\n",
    "    print(f\"{name:6s} [{m['type']:6s}] {sorted(m['capabilities'])}\")\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "b6b61108",
   "metadata": {},
   "source": [
    "## Step 2 - Give a goal, and watch it plan\n",
    "\n",
    "The planner turns one natural-language goal into a set of sub-tasks, each tagged\n",
    "with the capability it needs, then assigns each to a capable online robot.\n",
    "\n",
    "We show the **rule-based planner** first: it needs no model provider, so the whole\n",
    "pipeline is provable end-to-end before an LLM is introduced. The\n",
    "*allocation* (the safety-relevant \"which robot runs this\" decision) stays\n",
    "deterministic and auditable no matter which planner produces the sub-tasks. Step 2b\n",
    "swaps in a Strands agent that produces the same sub-task shape by reasoning instead\n",
    "of keyword matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c993a187",
   "metadata": {},
   "outputs": [],
   "source": [
    "def decompose_rule(goal):\n",
    "    \"\"\"Keyword decomposition: the deterministic floor, no model needed.\"\"\"\n",
    "    g = goal.lower()\n",
    "    tasks = []\n",
    "    if any(w in g for w in (\"move\", \"transport\", \"bring\", \"carry\", \"fetch\")):\n",
    "        tasks.append({\"description\": \"transport the bin to the work zone\", \"capability\": \"transport\"})\n",
    "    if any(w in g for w in (\"sort\", \"pick\", \"place\", \"grasp\", \"load\")):\n",
    "        tasks.append({\"description\": \"pick and place the parts\", \"capability\": \"grasp\"})\n",
    "    if not tasks:\n",
    "        tasks.append({\"description\": goal.strip(), \"capability\": \"manipulate\"})\n",
    "    return tasks\n",
    "\n",
    "def allocate(tasks, fleet):\n",
    "    \"\"\"Assign each sub-task to a capable ONLINE robot, load-balanced.\n",
    "\n",
    "    A robot may hold more than one sub-task, so we do not claim it after the\n",
    "    first. Among capable online robots pick the least-loaded, ties by fleet order.\n",
    "    An unmatched task is left UNASSIGNED - the safe outcome, never mis-assigned.\n",
    "    \"\"\"\n",
    "    load = {name: 0 for name in fleet}\n",
    "    order = list(fleet)\n",
    "    for t in tasks:\n",
    "        candidates = [n for n, m in fleet.items() if m[\"online\"] and t[\"capability\"] in m[\"capabilities\"]]\n",
    "        if not candidates:\n",
    "            t[\"assigned_to\"] = None\n",
    "            continue\n",
    "        chosen = min(candidates, key=lambda n: (load[n], order.index(n)))\n",
    "        t[\"assigned_to\"] = chosen\n",
    "        load[chosen] += 1\n",
    "    return tasks\n",
    "\n",
    "def render_plan(goal, tasks):\n",
    "    assigned = sum(t[\"assigned_to\"] is not None for t in tasks)\n",
    "    print(f'Goal: \"{goal}\"')\n",
    "    print(f\"Plan ({assigned}/{len(tasks)} assigned):\")\n",
    "    for i, t in enumerate(tasks, 1):\n",
    "        who = t[\"assigned_to\"] or \"UNASSIGNED (no capable online robot)\"\n",
    "        print(f\"  {i}. [{t['capability']}] {t['description']} -> {who}\")\n",
    "\n",
    "GOAL = \"Move the parts bin to the work zone and sort the parts into trays.\"\n",
    "tasks = allocate(decompose_rule(GOAL), fleet)\n",
    "render_plan(GOAL, tasks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac464f1b",
   "metadata": {},
   "source": [
    "### Step 2b (optional) - Let a Strands agent do the decomposition\n",
    "\n",
    "Same sub-task shape, produced by reasoning instead of keyword matching. This cell\n",
    "needs a Strands model provider (for example Amazon Bedrock credentials). If none is\n",
    "available it falls back to the rule-based planner, so the notebook still runs.\n",
    "\n",
    "Only the *decomposition* is delegated to the model; `allocate(...)` is unchanged, so\n",
    "the assignment stays deterministic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e55c081",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "def decompose_agent(goal, fleet, model_id=None, region=None):\n",
    "    from strands import Agent\n",
    "    from strands.models import BedrockModel\n",
    "\n",
    "    system = (\n",
    "        \"You decompose a robot-fleet goal into sub-tasks. Respond with ONLY a JSON \"\n",
    "        \"array. Each element has 'description' (short imperative) and 'capability' \"\n",
    "        \"(one of: manipulate, grasp, navigate, transport, inspect). Produce the \"\n",
    "        \"minimum sub-tasks that cover the goal. No prose.\"\n",
    "    )\n",
    "    model = BedrockModel(model_id=model_id or \"us.anthropic.claude-opus-4-8\",\n",
    "                         region_name=region or os.getenv(\"AWS_REGION\", \"us-east-1\"))\n",
    "    caps = sorted({c for m in fleet.values() for c in m['capabilities']})\n",
    "    # callback_handler=None silences the SDK's default stdout streaming.\n",
    "    agent = Agent(model=model, system_prompt=system, callback_handler=None)\n",
    "    raw = str(agent(f'Fleet capabilities: {caps}. Goal: \"{goal}\". Return the JSON array.')).strip()\n",
    "    s, e = raw.find(\"[\"), raw.rfind(\"]\")\n",
    "    data = json.loads(raw[s:e + 1]) if s != -1 and e != -1 else []\n",
    "    out = [{\"description\": str(d.get(\"description\", \"task\")), \"capability\": str(d[\"capability\"])}\n",
    "           for d in data if isinstance(d, dict) and d.get(\"capability\")]\n",
    "    if not out:\n",
    "        raise ValueError(f\"agent returned no parseable sub-tasks: {raw[:200]!r}\")\n",
    "    return out\n",
    "\n",
    "try:\n",
    "    agent_tasks = allocate(decompose_agent(GOAL, fleet), fleet)\n",
    "    print(\"Agent planner:\\n\")\n",
    "    render_plan(GOAL, agent_tasks)\n",
    "except Exception as e:\n",
    "    print(f\"Agent planner unavailable ({type(e).__name__}: {e});\\nfalling back to rule-based plan above.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f5f8019",
   "metadata": {},
   "source": [
    "## Step 3 - Dispatch in a synchronized control loop\n",
    "\n",
    "With assignments in hand, drive every assigned robot together in one control loop.\n",
    "A robot that owns several sub-tasks runs one policy with its instructions joined, so\n",
    "`run_multi_policy` gets exactly one entry per robot.\n",
    "\n",
    "Here each robot runs a `MockPolicy` so the loop runs anywhere with no model weights.\n",
    "On hardware you would swap in a real policy per robot; the dispatch code does not change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a79b765",
   "metadata": {},
   "outputs": [],
   "source": [
    "per_robot = {}\n",
    "for t in tasks:\n",
    "    if t[\"assigned_to\"]:\n",
    "        per_robot.setdefault(t[\"assigned_to\"], []).append(t[\"description\"])\n",
    "\n",
    "policies = {name: MockPolicy() for name in per_robot}\n",
    "instructions = {name: \"; then \".join(descs) for name, descs in per_robot.items()}\n",
    "\n",
    "result = sim.run_multi_policy(\n",
    "    policies=policies,\n",
    "    instructions=instructions,\n",
    "    n_steps=30,\n",
    "    control_frequency=50.0,\n",
    "    action_horizon=4,\n",
    ")\n",
    "print(result.get(\"content\", [{}])[0].get(\"text\", result))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47a5fc1d",
   "metadata": {},
   "source": [
    "## Step 4 - Handle failure without mis-assigning work\n",
    "\n",
    "A fleet plan is only credible if it degrades safely. Take the mobile base offline and\n",
    "re-plan. No other online robot advertises `transport`, so that sub-task is left\n",
    "**unassigned** rather than handed to a grasping arm that cannot carry a bin. The\n",
    "grasp sub-task still routes to an arm.\n",
    "\n",
    "Leaving it unassigned is the correct outcome: silently reassigning it would produce a\n",
    "confident failure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3527e406",
   "metadata": {},
   "outputs": [],
   "source": [
    "fleet[\"base\"][\"online\"] = False\n",
    "print(\"base is now OFFLINE. Re-planning...\\n\")\n",
    "\n",
    "tasks_after = allocate(decompose_rule(GOAL), fleet)\n",
    "render_plan(GOAL, tasks_after)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66aa7654",
   "metadata": {},
   "source": [
    "## Step 5 (optional) - Dispatch over the mesh, with a human in the loop\n",
    "\n",
    "The steps above dispatched in-process. In a real deployment the robots are separate\n",
    "peers on a [Zenoh](https://zenoh.io/) mesh, and the orchestrator sends each\n",
    "assignment as a mesh command. **Every physically-actuating action pauses for a human\n",
    "approval**, delivered out-of-band of the model's tool arguments, so a prompt-injection\n",
    "attempt cannot smuggle its own approval flag into the command body.\n",
    "\n",
    "This cell needs the `mesh` extra and stands up a second peer, so it is heavier than\n",
    "the default path. The approval function here auto-approves for an unattended run; in\n",
    "an operator setting you would prompt a human. A declined action is never dispatched\n",
    "and, per the mesh's rate-limit semantics, consumes no rate-limit slot.\n",
    "\n",
    "> This cell is illustrative of the mesh API shape. Run it only with the `mesh` extra\n",
    "> installed; it is safe to skip and does not affect Steps 1-4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd8f03a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def auto_approve(action, target, instruction):\n",
    "    \"\"\"Out-of-band approver. Returns a bool the model cannot influence.\"\"\"\n",
    "    print(f\"[HITL] auto-approved {action} -> {target}: {instruction!r}\")\n",
    "    return True\n",
    "\n",
    "def dispatch_over_mesh(sim, orch_mesh, per_robot, approve=auto_approve, n_steps=30):\n",
    "    sim_peer = sim.peer_id\n",
    "    dispatched, denied = [], []\n",
    "    for name, descs in per_robot.items():\n",
    "        instruction = \"; then \".join(descs)\n",
    "        # HITL gate BEFORE the command goes on the wire.\n",
    "        if not approve(\"start\", name, instruction):\n",
    "            denied.append(name)\n",
    "            continue\n",
    "        r = orch_mesh.send(\n",
    "            sim_peer,\n",
    "            {\"action\": \"start\", \"instruction\": instruction,\n",
    "             \"policy_provider\": \"mock\", \"robot_name\": name, \"n_steps\": n_steps},\n",
    "            timeout=15.0,\n",
    "        )\n",
    "        ok = isinstance(r, dict) and r.get(\"result\", {}).get(\"status\") == \"success\"\n",
    "        dispatched.append((name, ok))\n",
    "    return {\"dispatched\": dispatched, \"denied\": denied}\n",
    "\n",
    "# To run this end to end, build the fleet with the mesh joined:\n",
    "#\n",
    "#   os.environ.setdefault(\"STRANDS_MESH_LOCAL_DEV\", \"1\")  # dev-only: no mTLS\n",
    "#   from strands_robots.mesh import init_mesh\n",
    "#   sim_mesh = init_mesh(sim, peer_id=None, peer_type=\"sim\")\n",
    "#   sim.mesh, sim.peer_id = sim_mesh, sim_mesh.peer_id\n",
    "#   ... add robots ...\n",
    "#   orch_mesh = init_mesh(_Owner(), peer_id=\"fleet-orchestrator\")  # a SEPARATE peer\n",
    "#   result = dispatch_over_mesh(sim, orch_mesh, per_robot)\n",
    "#\n",
    "# A peer cannot RPC its own peer_id in-process, so the orchestrator must be a\n",
    "# distinct peer from the sim that hosts the robots.\n",
    "print(\"Mesh dispatch helper defined. See the commented setup to run it with the mesh extra.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e621456",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "\n",
    "The `Simulation` owns a thread pool, a MuJoCo world, and a temp directory. Always\n",
    "destroy it when done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dfdb3f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.destroy()\n",
    "print(\"Fleet torn down.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "153c139a",
   "metadata": {},
   "source": [
    "## Where to go from here\n",
    "\n",
    "- **Hardware.** The `Robot()` abstraction is the same on real hardware: a physical\n",
    "  SO-101 joins the same mesh, answers the same commands, and takes the same\n",
    "  assignments. Natural-language control of a physical arm through this stack works\n",
    "  today.\n",
    "- **Real policies.** Swap `MockPolicy` for a policy trained on your robot and scene.\n",
    "  Completing a manipulation task autonomously needs a fine-tune on teleoperated\n",
    "  demonstrations - the record-train-deploy loop covered in the streaming data loop post.\n",
    "- **Deploy the planner.** The planner is an agent that reasons and dispatches; it\n",
    "  holds no model weights, runs on CPU, and deploys to Amazon Bedrock AgentCore the\n",
    "  way any Strands agent does. Heavy vision-language-action inference runs on a\n",
    "  separate GPU, called over the network.\n",
    "\n",
    "See the [Strands Robots documentation](https://strands-labs.github.io/robots/) for\n",
    "the robot catalog, simulation, policy providers, and the mesh in depth."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
