|
| 1 | +"""Drive a tabletop pick with a reactive behavior tree that recovers on failure. |
| 2 | +
|
| 3 | +`01_pick_and_retry.py` recovers from a grasp miss with imperative `if/else` |
| 4 | +control flow. This example does the *same* tabletop task on the *same* world |
| 5 | +(`Tabletop2D`), but the recovery is expressed **declaratively** as a reactive |
| 6 | +behavior tree (BT) — the structure roboticists actually reach for when a robot |
| 7 | +has to keep retrying and re-perceiving until it succeeds. |
| 8 | +
|
| 9 | +The tree is ticked once per control step. Each tick walks the tree and yields |
| 10 | +exactly one environment action (the leaf that is `RUNNING`): |
| 11 | +
|
| 12 | + Fallback "pick the block" |
| 13 | + ├── Condition object_in_gripper? # already holding -> whole tree SUCCESS |
| 14 | + └── Fallback "grasp or recover" |
| 15 | + ├── Sequence "confident grasp" |
| 16 | + │ ├── Condition belief_confident? # localized and uncertainty small? |
| 17 | + │ └── Action grasp_at_belief # pick(belief_mean) |
| 18 | + └── Action relook_to_refine # recovery: move to a new viewpoint |
| 19 | +
|
| 20 | +A Fallback (a.k.a. Selector) ticks its children left to right and stops at the |
| 21 | +first that does not FAIL, so it reads as "try the primary thing; if it isn't |
| 22 | +possible, do the recovery." The single `relook_to_refine` recovery leaf covers |
| 23 | +*both* failure modes the world throws: |
| 24 | +
|
| 25 | + * **occlusion** — the object starts behind an occluder, so `belief_confident?` |
| 26 | + fails (nothing localized yet) and the tree falls through to `relook`, which |
| 27 | + scans a fresh viewpoint until the detector sees the block. |
| 28 | + * **grasp miss** — a missed grasp grows the belief radius back above the |
| 29 | + confidence threshold, so on the next tick `belief_confident?` fails again and |
| 30 | + the tree falls back to `relook` to gather a better view before re-grasping. |
| 31 | +
|
| 32 | +That is the lesson: a grasp miss is not handled by a special-case branch; it |
| 33 | +simply makes a precondition false, and the *same* declarative fallback re-runs |
| 34 | +active perception before retrying. Recovery is a property of the tree's shape, |
| 35 | +not of a hand-written recovery routine. |
| 36 | +
|
| 37 | +Success: the block is lifted (`gripper.holding` set) before max_steps. |
| 38 | +Failure: grasp_miss (recoverable until the world's attempt budget is spent, |
| 39 | +then terminal) and tree_exhausted (terminal - ran out of steps still empty). |
| 40 | +
|
| 41 | +References: |
| 42 | + * M. Colledanchise and P. Ogren, "Behavior Trees in Robotics and AI: An |
| 43 | + Introduction," CRC Press, 2018. arXiv:1709.00084. |
| 44 | + * BehaviorTree.CPP (Faconti, Colledanchise) https://www.behaviortree.dev/ and |
| 45 | + py_trees (Stonier) https://github.com/splintered-reality/py_trees - the |
| 46 | + Sequence/Fallback tick semantics mirrored here, also used by the ROS 2 |
| 47 | + Nav2 BT Navigator for reactive recovery. |
| 48 | +""" |
| 49 | + |
| 50 | +from __future__ import annotations |
| 51 | + |
| 52 | +import argparse |
| 53 | +import sys |
| 54 | +from enum import Enum |
| 55 | +from pathlib import Path |
| 56 | +from typing import Any, Callable |
| 57 | + |
| 58 | +import numpy as np |
| 59 | + |
| 60 | +ROOT = Path(__file__).resolve().parents[2] |
| 61 | +if str(ROOT) not in sys.path: |
| 62 | + sys.path.insert(0, str(ROOT)) |
| 63 | + |
| 64 | +from pir.core.types import Failure, Trace |
| 65 | +from pir.worlds.tabletop_2d import Tabletop2D |
| 66 | + |
| 67 | + |
| 68 | +# --- A minimal reactive behavior tree -------------------------------------- |
| 69 | +# |
| 70 | +# Three return statuses and two composites (Sequence, Fallback) plus leaf |
| 71 | +# Condition/Action nodes are enough to express the whole recovery policy. This |
| 72 | +# mirrors the core of py_trees / BehaviorTree.CPP, kept small enough to read. |
| 73 | + |
| 74 | + |
| 75 | +class Status(Enum): |
| 76 | + SUCCESS = "success" |
| 77 | + FAILURE = "failure" |
| 78 | + RUNNING = "running" |
| 79 | + |
| 80 | + |
| 81 | +class Node: |
| 82 | + """Base class: a node is ticked and returns a Status.""" |
| 83 | + |
| 84 | + name = "node" |
| 85 | + |
| 86 | + def tick(self, bb: "Blackboard") -> Status: # pragma: no cover - overridden |
| 87 | + raise NotImplementedError |
| 88 | + |
| 89 | + |
| 90 | +class Sequence(Node): |
| 91 | + """Tick children in order; FAIL/RUNNING short-circuits, all SUCCESS -> SUCCESS. |
| 92 | +
|
| 93 | + Reactive: re-ticked from the first child every control step, so a condition |
| 94 | + that flips to FAILURE immediately re-routes the tree. |
| 95 | + """ |
| 96 | + |
| 97 | + def __init__(self, name: str, children: list[Node]) -> None: |
| 98 | + self.name = name |
| 99 | + self.children = children |
| 100 | + |
| 101 | + def tick(self, bb: "Blackboard") -> Status: |
| 102 | + for child in self.children: |
| 103 | + status = child.tick(bb) |
| 104 | + if status is not Status.SUCCESS: |
| 105 | + return status |
| 106 | + return Status.SUCCESS |
| 107 | + |
| 108 | + |
| 109 | +class Fallback(Node): |
| 110 | + """Tick children in order; first non-FAILURE wins, all FAILURE -> FAILURE. |
| 111 | +
|
| 112 | + This is where recovery lives: the primary branch comes first, the recovery |
| 113 | + branch comes after, and the tree falls through to recovery only when the |
| 114 | + primary branch reports FAILURE. |
| 115 | + """ |
| 116 | + |
| 117 | + def __init__(self, name: str, children: list[Node]) -> None: |
| 118 | + self.name = name |
| 119 | + self.children = children |
| 120 | + |
| 121 | + def tick(self, bb: "Blackboard") -> Status: |
| 122 | + for child in self.children: |
| 123 | + status = child.tick(bb) |
| 124 | + if status is not Status.FAILURE: |
| 125 | + return status |
| 126 | + return Status.FAILURE |
| 127 | + |
| 128 | + |
| 129 | +class Condition(Node): |
| 130 | + """A leaf that reads the blackboard and returns SUCCESS or FAILURE at once.""" |
| 131 | + |
| 132 | + def __init__(self, name: str, predicate: Callable[["Blackboard"], bool]) -> None: |
| 133 | + self.name = name |
| 134 | + self.predicate = predicate |
| 135 | + |
| 136 | + def tick(self, bb: "Blackboard") -> Status: |
| 137 | + return Status.SUCCESS if self.predicate(bb) else Status.FAILURE |
| 138 | + |
| 139 | + |
| 140 | +class Action(Node): |
| 141 | + """A leaf that proposes one environment action and reports RUNNING. |
| 142 | +
|
| 143 | + The environment executes the proposed action after the tick, so a leaf is |
| 144 | + "running" for exactly one control step; the outcome is observed on the next |
| 145 | + tick through the conditions above it. |
| 146 | + """ |
| 147 | + |
| 148 | + def __init__(self, name: str, propose: Callable[["Blackboard"], dict[str, Any]]) -> None: |
| 149 | + self.name = name |
| 150 | + self.propose = propose |
| 151 | + |
| 152 | + def tick(self, bb: "Blackboard") -> Status: |
| 153 | + bb.action = self.propose(bb) |
| 154 | + bb.active_leaf = self.name |
| 155 | + return Status.RUNNING |
| 156 | + |
| 157 | + |
| 158 | +class Blackboard: |
| 159 | + """Shared memory the tree reads and writes (py_trees-style).""" |
| 160 | + |
| 161 | + def __init__(self) -> None: |
| 162 | + self.obs: dict[str, Any] = {} |
| 163 | + self.action: dict[str, Any] | None = None |
| 164 | + self.active_leaf: str = "" |
| 165 | + |
| 166 | + |
| 167 | +# --- The agent: belief tracking + the tree that decides what to do ---------- |
| 168 | + |
| 169 | + |
| 170 | +class BehaviorTreeAgent: |
| 171 | + """Tracks a spatial belief and lets a reactive BT choose look vs. grasp.""" |
| 172 | + |
| 173 | + confidence_radius = 0.085 |
| 174 | + |
| 175 | + def __init__(self) -> None: |
| 176 | + self.viewpoints = [ |
| 177 | + np.array([0.84, 0.52]), # right of the occluder -> object visible |
| 178 | + np.array([0.78, 0.22]), |
| 179 | + np.array([0.20, 0.84]), # left but above the occluder -> still visible |
| 180 | + ] |
| 181 | + self.tree = self._build_tree() |
| 182 | + self.reset() |
| 183 | + |
| 184 | + def reset(self) -> None: |
| 185 | + self.belief_mean: np.ndarray | None = None |
| 186 | + self.belief_radius = 0.14 |
| 187 | + self.look_count = 0 |
| 188 | + self.retry_count = 0 |
| 189 | + self._last_integrated_time: int | None = None |
| 190 | + |
| 191 | + # The tree is pure structure; all state lives on the agent / blackboard. |
| 192 | + def _build_tree(self) -> Node: |
| 193 | + confident_grasp = Sequence( |
| 194 | + "confident_grasp", |
| 195 | + [ |
| 196 | + Condition("belief_confident?", self._belief_confident), |
| 197 | + Action("grasp_at_belief", self._grasp_at_belief), |
| 198 | + ], |
| 199 | + ) |
| 200 | + grasp_or_recover = Fallback( |
| 201 | + "grasp_or_recover", |
| 202 | + [confident_grasp, Action("relook_to_refine", self._relook_to_refine)], |
| 203 | + ) |
| 204 | + return Fallback( |
| 205 | + "pick_the_block", |
| 206 | + [Condition("object_in_gripper?", self._object_in_gripper), grasp_or_recover], |
| 207 | + ) |
| 208 | + |
| 209 | + # --- conditions --- |
| 210 | + def _object_in_gripper(self, bb: Blackboard) -> bool: |
| 211 | + return (bb.obs.get("gripper") or {}).get("holding") is not None |
| 212 | + |
| 213 | + def _belief_confident(self, bb: Blackboard) -> bool: |
| 214 | + return self.belief_mean is not None and self.belief_radius <= self.confidence_radius |
| 215 | + |
| 216 | + # --- actions --- |
| 217 | + def _grasp_at_belief(self, bb: Blackboard) -> dict[str, Any]: |
| 218 | + assert self.belief_mean is not None |
| 219 | + return {"type": "pick", "position": np.clip(self.belief_mean, 0.0, 1.0)} |
| 220 | + |
| 221 | + def _relook_to_refine(self, bb: Blackboard) -> dict[str, Any]: |
| 222 | + target = self.viewpoints[self.look_count % len(self.viewpoints)] |
| 223 | + self.look_count += 1 |
| 224 | + return {"type": "look", "target": target} |
| 225 | + |
| 226 | + # --- closed-loop hooks --- |
| 227 | + def act(self, obs: dict[str, Any]) -> dict[str, Any]: |
| 228 | + self._integrate_observation(obs) |
| 229 | + self.last_bb = Blackboard() |
| 230 | + self.last_bb.obs = obs |
| 231 | + status = self.tree.tick(self.last_bb) |
| 232 | + self.last_status = status |
| 233 | + # The root only reports SUCCESS once the block is held; until then a leaf |
| 234 | + # action is always RUNNING, so bb.action is set. Guard defensively. |
| 235 | + if self.last_bb.action is None: |
| 236 | + return {"type": "noop"} |
| 237 | + return self.last_bb.action |
| 238 | + |
| 239 | + def update(self, obs: dict[str, Any], reward: float, info: dict[str, Any]) -> None: |
| 240 | + self._integrate_observation(obs) |
| 241 | + |
| 242 | + failure = info.get("failure") |
| 243 | + if isinstance(failure, Failure) and failure.kind == "grasp_miss": |
| 244 | + # A miss does not trigger a special branch: it just grows uncertainty |
| 245 | + # back above the confidence threshold, so `belief_confident?` fails and |
| 246 | + # the Fallback re-runs `relook_to_refine` before the next grasp. |
| 247 | + self.retry_count += 1 |
| 248 | + self.belief_radius = max(self.belief_radius, self.confidence_radius) + 0.04 |
| 249 | + info["retry_count"] = self.retry_count |
| 250 | + elif info.get("success"): |
| 251 | + self.belief_radius = 0.025 |
| 252 | + # Confirm the tree now short-circuits at `object_in_gripper?`, so the |
| 253 | + # root reports SUCCESS rather than the grasp leaf's RUNNING. |
| 254 | + confirm = Blackboard() |
| 255 | + confirm.obs = obs |
| 256 | + self.last_status = self.tree.tick(confirm) |
| 257 | + |
| 258 | + info["bt_status"] = getattr(self, "last_status", Status.RUNNING).value |
| 259 | + info["bt_leaf"] = getattr(self, "last_bb", Blackboard()).active_leaf |
| 260 | + info["belief_mean"] = None if self.belief_mean is None else self.belief_mean.copy() |
| 261 | + info["belief_radius"] = float(self.belief_radius) |
| 262 | + info["retry_count"] = int(self.retry_count) |
| 263 | + |
| 264 | + def _integrate_observation(self, obs: dict[str, Any]) -> None: |
| 265 | + obs_time = int(obs.get("time", -1)) |
| 266 | + if obs_time == self._last_integrated_time: |
| 267 | + return |
| 268 | + self._last_integrated_time = obs_time |
| 269 | + |
| 270 | + detections = obs.get("detections", []) |
| 271 | + if not detections: |
| 272 | + return |
| 273 | + |
| 274 | + position = np.asarray(detections[0]["position"], dtype=float) |
| 275 | + confidence = float(detections[0].get("confidence", 0.5)) |
| 276 | + if self.belief_mean is None: |
| 277 | + self.belief_mean = position.copy() |
| 278 | + else: |
| 279 | + alpha = float(np.clip(0.35 + 0.45 * confidence, 0.35, 0.80)) |
| 280 | + self.belief_mean = alpha * self.belief_mean + (1.0 - alpha) * position |
| 281 | + self.belief_radius = max(0.035, self.belief_radius * 0.72) |
| 282 | + |
| 283 | + |
| 284 | +def run(seed: int = 3, render: bool = True, max_steps: int = 40) -> Trace: |
| 285 | + env = Tabletop2D(seed=seed) |
| 286 | + obs = env.reset(seed=seed) |
| 287 | + agent = BehaviorTreeAgent() |
| 288 | + agent.reset() |
| 289 | + trace = Trace() |
| 290 | + |
| 291 | + for _ in range(max_steps): |
| 292 | + action = agent.act(obs) |
| 293 | + result = env.step(action) |
| 294 | + obs, reward, done, info = result.as_tuple() |
| 295 | + agent.update(obs, reward, info) |
| 296 | + trace.append(obs, action, reward, info) |
| 297 | + |
| 298 | + if render: |
| 299 | + env.render(agent=agent, info=info) |
| 300 | + |
| 301 | + if done: |
| 302 | + break |
| 303 | + |
| 304 | + if not (trace.infos and trace.infos[-1].get("success")): |
| 305 | + trace.infos[-1]["failure"] = trace.infos[-1].get("failure") or Failure( |
| 306 | + "tree_exhausted", "ran out of steps without lifting the object", False |
| 307 | + ) |
| 308 | + |
| 309 | + return trace |
| 310 | + |
| 311 | + |
| 312 | +def main() -> None: |
| 313 | + parser = argparse.ArgumentParser() |
| 314 | + parser.add_argument("--seed", type=int, default=3) |
| 315 | + parser.add_argument("--max-steps", type=int, default=40) |
| 316 | + parser.add_argument("--no-render", action="store_true") |
| 317 | + args = parser.parse_args() |
| 318 | + |
| 319 | + trace = run(seed=args.seed, render=not args.no_render, max_steps=args.max_steps) |
| 320 | + picked = bool(trace.infos and trace.infos[-1].get("success")) |
| 321 | + leaves = [info.get("bt_leaf", "") for info in trace.infos] |
| 322 | + print( |
| 323 | + f"picked={picked} steps={len(trace.actions)} " |
| 324 | + f"retries={trace.summary().retry_count} " |
| 325 | + f"relooks={leaves.count('relook_to_refine')} grasps={leaves.count('grasp_at_belief')}" |
| 326 | + ) |
| 327 | + |
| 328 | + if not args.no_render: |
| 329 | + import matplotlib.pyplot as plt |
| 330 | + |
| 331 | + plt.ioff() |
| 332 | + plt.show() |
| 333 | + |
| 334 | + |
| 335 | +if __name__ == "__main__": |
| 336 | + main() |
0 commit comments