Publish tiny_computer_agent/ from marimo

tiny_computer_agent/__init__.py

@@ -1,17 +1,17 @@
 """Tiny Computer Agent - an agentic harness that drives a real computer with a
 tiny, fully local model.
 
-The harness is deliberately boring: it owns the loop, the closed action
-vocabulary and the safety gates, so the model only ever answers one question -
-"which action, from this menu, comes next?"  Nothing here needs a network call
-at run time, and the core is standard-library only.
+The harness is deliberately boring: it owns the loop, the closed action vocabulary,
+perception and the safety gates, so the model only ever answers one question -
+"which item from this short menu, and which verb?"  Nothing here needs a network
+call at run time, and the core imports nothing outside the standard library.
 
-    from tiny_computer_agent import (
-        AgentLoop, DryRunExecutor, Element, FakeScreen, MockBrain, Policy,
-    )
+    from tiny_computer_agent import AgentLoop, GrabScreen, IndexedScreen, OcrDetector
+    from tiny_computer_agent.models import TextPlannerBrain
 
-    screen = FakeScreen([(Element("Save", x=320, y=120, width=80, height=24),)])
-    loop = AgentLoop(screen=screen, brain=MockBrain([{"kind": "click", "label": "Save"}]))
+    screen = IndexedScreen(GrabScreen(), OcrDetector())      # pixels -> labels
+    brain = TextPlannerBrain()                               # 0.5B, CPU-friendly
+    loop = AgentLoop(screen=screen, brain=brain)             # dry run by default
     episode = loop.run("save the document")
     print(episode.render())
 """
@@ -34,13 +34,28 @@ from tiny_computer_agent.actions import (
 from tiny_computer_agent.brain import (
     Brain,
     BrainError,
+    DecisionMenu,
     GoalUnreachable,
+    MenuItem,
     MockBrain,
+    ModelBrain,
     PlannerBrain,
-    VLMBrain,
     extract_json_object,
+    extract_json_objects,
+    parse_decision,
 )
 from tiny_computer_agent.loop import AgentLoop, Episode, Step
+from tiny_computer_agent.models import (
+    DEFAULT_TEXT_MODEL,
+    DEFAULT_VISION_MODEL,
+    MODEL_BRAINS,
+    TextPlannerBrain,
+    VLMBrain,
+    make_model_brain,
+    model_cached,
+    model_dependencies,
+    resolve_device,
+)
 from tiny_computer_agent.observe import (
     Element,
     FakeScreen,
@@ -48,16 +63,30 @@ from tiny_computer_agent.observe import (
     Observation,
     ScreenSource,
     ScreenUnavailable,
+    normalise_label,
+)
+from tiny_computer_agent.perception import (
+    Detector,
+    DetectorUnavailable,
+    IndexedScreen,
+    OcrDetector,
+    detector_available,
+    make_detector,
 )
 from tiny_computer_agent.safety import Policy, PolicyViolation
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = [
     "ACTION_KINDS", "ACTION_MENU", "Action", "ActionError", "ActionResult",
-    "AgentLoop", "Brain", "BrainError", "DryRunExecutor", "Element", "Episode",
-    "Executor", "FakeScreen", "GoalUnreachable", "GrabScreen", "InputUnavailable", "MockBrain",
-    "Observation", "PlannerBrain", "Policy", "PolicyViolation", "PyAutoGuiExecutor",
-    "ScreenSource", "ScreenUnavailable", "Step", "VLMBrain", "extract_json_object",
-    "parse_action", "resolve_point", "__version__",
+    "AgentLoop", "Brain", "BrainError", "DEFAULT_TEXT_MODEL", "DEFAULT_VISION_MODEL",
+    "DecisionMenu", "Detector", "DetectorUnavailable", "DryRunExecutor", "Element",
+    "Episode", "Executor", "FakeScreen", "GoalUnreachable", "GrabScreen",
+    "IndexedScreen", "InputUnavailable", "MODEL_BRAINS", "MenuItem", "MockBrain",
+    "ModelBrain", "Observation", "OcrDetector", "PlannerBrain", "Policy",
+    "PolicyViolation", "PyAutoGuiExecutor", "ScreenSource", "ScreenUnavailable",
+    "Step", "TextPlannerBrain", "VLMBrain", "detector_available",
+    "extract_json_object", "extract_json_objects", "make_detector", "make_model_brain",
+    "model_cached", "model_dependencies", "normalise_label", "parse_action",
+    "parse_decision", "resolve_device", "resolve_point", "__version__",
 ]

tiny_computer_agent/__main__.py

@@ -1,13 +1,15 @@
 """Command line entry point.
 
-    python -m tiny_computer_agent                        # scripted demo, dry-run
-    python -m tiny_computer_agent --live                 # your real screen, dry-run
-    python -m tiny_computer_agent "click 'Save'"        # planner brain
-    python -m tiny_computer_agent --script '[{"kind": "done"}]'
-
-Real mouse and keyboard input additionally requires ``--real-input`` *and*
-``TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT=1`` in the environment. That is deliberate:
-dry-run is the default so you can read a whole episode first.
+    python -m tiny_computer_agent                     # scripted demo, dry-run
+    python -m tiny_computer_agent --doctor            # what can this machine do?
+    python -m tiny_computer_agent --image example     # replay a screenshot, dry-run
+    python -m tiny_computer_agent --image shot.png    # replay your own screenshot
+    python -m tiny_computer_agent --live              # your real screen, dry-run
+    python -m tiny_computer_agent --live --real-input "click 'OK'"
+
+Real mouse and keyboard input additionally requires
+``TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT=1``. Dry run is the default, so you can read a
+whole episode before anything touches your pointer.
 """
 
 from __future__ import annotations
@@ -15,16 +17,24 @@ from __future__ import annotations
 import argparse
 import json
 import os
+import platform
 import sys
 
 from tiny_computer_agent.act import DryRunExecutor, PyAutoGuiExecutor
 from tiny_computer_agent.brain import MockBrain, PlannerBrain
+from tiny_computer_agent.demo import fake_desktop
+from tiny_computer_agent.desktop import (
+    BRAIN_NAMES, DEFAULT_KILL_SWITCH, Check, display_available, run_checks,
+)
 from tiny_computer_agent.loop import AgentLoop, Step
-from tiny_computer_agent.observe import Element, FakeScreen, GrabScreen
+from tiny_computer_agent.observe import (
+    Element, FakeScreen, GrabScreen, ImageScreen, ScreenUnavailable,
+)
+from tiny_computer_agent.perception import IndexedScreen, detector_available, make_detector
 from tiny_computer_agent.safety import Policy
 
 ALLOW_REAL_INPUT_ENV = "TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT"
-DEFAULT_KILL_SWITCH = "/tmp/tiny-computer-agent.STOP"
+EXAMPLE_IMAGE = "example"
 
 
 def demo_screen() -> FakeScreen:
@@ -50,24 +60,114 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--max-steps", type=int, default=6)
     parser.add_argument("--script", default=None,
                         help="JSON list of actions to replay instead of using a brain")
+    parser.add_argument("--doctor", action="store_true",
+                        help="report what this machine can and cannot do, then exit")
+    parser.add_argument("--image", default=None,
+                        help="replay a screenshot instead of the screen; "
+                             f"'{EXAMPLE_IMAGE}' uses the built-in one")
+    parser.add_argument("--brain", default="planner", choices=BRAIN_NAMES,
+                        help="planner (rule-based, no weights) | text-planner | vlm")
+    parser.add_argument("--detector", default="auto", choices=("auto", "ocr", "grid"),
+                        help="how to build the element index (auto = OCR, grid fallback)")
     parser.add_argument("--live", action="store_true",
                         help="capture your real screen instead of the scripted demo")
     parser.add_argument("--real-input", action="store_true",
                         help="actually move the mouse and type (requires the allow env var)")
+    parser.add_argument("--monitor", type=int, default=1, help="1 is the primary monitor")
+    parser.add_argument("--scale", type=float, default=None,
+                        help="screen pixels per captured pixel (HiDPI, secondary monitors)")
+    parser.add_argument("--no-ocr", action="store_true",
+                        help="skip OCR and index a numbered grid instead")
+    parser.add_argument("--max-width", type=int, default=1600,
+                        help="downscale captures wider than this before indexing")
+    parser.add_argument("--allow-grid", action="store_true",
+                        help="permit clicks on numbered grid cells with --real-input")
     parser.add_argument("--kill-switch", default=DEFAULT_KILL_SWITCH,
                         help="create this file to abort the run before the next action")
     return parser
 
 
+def report(checks: list[Check]) -> int:
+    """Print a preflight report; non-zero when something required is missing."""
+    blocking = [c for c in checks if c.required and not c.ok]
+    lines = [
+        f"tiny-computer-agent preflight ({platform.system()} {platform.machine()}, "
+        f"python {platform.python_version()})",
+        "",
+        *(c.render() for c in checks),
+        "",
+    ]
+    if blocking:
+        lines.append(f"NOT READY: {len(blocking)} blocking problem(s): "
+                     + ", ".join(c.name for c in blocking))
+    else:
+        lines.append("READY: the harness can drive this desktop.")
+    print("\n".join(lines))
+    return 1 if blocking else 0
+
+
+def _probe_live(args: argparse.Namespace) -> Check:
+    """Actually grab one frame, because only a real capture proves it works."""
+    try:
+        observation = GrabScreen(
+            monitor=args.monitor, max_width=args.max_width, scale=args.scale
+        ).capture()
+        frame = observation.frame
+        return Check("live capture", True, f"captured {frame.width}x{frame.height} pixels")
+    except ScreenUnavailable as exc:
+        return Check("live capture", False, str(exc),
+                     fix="run this on the desktop session itself, not a headless shell")
+
+
+def _build_screen(args: argparse.Namespace):
+    """(screen, description). Raises ScreenUnavailable / OSError on failure."""
+    detector_name = "grid" if args.no_ocr else args.detector
+    if args.image is not None:
+        source = ImageScreen(fake_desktop() if args.image == EXAMPLE_IMAGE else args.image)
+        return IndexedScreen(source, make_detector(detector_name)), f"screenshot replay: {args.image}"
+    if args.live:
+        source = GrabScreen(monitor=args.monitor, max_width=args.max_width, scale=args.scale)
+        return IndexedScreen(source, make_detector(detector_name)), f"real desktop (monitor {args.monitor})"
+    return demo_screen(), "scripted demo"
+
+
 def main(argv: list[str] | None = None) -> int:
     args = build_parser().parse_args(argv)
+    detector_name = "grid" if args.no_ocr else args.detector
+
+    if args.doctor:
+        checks = run_checks(brain=args.brain, detector=detector_name,
+                            real_input=args.real_input, kill_switch_path=args.kill_switch)
+        if args.live:
+            checks.append(_probe_live(args))
+        return report(checks)
 
     if args.real_input and os.environ.get(ALLOW_REAL_INPUT_ENV) != "1":
         print(f"refusing --real-input: set {ALLOW_REAL_INPUT_ENV}=1 to confirm you "
               f"understand this will drive your actual mouse and keyboard.", file=sys.stderr)
         return 2
 
-    screen = GrabScreen() if args.live else demo_screen()
+    if args.brain == "vlm" and not (args.live or args.image):
+        print("the vlm brain needs --live (or --image) so it has pixels to look at; "
+              "use --brain planner to exercise the harness offline.", file=sys.stderr)
+        return 2
+
+    if args.live:
+        available, detail = display_available()
+        if not available:
+            print(f"live capture needs a display session: {detail}")
+            print(f"refusing to run live: {detail}. Try --image to replay a screenshot, "
+                  f"or --doctor to diagnose.", file=sys.stderr)
+            return 2
+
+    try:
+        screen, description = _build_screen(args)
+    except ScreenUnavailable as exc:
+        print(f"cannot see the screen: {exc}", file=sys.stderr)
+        return 3
+    except Exception as exc:
+        print(f"could not start the screen: {type(exc).__name__}: {exc}", file=sys.stderr)
+        return 2
 
     if args.script:
         try:
@@ -76,26 +176,64 @@ def main(argv: list[str] | None = None) -> int:
             print(f"--script is not valid JSON: {exc}", file=sys.stderr)
             return 2
         brain = MockBrain(script)
+    elif args.brain == "planner":
+        brain = PlannerBrain()
     else:
-        brain = PlannerBrain(done_when_text_contains=None)
+        from tiny_computer_agent.models import make_model_brain
+
+        try:
+            brain = make_model_brain(args.brain)
+        except Exception as exc:
+            print(f"could not load the {args.brain} brain: {type(exc).__name__}: {exc}",
+                  file=sys.stderr)
+            return 2
 
     policy = Policy(
         dry_run=not args.real_input,
         step_budget=args.max_steps,
         kill_switch_path=args.kill_switch,
+        allow_grid_clicks=args.allow_grid,
     )
     executor = PyAutoGuiExecutor() if args.real_input else DryRunExecutor()
 
     mode = "REAL INPUT (mouse and keyboard live)" if args.real_input else "dry-run"
-    print(f"goal: {args.goal}\nmode: {mode}  |  screen: "
-          f"{'real desktop' if args.live else 'scripted demo'}  |  "
-          f"kill switch: create {args.kill_switch} to stop\n")
+    print(f"goal: {args.goal}")
+    print(f"mode: {mode}  |  screen: {description}  |  brain: {args.brain}  |  "
+          f"kill switch: create {args.kill_switch} to stop")
+    if args.image is not None or args.live:
+        _, index_detail = detector_available(detector_name)
+        print(f"index: {detector_name} - {index_detail}")
+    if args.image is not None:
+        # One read, so the header can state how big the action space is.
+        try:
+            preview = screen.capture()
+        except ScreenUnavailable as exc:
+            print(f"cannot see the screen: {exc}", file=sys.stderr)
+            return 3
+        print(f"elements detected: {len(preview.elements)}")
+    print()
+
+    warned = False
 
     def on_step(step: Step) -> None:
+        nonlocal warned
+        elements = step.observation.elements
+        if not warned and elements and all(e.kind == "grid" for e in elements):
+            warned = True
+            _, index_detail = detector_available(detector_name)
+            print(f"warning: no text was found on screen, so the index is "
+                  f"{len(elements)} numbered grid cells and the agent is choosing "
+                  f"blind ({index_detail})", file=sys.stderr)
         print("  " + step.line())
 
-    episode = AgentLoop(screen=screen, brain=brain, executor=executor,
-                        policy=policy, on_step=on_step).run(args.goal)
+    try:
+        episode = AgentLoop(screen=screen, brain=brain, executor=executor,
+                            policy=policy, on_step=on_step).run(args.goal)
+    except ScreenUnavailable as exc:
+        print(f"cannot see the screen: {exc}", file=sys.stderr)
+        print("run `python -m tiny_computer_agent --doctor --live` for details.", file=sys.stderr)
+        return 3
+
     print()
     print(episode.render())
     return 0 if episode.success else 1

tiny_computer_agent/brain.py

@@ -1,23 +1,29 @@
-"""Decision making.
+"""Decision making: menu-mediated, model-agnostic.
 
-Three brains ship here, and only one of them is a model:
+The design lesson from testing tiny models on real screenshots: a 256M VLM can
+*read* a button label from pixels but cannot reliably reason, and a 0.5B text model
+writes perfect JSON but picks the wrong widget when handed noisy free-form text.
+Both do much better when the harness hands them a **short numbered menu** and asks
+for a number. That is what this module builds, and it keeps the model's output space
+finite no matter how small the model is.
+
+Brains shipped here:
 
 * :class:`MockBrain`    - replays a scripted action list (tests, demos, replays)
-* :class:`PlannerBrain` - rule-based: clicks the element whose label matches the
-  goal. No model, no download, yet it solves real "click the thing labelled X"
-  tasks, because the *element index* does the grounding a VLM would otherwise do
-* :class:`VLMBrain`     - the seam for a real tiny VLM. Prompt building and reply
-  validation are done; you implement ``_generate()``
+* :class:`PlannerBrain` - rule-based, no model: clicks the element matching the goal
+* :class:`ModelBrain`   - base for real models: menu, prompt, lenient parsing
+  (see :mod:`tiny_computer_agent.models` for the text and vision backends)
 """
 
 from __future__ import annotations
 
 import json
 import re
-from typing import Any, Mapping, Protocol, Sequence
+from dataclasses import dataclass
+from typing import Any, Iterator, Mapping, Protocol, Sequence
 
-from tiny_computer_agent.actions import ACTION_MENU, Action, ActionError, parse_action
-from tiny_computer_agent.observe import Observation
+from tiny_computer_agent.actions import ACTION_KINDS, Action, ActionError, parse_action
+from tiny_computer_agent.observe import Element, Observation, normalise_label
 
 
 class BrainError(RuntimeError):
@@ -33,19 +39,217 @@ class GoalUnreachable(BrainError):
     """
 
 
+def extract_json_objects(raw: str) -> list[dict[str, Any]]:
+    """Every flat JSON object in a reply, in the order they appear.
+
+    Regex-based and non-nesting: replies from tiny models are one-level objects.
+    Callers should prefer the **last** object, because a model that echoes its own
+    prompt before answering puts the prompt's example first and its answer last.
+    """
+    objects: list[dict[str, Any]] = []
+    for match in re.finditer(r"\{[^{}]*\}", raw or "", re.DOTALL):
+        try:
+            payload = json.loads(match.group(0))
+        except json.JSONDecodeError:
+            continue
+        if isinstance(payload, dict):
+            objects.append(payload)
+    return objects
+
+
 def extract_json_object(raw: str) -> dict[str, Any]:
-    """Pull the first JSON object out of a reply, ignoring any surrounding prose."""
-    start = raw.find("{")
-    end = raw.rfind("}")
-    if start == -1 or end <= start:
-        raise BrainError(f"no JSON object in the model reply: {raw[:200]!r}")
-    try:
-        payload = json.loads(raw[start : end + 1])
-    except json.JSONDecodeError as exc:
-        raise BrainError(f"malformed JSON in the model reply: {exc}") from exc
-    if not isinstance(payload, dict):
-        raise BrainError(f"expected a JSON object, got {type(payload).__name__}")
-    return payload
+    """The decision out of a reply: the *last* JSON object, not the first."""
+    objects = extract_json_objects(raw)
+    if not objects:
+        raise BrainError(f"no JSON object in the model reply: {(raw or '')[:200]!r}")
+    return objects[-1]
+
+
+@dataclass(frozen=True)
+class MenuItem:
+    """One numbered choice shown to the model."""
+
+    index: int
+    element: Element
+
+    @property
+    def label(self) -> str:
+        return self.element.label
+
+
+class DecisionMenu:
+    """A short, numbered, closed choice set built from the element index.
+
+    Elements are presented in reading order (top to bottom, then left to right),
+    and capped: a 20-item menu is as good as free-form text to a small model.
+    """
+
+    def __init__(self, observation: Observation, *, max_items: int = 12) -> None:
+        ordered = sorted(observation.elements, key=lambda e: (e.y, e.x))[:max_items]
+        self.items: tuple[MenuItem, ...] = tuple(
+            MenuItem(index=i + 1, element=e) for i, e in enumerate(ordered)
+        )
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __bool__(self) -> bool:
+        return bool(self.items)
+
+    def __iter__(self) -> Iterator[MenuItem]:
+        return iter(self.items)
+
+    def item(self, value: Any) -> MenuItem | None:
+        """Look up a 1-based index, tolerating the strings models actually send."""
+        try:
+            number = int(float(str(value).strip()))
+        except (TypeError, ValueError):
+            return None
+        if 1 <= number <= len(self.items):
+            return self.items[number - 1]
+        return None
+
+    def find(self, label: str) -> Element | None:
+        """Resolve a label against the menu (normalised, like Observation.find)."""
+        wanted = normalise_label(label)
+        if not wanted:
+            return None
+        for item in self.items:
+            if normalise_label(item.label) == wanted:
+                return item.element
+        for item in self.items:
+            if wanted in normalise_label(item.label):
+                return item.element
+        return None
+
+    def mentioned_in(self, text: str) -> tuple[MenuItem, ...]:
+        """Menu entries whose label appears in the text, in order of first mention."""
+        haystack = normalise_label(text)
+        found: list[tuple[int, MenuItem]] = []
+        for item in self.items:
+            needle = normalise_label(item.label)
+            if needle and (position := haystack.find(needle)) != -1:
+                found.append((position, item))
+        return tuple(item for _, item in sorted(found, key=lambda pair: pair[0]))
+
+    def render(self) -> str:
+        return "\n".join(
+            f"  {item.index}. {item.label}  (at {item.element.center[0]},{item.element.center[1]})"
+            for item in self.items
+        )
+
+
+def _recover_from_text(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None:
+    """Salvage a decision from a payload that named the right thing in the wrong field.
+
+    Observed for real: a 0.5B planner replied ``{"kind": "click", "index": 0,
+    "reason": "Save"}`` - the correct answer, in the wrong place. If the index is
+    unusable but exactly one menu label is mentioned anywhere in the payload, take
+    the label. It is still only ever a label the detector saw in this frame, so the
+    deny list, the dry run and the step budget all continue to apply.
+    """
+    text = json.dumps(payload, ensure_ascii=False)
+    mentioned = menu.mentioned_in(text)
+    if len({item.label for item in mentioned}) == 1:
+        item = mentioned[0]
+        return Action(
+            "click", label=item.label,
+            reason=f"recovered {item.label!r} from the model's reply: {text[:80]}",
+        )
+    return None
+
+
+def _action_from_payload(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None:
+    """Interpret one JSON object, tolerating the ways small models mangle the shape."""
+    kind = payload.get("kind")
+    if isinstance(kind, str) and kind.strip().lower() not in ACTION_KINDS:
+        kind = None  # e.g. {"kind": "file"} - a garbled click; fall back to index/label
+
+    if payload.get("index") is not None:
+        item = menu.item(payload["index"])
+        if item is not None:
+            return Action("click", label=item.label, reason=f"menu index {payload['index']}")
+        recovered = _recover_from_text(payload, menu)
+        if recovered is not None:
+            return recovered
+        raise BrainError(
+            f"the model chose index {payload['index']!r}, which is not in the "
+            f"menu of {len(menu)} items"
+        )
+
+    if kind in (None, "click") and "label" in payload:
+        element = menu.find(str(payload["label"]))
+        if element is None:
+            raise BrainError(f"the model named {payload['label']!r}, which is not on screen")
+        return Action("click", label=element.label, reason=f"the model named {element.label!r}")
+
+    if kind is not None:
+        try:
+            return parse_action(dict(payload, kind=kind))
+        except ActionError as exc:
+            # a valid-looking verb with no usable target: maybe the label is in the prose
+            return _recover_from_text(payload, menu) or (_raise_brain_error(exc))
+
+    return _recover_from_text(payload, menu)
+
+
+def _raise_brain_error(exc: Exception) -> Action:
+    raise BrainError(f"the model proposed an invalid action: {exc}") from exc
+
+
+_DONE = re.compile(r"\b(done|finished|complete[d]?)\b", re.IGNORECASE)
+
+
+def parse_decision(raw: str, observation: Observation, *, max_menu_items: int = 12) -> Action:
+    """Turn a small model's reply into a validated Action - leniently, but safely.
+
+    Tried in order:
+
+    1. the **last** JSON object in the reply (``index``, ``label`` or a full action)
+    2. a label the reply names (most specific evidence of intent)
+    3. a bare menu number
+    4. a "done" word
+
+    A reply that mentions three or more menu entries is rejected as an echo of the
+    prompt rather than a decision - the failure mode of the weakest models, and the
+    one most likely to make an agent click something arbitrary.
+    """
+    text = " ".join((raw or "").split())
+    if not text:
+        raise BrainError("the model returned an empty reply")
+    menu = DecisionMenu(observation, max_items=max_menu_items)
+
+    for payload in reversed(extract_json_objects(text)):
+        action = _action_from_payload(payload, menu)
+        if action is not None:
+            return action
+
+    mentioned = menu.mentioned_in(text)
+    if len(mentioned) >= 3:
+        raise BrainError(
+            f"the reply mentions {len(mentioned)} menu entries, so it looks like an echo "
+            f"of the prompt rather than a decision: {text[:120]!r}"
+        )
+    if mentioned:
+        item = mentioned[0]
+        note = f"the reply names {item.label!r}"
+        if len(mentioned) > 1:
+            note += f" (ambiguous: {len(mentioned)} labels mentioned; took the first)"
+        return Action("click", label=item.label, reason=note)
+
+    number = re.search(r"\b(\d{1,2})\b", text)
+    if number is not None:
+        item = menu.item(number.group(1))
+        if item is not None:
+            return Action(
+                "click", label=item.label,
+                reason=f"menu index {number.group(1)} from reply {text[:40]!r}",
+            )
+
+    if _DONE.search(text):
+        return Action("done", reason=_DONE.search(text).group(0).lower())
+
+    raise BrainError(f"could not read a decision out of the reply: {text[:120]!r}")
 
 
 class Brain(Protocol):
@@ -82,11 +286,11 @@ class MockBrain:
 class PlannerBrain:
     """Rule-based baseline: click the element matching the goal, then stop.
 
-    Its purpose is twofold - it proves the harness works end to end with no model
-    at all, and it shows how much of a "computer use" task is really an indexing
-    problem. It waits when the target is not on screen yet, refuses to click the
-    same label twice, and gives up rather than flailing. It will not generalise
-    the way a VLM can.
+    Its purpose is threefold - it proves the harness works end to end with no model
+    at all, it is a fallback when a model download is undesirable, and it shows how
+    much of a "computer use" task is really an indexing problem. It waits when the
+    target is not on screen yet, refuses to click the same label twice, and gives up
+    rather than flailing. It will not generalise the way a model can.
     """
 
     name = "planner"
@@ -124,7 +328,6 @@ class PlannerBrain:
                           reason=f"matched {target!r} in the element index")
         raise GoalUnreachable(f"no element name found in the goal {goal!r}")
 
-
     @staticmethod
     def _target_from(goal: str) -> str | None:
         """Best effort: quoted text wins, else the last capitalised word."""
@@ -135,40 +338,97 @@ class PlannerBrain:
         return words[-1].strip(".,!?") if words else None
 
 
-class VLMBrain:
-    """Seam for a real (small) vision-language model.
+class ModelBrain:
+    """Shared plumbing for model-backed brains: menu, prompt, lenient parsing.
 
-    ``build_prompt`` and reply validation are implemented; subclass this and
-    implement ``_generate()`` to return the raw model text. Swap the model by
-    changing ``model_id`` - nothing else in the harness cares which weights you use.
+    Subclasses implement :meth:`_generate`. ``requires_index`` marks brains that
+    cannot decide without an element index (a text-only planner): for those, an
+    empty screen is a hard stop rather than a prompt.
     """
 
-    name = "vlm"
+    name = "model"
+    requires_index = False
 
-    def __init__(self, model_id: str = "vikhyatk/moondream2", *,
-                 max_new_tokens: int = 128) -> None:
+    def __init__(self, model_id: str, *, max_menu_items: int = 12,
+                 max_new_tokens: int = 64) -> None:
         self.model_id = model_id
+        self.max_menu_items = max_menu_items
         self.max_new_tokens = max_new_tokens
+        self.last_prompt: str | None = None
+        self.last_reply: str | None = None
+        # "what I already tried", newest last; each entry's outcome is filled in at
+        # the next decision, when the new observation can be compared to the old one
+        self.history: list[str] = []
+        self._last_fingerprint: tuple | None = None
+
+    @staticmethod
+    def _fingerprint(observation: Observation) -> tuple:
+        return (tuple(element.label for element in observation.elements),
+                observation.text[:200])
+
+    def build_prompt(self, goal: str, menu: DecisionMenu, observation: Observation,
+                     *, history: Sequence[str] = (), screen_changed: bool = True) -> str:
+        """A short prompt with a strict output shape.
+
+        Deliberately contains placeholders (``<number>``) rather than worked
+        examples: the weakest models echo the prompt back, and an echoed
+        placeholder fails validation instead of being mistaken for a decision.
+        """
+        lines = [
+            "You are controlling a computer. Choose exactly one action.",
+            "",
+            "Reply with one JSON object and nothing else, in this shape:",
+            '  {"kind": "click", "index": <number from the list>, "reason": "<brief>"}',
+            '  {"kind": "done", "reason": "<brief>"}',
+            "Use done only when the goal is already met.",
+            "",
+            f"Goal: {goal}",
+        ]
+        if observation.text:
+            lines += ["", f"Text on screen: {observation.text[:300]}"]
+        lines += ["", f"Clickable items ({len(menu)}):"]
+        lines.append(menu.render() if menu else "  (none detected)")
+        if history:
+            lines += ["", "Already attempted - do NOT repeat these, they did not work:"]
+            lines += [f"  - {entry}" for entry in history]
+        if not screen_changed:
+            lines += ["", "The screen did not change after your last action. "
+                          "Choose a DIFFERENT action."]
+        return "\n".join(lines)
 
-    def build_prompt(self, goal: str, observation: Observation) -> str:
-        return (
-            "You are controlling a computer. Choose the single next action.\n\n"
-            f"Goal: {goal}\n\n"
-            f"{observation.summary()}\n\n"
-            f"{ACTION_MENU}"
+    def decide(self, goal: str, observation: Observation) -> Action:
+        menu = DecisionMenu(observation, max_items=self.max_menu_items)
+        if self.requires_index and not menu:
+            raise GoalUnreachable(
+                "no clickable elements on screen - is screen capture and OCR working?"
+            )
+        fingerprint = self._fingerprint(observation)
+        screen_changed = (self._last_fingerprint is None
+                          or fingerprint != self._last_fingerprint)
+
+        # Finalise the previous attempt first: only now do we know whether it changed
+        # anything, and "you already tried this and nothing happened" is the single most
+        # useful thing to tell a small model.
+        if self.history:
+            previous = self.history[-1].split(" -> ")[0]
+            self.history[-1] = (
+                f"{previous} -> "
+                + ("the screen changed" if screen_changed else "the screen did NOT change")
+            )
+
+        self.last_prompt = self.build_prompt(
+            goal, menu, observation, history=tuple(self.history),
+            screen_changed=screen_changed,
         )
+        self.last_reply = self._generate(self.last_prompt, observation)
+        action = parse_decision(self.last_reply, observation,
+                                max_menu_items=self.max_menu_items)
+        self.history.append(action.describe())
+        del self.history[:-3]
+        self._last_fingerprint = fingerprint
+        return action
 
     def _generate(self, prompt: str, observation: Observation) -> str:
         raise NotImplementedError(
-            "No vision backend is wired up yet. Subclass VLMBrain and implement "
-            "_generate(prompt, observation) to return raw model text, or use "
-            "MockBrain / PlannerBrain while the harness is being tested."
+            "subclass ModelBrain and implement _generate(prompt, observation)"
         )
-
-    def decide(self, goal: str, observation: Observation) -> Action:
-        raw = self._generate(self.build_prompt(goal, observation), observation)
-        payload = extract_json_object(raw)
-        try:
-            return parse_action(payload)
-        except ActionError as exc:
-            raise BrainError(f"model proposed an invalid action: {exc}") from exc

tiny_computer_agent/demo.py

@@ -0,0 +1,60 @@
+"""Synthetic screens for offline demos and tests.
+
+No display required: this module draws a small fake desktop - a top bar with Files
+and Terminal, plus a Save dialog with a file-name field and Save/Cancel buttons - so
+the whole pipeline (pixels, OCR, index, menu, model, action) can be exercised on a
+headless machine, in CI, or before you let an agent near your real desktop.
+"""
+
+from __future__ import annotations
+
+import glob
+from pathlib import Path
+from typing import Any
+
+
+def _font(size: int) -> Any:
+    from PIL import ImageFont
+
+    candidates = [
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
+        "/System/Library/Fonts/Supplemental/Arial.ttf",
+        "C:/Windows/Fonts/arial.ttf",
+    ] + glob.glob("/usr/share/fonts/**/*Regular*.ttf", recursive=True)
+    for path in candidates:
+        try:
+            return ImageFont.truetype(path, size)
+        except Exception:
+            continue
+    return ImageFont.load_default()
+
+
+def fake_desktop(width: int = 720, height: int = 460) -> Any:
+    """Draw a screenshot-like image: a desktop with an unsaved Save dialog."""
+    from PIL import Image, ImageDraw
+
+    image = Image.new("RGB", (width, height), (245, 246, 248))
+    draw = ImageDraw.Draw(image)
+    draw.rectangle([0, 0, width, 28], fill=(38, 40, 48))
+    draw.text((12, 7), "Files", fill="white", font=_font(15))
+    draw.text((90, 7), "Terminal", fill="white", font=_font(15))
+    draw.rectangle([180, 90, 540, 330], fill="white", outline=(120, 120, 130), width=2)
+    draw.text((200, 108), "Untitled document - unsaved changes",
+              fill=(30, 30, 30), font=_font(16))
+    draw.text((200, 170), "File name", fill=(60, 60, 60), font=_font(15))
+    draw.rectangle([200, 195, 520, 225], fill=(252, 252, 252),
+                   outline=(120, 120, 130), width=2)
+    draw.rectangle([200, 265, 380, 300], fill=(66, 111, 220),
+                   outline=(40, 70, 160), width=2)
+    draw.text((275, 274), "Save", fill="white", font=_font(16))
+    draw.rectangle([400, 265, 520, 300], fill=(230, 230, 235),
+                   outline=(120, 120, 130), width=2)
+    draw.text((440, 274), "Cancel", fill=(30, 30, 30), font=_font(16))
+    return image
+
+
+def write_fake_desktop(path: str | Path) -> Path:
+    """Save :func:`fake_desktop` to a PNG (handy for ``--image``)."""
+    target = Path(path)
+    fake_desktop().save(target)
+    return target

tiny_computer_agent/desktop.py

@@ -0,0 +1,197 @@
+"""Wiring for a real desktop: preflight checks, then one-call agent construction.
+
+This module exists because "make it work on my computer" is mostly a diagnosis
+problem. Rather than failing deep inside an episode, :func:`doctor` reports exactly
+which piece is missing - display, capture backend, input backend, OCR engine, model
+runtime - and what to do about each.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import platform
+from dataclasses import dataclass
+from typing import Any, Callable, Iterable, Sequence
+
+from tiny_computer_agent.act import DryRunExecutor, Executor, PyAutoGuiExecutor
+from tiny_computer_agent.brain import Brain, PlannerBrain
+from tiny_computer_agent.loop import AgentLoop, Step
+from tiny_computer_agent.observe import GrabScreen, ScreenSource
+from tiny_computer_agent.perception import (
+    Detector, IndexedScreen, detector_available, make_detector,
+)
+from tiny_computer_agent.safety import Policy
+
+DEFAULT_KILL_SWITCH = "/tmp/tiny-computer-agent.STOP"
+BRAIN_NAMES = ("planner", "text-planner", "vlm")
+
+
+@dataclass(frozen=True)
+class Check:
+    """One preflight result: what was checked, whether it works, and how to fix it."""
+
+    name: str
+    ok: bool
+    detail: str
+    fix: str = ""
+    required: bool = True
+
+    def render(self) -> str:
+        marker = "ok  " if self.ok else ("FAIL" if self.required else "warn")
+        line = f"  [{marker}] {self.name}: {self.detail}"
+        if not self.ok and self.fix:
+            line += f"\n         fix: {self.fix}"
+        return line
+
+
+def _installed(module: str) -> bool:
+    return importlib.util.find_spec(module) is not None
+
+
+def display_available() -> tuple[bool, str]:
+    """Whether a windowing session looks reachable from this process."""
+    system = platform.system()
+    if system in ("Darwin", "Windows"):
+        return True, f"{system} desktop session assumed"
+    display = os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")
+    if display:
+        return True, f"DISPLAY={display}"
+    return False, "no DISPLAY or WAYLAND_DISPLAY set (headless container or SSH session)"
+
+
+def run_checks(*, brain: str = "planner", model_id: str | None = None,
+               detector: str = "ocr", real_input: bool = False,
+               kill_switch_path: str | None = DEFAULT_KILL_SWITCH) -> list[Check]:
+    """Every precondition for driving a real desktop, in one list."""
+    checks: list[Check] = []
+
+    ok, detail = display_available()
+    checks.append(Check(
+        "display server", ok, detail,
+        fix="run this on the desktop session itself (a headless server cannot be driven)",
+    ))
+
+    checks.append(Check(
+        "pillow", _installed("PIL"),
+        "installed" if _installed("PIL") else "not installed",
+        fix="pip install 'tiny-computer-agent[desktop]'",
+    ))
+
+    has_mss = _installed("mss")
+    checks.append(Check(
+        "screen capture", has_mss or _installed("PIL"),
+        "mss (fast, multi-monitor)" if has_mss else "pillow ImageGrab fallback",
+        fix="pip install mss for faster multi-monitor capture",
+        required=False,
+    ))
+
+    has_pyautogui = _installed("pyautogui")
+    checks.append(Check(
+        "input backend (pyautogui)", has_pyautogui,
+        "installed" if has_pyautogui else "not installed",
+        fix="pip install 'tiny-computer-agent[desktop]'",
+        required=real_input,
+    ))
+
+    detector_ok, detector_detail = detector_available(detector)
+    checks.append(Check(
+        f"element detector ({detector})", detector_ok, detector_detail,
+        fix="pip install 'tiny-computer-agent[ocr]'",
+    ))
+
+    if brain == "planner":
+        checks.append(Check(
+            "model runtime", True,
+            "not needed: the planner is rule-based (no weights, no download)",
+            required=False,
+        ))
+    else:
+        from tiny_computer_agent.models import (
+            DEFAULT_TEXT_MODEL, DEFAULT_VISION_MODEL, model_cached, model_dependencies,
+        )
+
+        deps_ok, deps_detail = model_dependencies()
+        model = model_id or (DEFAULT_TEXT_MODEL if brain == "text-planner" else DEFAULT_VISION_MODEL)
+        checks.append(Check(
+            "model runtime", deps_ok, deps_detail,
+            fix="pip install 'tiny-computer-agent[vision]'",
+        ))
+        checks.append(Check(
+            f"model weights ({model})",
+            model_cached(model),
+            "cached locally" if model_cached(model) else "not downloaded yet",
+            fix=f"first run will download {model} from the Hub (one time)",
+            required=False,
+        ))
+
+    checks.append(Check(
+        "execution mode",
+        True,
+        "REAL mouse and keyboard input" if real_input else "dry run (nothing is clicked)",
+        fix="pass --real-input AND set TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT=1",
+        required=False,
+    ))
+
+    checks.append(Check(
+        "kill switch", True,
+        f"create {kill_switch_path} to abort the next action"
+        if kill_switch_path else "disabled",
+        required=False,
+    ))
+    return checks
+
+
+def doctor(*, brain: str = "planner", model_id: str | None = None,
+           detector: str = "ocr", real_input: bool = False,
+           kill_switch_path: str | None = DEFAULT_KILL_SWITCH) -> str:
+    """A human-readable preflight report."""
+    checks = run_checks(brain=brain, model_id=model_id, detector=detector,
+                        real_input=real_input, kill_switch_path=kill_switch_path)
+    blocking = [c for c in checks if c.required and not c.ok]
+    lines = [
+        f"tiny-computer-agent preflight ({platform.system()} {platform.machine()}, python {platform.python_version()})",
+        "",
+        *(c.render() for c in checks),
+        "",
+    ]
+    if blocking:
+        lines.append(f"NOT READY: {len(blocking)} blocking problem(s): "
+                     + ", ".join(c.name for c in blocking))
+    else:
+        lines.append("READY: the harness can drive this desktop.")
+        if not real_input:
+            lines.append("Dry run is on: nothing will be clicked. Read the plan first.")
+    return "\n".join(lines)
+
+
+def build_agent(*, brain: Brain | str | None = None, screen: ScreenSource | None = None,
+                detector: Detector | str | None = None, dry_run: bool = True,
+                max_steps: int = 10, kill_switch_path: str | None = DEFAULT_KILL_SWITCH,
+                allow_grid_clicks: bool = False,
+                on_step: Callable[[Step], None] | None = None,
+                detector_options: dict[str, Any] | None = None) -> AgentLoop:
+    """Wire a real desktop into an :class:`AgentLoop` in one call.
+
+    Defaults are the safe ones: dry run, a step budget of 10, and a kill-switch file
+    on disk that aborts the next action if it appears. ``screen``/``detector``/
+    ``brain`` can be injected, which is how the tests drive this without a desktop.
+    """
+    if detector is None or isinstance(detector, str):
+        detector = make_detector(detector or "ocr", **(detector_options or {}))
+    if screen is None:
+        screen = IndexedScreen(GrabScreen(), detector)
+    if brain is None or isinstance(brain, str):
+        name = brain or "planner"
+        if name == "planner":
+            brain = PlannerBrain()
+        else:
+            from tiny_computer_agent.models import make_model_brain
+
+            brain = make_model_brain(name)
+    executor: Executor = DryRunExecutor() if dry_run else PyAutoGuiExecutor()
+    policy = Policy(dry_run=dry_run, step_budget=max_steps,
+                    kill_switch_path=kill_switch_path,
+                    allow_grid_clicks=allow_grid_clicks)
+    return AgentLoop(screen=screen, brain=brain, executor=executor, policy=policy,
+                     on_step=on_step)

tiny_computer_agent/loop.py

@@ -74,14 +74,16 @@ class AgentLoop:
 
     * every action is one of six validated verbs
     * nothing executes until policy has passed it
-    * the episode ends on a step budget, a kill switch, or a brain error
+    * the episode ends on a step budget, a kill switch, a brain error, or a
+      failed screen capture (reported as a step, never as a traceback)
     * the screen is re-read after every action, so the next decision sees the
       consequence of the last one
     """
 
     def __init__(self, *, screen: ScreenSource, brain: Brain,
                  executor: Executor | None = None, policy: Policy | None = None,
-                 on_step: Callable[[Step], None] | None = None) -> None:
+                 on_step: Callable[[Step], None] | None = None,
+                 max_repeats: int = 1) -> None:
         self.screen = screen
         self.brain = brain
         self.policy = policy or Policy()
@@ -89,11 +91,21 @@ class AgentLoop:
             executor = DryRunExecutor() if self.policy.dry_run else PyAutoGuiExecutor()
         self.executor = executor
         self.on_step = on_step
+        self.max_repeats = max_repeats
 
     def run(self, goal: str, *, max_steps: int | None = None) -> Episode:
         budget = max_steps or self.policy.step_budget
         episode = Episode(goal=goal, stop_reason=f"step budget of {budget} exhausted")
-        observation = self.screen.capture()
+        # A small model that repeats an action against an unchanged screen is stuck,
+        # not persistent. Tolerate `max_repeats` such repeats, then stop - burning the
+        # whole budget on the same click teaches the user nothing.
+        previous: tuple[tuple, str] | None = None
+        repeats = 0
+        try:
+            observation = self.screen.capture()
+        except Exception as exc:
+            episode.stop_reason = f"capture failed: {type(exc).__name__}: {exc}"
+            return episode
 
         for index in range(budget):
             if self.policy.should_stop():
@@ -129,6 +141,21 @@ class AgentLoop:
                 episode.stop_reason = f"policy refusal: {exc}"
                 break
 
+            signature = (self._fingerprint(observation), action.describe())
+            repeats = repeats + 1 if signature == previous else 0
+            previous = signature
+            if repeats > self.max_repeats:
+                self._record(episode, Step(
+                    index, observation, action=action,
+                    error=f"no progress: repeating {action.describe()} "
+                          f"against an unchanged screen",
+                ))
+                episode.stop_reason = (
+                    f"no progress after {repeats} identical actions on an unchanged "
+                    f"screen ({action.kind} {action.label or action.text or ''})".strip()
+                )
+                break
+
             try:
                 result = self.executor.execute(action, observation)
             except Exception as exc:
@@ -142,10 +169,22 @@ class AgentLoop:
                 episode.stop_reason = f"action failed: {result.note}"
                 break
 
-            observation = self.screen.capture()  # verify: what did that change?
+            try:
+                observation = self.screen.capture()  # verify: what did that change?
+            except Exception as exc:
+                self._record(episode, Step(index + 1, observation,
+                                           error=f"capture failed: {type(exc).__name__}: {exc}"))
+                episode.stop_reason = "perception failed mid-episode"
+                break
 
         return episode
 
+    @staticmethod
+    def _fingerprint(observation: Observation) -> tuple:
+        """What the agent can see: enough to tell "nothing changed" from "it worked"."""
+        return (tuple(element.label for element in observation.elements),
+                observation.text[:200])
+
     def _record(self, episode: Episode, step: Step) -> None:
         episode.steps.append(step)
         if self.on_step is not None:

tiny_computer_agent/models.py

@@ -0,0 +1,187 @@
+"""Real (tiny) model backends, small enough to run on a laptop CPU.
+
+Two of them, because they trade off differently:
+
+* :class:`TextPlannerBrain` - a ~0.5B *text* model over the OCR index. No vision,
+  no big download, and the element index does all the grounding, which is why it
+  works at this size at all.
+* :class:`VLMBrain` - a tiny *vision* model that sees the screenshot itself
+  (256M-class models can read a button label from pixels, but reason poorly, so the
+  menu still carries most of the decision).
+
+Both inherit menu building, prompt shape and lenient parsing from
+:class:`~tiny_computer_agent.brain.ModelBrain`, so swapping weights never changes the
+harness contract.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from tiny_computer_agent.brain import BrainError, ModelBrain
+from tiny_computer_agent.observe import Observation
+
+DEFAULT_TEXT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+DEFAULT_VISION_MODEL = "HuggingFaceTB/SmolVLM-256M-Instruct"
+
+
+def resolve_device(prefer: str | None = None) -> str:
+    """Pick a torch device: explicit preference, then CUDA, then CPU."""
+    try:
+        import torch
+    except ImportError:  # pragma: no cover - depends on the machine
+        raise BrainError("model backends need torch: pip install 'tiny-computer-agent[vision]'")
+    if prefer:
+        return prefer
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def pick_dtype(device: str) -> Any:
+    import torch
+
+    if device.startswith("cuda"):
+        return torch.bfloat16
+    return torch.float32
+
+
+def model_dependencies() -> tuple[bool, str]:
+    """(available, detail) for preflight and the doctor command."""
+    try:
+        import transformers
+    except ImportError:
+        return False, "transformers is not installed (pip install 'tiny-computer-agent[vision]')"
+    try:
+        import torch
+    except ImportError:
+        return False, "torch is not installed (pip install 'tiny-computer-agent[vision]')"
+    return True, f"transformers {transformers.__version__}, torch {torch.__version__}"
+
+
+def model_cached(model_id: str) -> bool:
+    """True when the weights are already in the local Hugging Face cache."""
+    try:
+        from huggingface_hub import constants
+    except ImportError:  # pragma: no cover
+        return False
+    folder = os.path.join(constants.HF_HUB_CACHE, "models--" + model_id.replace("/", "--"))
+    return os.path.isdir(folder)
+
+
+class TextPlannerBrain(ModelBrain):
+    """Causal-LM planner over the OCR index: no pixels in, one menu number out."""
+
+    name = "text-planner"
+    requires_index = True
+
+    def __init__(self, model_id: str = DEFAULT_TEXT_MODEL, *, device: str | None = None,
+                 **kwargs: Any) -> None:
+        super().__init__(model_id, **kwargs)
+        self.device = device
+        self._tokenizer: Any = None
+        self._model: Any = None
+
+    def _load(self) -> None:
+        if self._model is not None:
+            return
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as exc:  # pragma: no cover - depends on the machine
+            raise BrainError(
+                "text planner needs transformers and torch: "
+                "pip install 'tiny-computer-agent[vision]'"
+            ) from exc
+        self.device = resolve_device(self.device)
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self._model = AutoModelForCausalLM.from_pretrained(
+            self.model_id, dtype=pick_dtype(self.device)
+        ).to(self.device).eval()
+
+    def unload(self) -> None:
+        self._model = None
+        self._tokenizer = None
+
+    def _generate(self, prompt: str, observation: Observation) -> str:
+        import torch
+
+        self._load()
+        messages = [{"role": "user", "content": prompt}]
+        text = self._tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=False
+        )
+        inputs = self._tokenizer(text, return_tensors="pt").to(self.device)
+        with torch.inference_mode():
+            output = self._model.generate(
+                **inputs,
+                max_new_tokens=self.max_new_tokens,
+                do_sample=False,
+                pad_token_id=self._tokenizer.eos_token_id,
+            )
+        return self._tokenizer.decode(
+            output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
+        ).strip()
+
+
+class VLMBrain(ModelBrain):
+    """Tiny vision-language model: sees the screenshot, answers with a menu number."""
+
+    name = "vlm"
+    requires_index = False
+
+    def __init__(self, model_id: str = DEFAULT_VISION_MODEL, *, device: str | None = None,
+                 **kwargs: Any) -> None:
+        super().__init__(model_id, **kwargs)
+        self.device = device
+        self._processor: Any = None
+        self._model: Any = None
+
+    def _load(self) -> None:
+        if self._model is not None:
+            return
+        try:
+            from transformers import AutoModelForImageTextToText, AutoProcessor
+        except ImportError as exc:  # pragma: no cover - depends on the machine
+            raise BrainError(
+                "the vision backend needs transformers and torch: "
+                "pip install 'tiny-computer-agent[vision]'"
+            ) from exc
+        self.device = resolve_device(self.device)
+        self._processor = AutoProcessor.from_pretrained(self.model_id)
+        self._model = AutoModelForImageTextToText.from_pretrained(
+            self.model_id, dtype=pick_dtype(self.device)
+        ).to(self.device).eval()
+
+    def unload(self) -> None:
+        self._model = None
+        self._processor = None
+
+    def _generate(self, prompt: str, observation: Observation) -> str:
+        import torch
+
+        image = observation.frame
+        if image is None or isinstance(image, str):
+            raise BrainError(
+                "the vision backend needs a real screenshot but the observation carries "
+                f"{type(image).__name__} - run with --live, or use the text planner"
+            )
+        self._load()
+        messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}]
+        text = self._processor.apply_chat_template(messages, add_generation_prompt=True)
+        inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device)
+        with torch.inference_mode():
+            output = self._model.generate(
+                **inputs, max_new_tokens=self.max_new_tokens, do_sample=False
+            )
+        return self._processor.batch_decode(
+            output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True
+        )[0].strip()
+
+
+MODEL_BRAINS = {"text-planner": TextPlannerBrain, "vlm": VLMBrain}
+
+
+def make_model_brain(name: str = "text-planner", **kwargs: Any) -> ModelBrain:
+    """Build a model brain by name, with an actionable error for unknown names."""
+    if name not in MODEL_BRAINS:
+        raise BrainError(f"unknown brain {name!r}; expected one of {tuple(MODEL_BRAINS)}")
+    return MODEL_BRAINS[name](**kwargs)

tiny_computer_agent/observe.py

@@ -1,15 +1,30 @@
-"""Perception: turn the world into an :class:`Observation` the brain can read.
+"""The Observation contract: screen frames and the element index built from them.
 
 The harness never asks the model to *find* anything. OCR or the accessibility tree
-produces a labelled element index; the model only chooses among those labels.
+produces a labelled element index (see :mod:`tiny_computer_agent.perception`); the
+model only chooses among those labels.
+
+Every :class:`Element` carries **absolute screen coordinates**, whoever built it.
 """
 
 from __future__ import annotations
 
+import sys
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Protocol, Sequence
 
 
+def normalise_label(text: str) -> str:
+    """Lower-case and keep only letters and digits.
+
+    OCR routinely loses or invents spaces ("Untitleddocument-unsavedchanges") and
+    models add punctuation ("Save."), so labels are compared on their alphanumeric
+    core - otherwise nothing a model says ever matches what was seen on screen.
+    """
+    return "".join(character for character in text.lower() if character.isalnum())
+
+
 @dataclass(frozen=True)
 class Element:
     """A clickable thing, located by the platform — not by the model."""
@@ -34,7 +49,8 @@ class Element:
         return (self.x, self.y)
 
     def matches(self, needle: str) -> bool:
-        return needle.strip().lower() in self.label.strip().lower()
+        wanted = normalise_label(needle)
+        return bool(wanted) and wanted in normalise_label(self.label)
 
 
 @dataclass(frozen=True)
@@ -47,9 +63,15 @@ class Observation:
     text: str = ""
 
     def find(self, label: str) -> Element | None:
-        """Exact label first, then substring; None if nothing matches."""
-        want = label.strip().lower()
-        exact = [e for e in self.elements if e.label.strip().lower() == want]
+        """Exact label first, then substring; None if nothing matches.
+
+        Comparison is on normalised text (see :func:`normalise_label`), so
+        ``find("Save")`` still matches a label the OCR rendered as ``"Save ."``.
+        """
+        wanted = normalise_label(label)
+        if not wanted:
+            return None
+        exact = [e for e in self.elements if normalise_label(e.label) == wanted]
         if exact:
             return exact[0]
         loose = [e for e in self.elements if e.matches(label)]
@@ -102,40 +124,130 @@ class FakeScreen:
         return observation
 
 
+class ImageScreen:
+    """A static screen source: replay one screenshot (or a sequence of them).
+
+    Lets the full pipeline - OCR, index, menu, model, action - be exercised with no
+    display at all, and lets you replay a captured failure offline. Accepts a path,
+    a PIL image, or a sequence of either.
+    """
+
+    def __init__(self, source: Any) -> None:
+        from PIL import Image
+
+        frames = [source] if isinstance(source, (str, Path)) or not isinstance(source, (list, tuple)) else list(source)
+        self._images: list[Any] = []
+        for frame in frames:
+            if isinstance(frame, (str, Path)):
+                self._images.append(Image.open(str(frame)).convert("RGB"))
+            else:
+                self._images.append(frame)
+        if not self._images:
+            raise ValueError("ImageScreen needs at least one image")
+        self.captures = 0
+
+    def capture(self) -> Observation:
+        index = min(self.captures, len(self._images) - 1)
+        self.captures += 1
+        return Observation(index=self.captures - 1, frame=self._images[index], elements=(), text="")
+
+
 class ScreenUnavailable(RuntimeError):
     """Raised when real capture is requested but the machine cannot provide it."""
 
 
 class GrabScreen:
-    """Real desktop capture via ``mss`` + Pillow.
-
-    Pixels only: nothing here builds an element index. Wiring OCR (tesseract) or
-    the accessibility tree (AT-SPI on Linux) is what turns frames into labels.
-    Needs a display, so it is untested in a headless sandbox.
+    """Capture the real desktop as an image.
+
+    Prefers ``mss`` (fast, multi-monitor) and falls back to Pillow's ``ImageGrab``.
+    Caps are deliberately absent: a frame pixel must be a *screen* pixel, which is
+    what lets every detector produce Elements with absolute screen coordinates.
+
+    Two options break that identity, so both are explicit: ``monitor`` may select a
+    display whose origin is not (0, 0), and ``max_width`` downscales a wide capture
+    (faster OCR, and a smaller image for a model to look at).
+    :meth:`map_to_screen` converts a frame coordinate back to a real one, and
+    :class:`~tiny_computer_agent.perception.IndexedScreen` applies it to every
+    detected element, so a click never lands on the wrong pixel.
     """
 
-    def __init__(self, *, monitor: int = 1, max_width: int = 1280) -> None:
+    def __init__(self, *, monitor: int = 1, max_width: int | None = None,
+                 scale: float | None = None) -> None:
+        if monitor < 0:
+            raise ValueError("monitor must be 0 (all displays) or a positive index")
+        if max_width is not None and max_width < 1:
+            raise ValueError("max_width must be positive")
+        if scale is not None and scale <= 0:
+            raise ValueError("scale must be positive")
         self.monitor = monitor
         self.max_width = max_width
+        self.scale = scale
+        self.origin: tuple[int, int] = (0, 0)
+        self.frame_scale: float = 1.0
         self._index = 0
 
     def capture(self) -> Observation:
+        image, origin = self._grab()
+        native_width = image.width
+        if self.max_width is not None and image.width > self.max_width:
+            ratio = self.max_width / image.width
+            image = image.resize((self.max_width, max(1, round(image.height * ratio))))
+        if self.scale is not None:
+            self.frame_scale = float(self.scale)
+        else:
+            self.frame_scale = native_width / image.width
+        self.origin = origin
+        observation = Observation(index=self._index, frame=image, elements=(), text="")
+        self._index += 1
+        return observation
+
+    def map_to_screen(self, element: Element) -> Element:
+        """Frame pixel -> absolute screen pixel, honouring the origin and any scale."""
+        if self.frame_scale == 1.0 and self.origin == (0, 0):
+            return element
+        factor = self.frame_scale
+        return Element(
+            label=element.label, kind=element.kind,
+            x=round(element.x * factor) + self.origin[0],
+            y=round(element.y * factor) + self.origin[1],
+            width=max(1, round(element.width * factor)),
+            height=max(1, round(element.height * factor)),
+        )
+
+    def _grab(self) -> tuple[Any, tuple[int, int]]:
+        """(image, screen position of the image's top-left pixel)."""
         try:
             import mss
+        except ImportError:
+            return self._grab_with_pillow()
+        try:
             from PIL import Image
-        except ImportError as exc:  # pragma: no cover - depends on the machine
+
+            with mss.mss() as sct:
+                monitors = sct.monitors
+                monitor = monitors[self.monitor] if self.monitor < len(monitors) else monitors[-1]
+                shot = sct.grab(monitor)
+                image = Image.frombytes("RGB", shot.size, shot.bgra, "raw", "BGRX")
+            return image, (int(monitor.get("left", 0)), int(monitor.get("top", 0)))
+        except Exception as exc:
             raise ScreenUnavailable(
-                "GrabScreen needs 'mss' and 'pillow': pip install \'tiny-computer-agent[desktop]\'"
+                f"screen capture failed (no display?): {type(exc).__name__}: {exc}"
             ) from exc
+
+    def _grab_with_pillow(self) -> tuple[Any, tuple[int, int]]:
         try:
-            with mss.mss() as sct:
-                shot = sct.grab(sct.monitors[self.monitor])
-                image = Image.frombytes("RGB", shot.size, shot.bgra, "raw", "BGRX")
-        except Exception as exc:  # pragma: no cover - depends on the machine
-            raise ScreenUnavailable(f"screen capture failed (no display?): {exc}") from exc
-        if image.width > self.max_width:
-            ratio = self.max_width / image.width
-            image = image.resize((self.max_width, int(image.height * ratio)))
-        observation = Observation(index=self._index, frame=image, elements=(), text="")
-        self._index += 1
-        return observation
+            from PIL import ImageGrab
+        except ImportError as exc:  # pragma: no cover - pillow is a desktop extra
+            raise ScreenUnavailable(
+                "screen capture needs mss or Pillow: pip install \'tiny-computer-agent[desktop]\'"
+            ) from exc
+        try:
+            if sys.platform.startswith("win"):
+                image = ImageGrab.grab(all_screens=self.monitor == 0)
+            else:
+                image = ImageGrab.grab()
+        except Exception as exc:
+            raise ScreenUnavailable(
+                f"screen capture failed (no display?): {type(exc).__name__}: {exc}"
+            ) from exc
+        return image.convert("RGB"), (0, 0)

tiny_computer_agent/perception.py

@@ -0,0 +1,225 @@
+"""Turning pixels into a labelled element index.
+
+This is the module that makes a *tiny* model viable. The model never has to guess
+where anything is, because OCR has already turned the screen into a short list of
+labels with coordinates - and the model only has to pick one.
+
+Filters matter as much as the engine. Raw OCR returns dialog prose, fragmentary
+debris and stray punctuation, and every extra line of noise is another chance for a
+small model to pick the wrong element, so the index is cleaned aggressively.
+
+Everything here is optional at import time: a missing OCR engine raises a clear
+:class:`DetectorUnavailable` instead of breaking the harness.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from tiny_computer_agent.observe import Element, Observation, ScreenSource
+
+
+class DetectorUnavailable(RuntimeError):
+    """Raised when an element detector is requested but cannot run on this machine."""
+
+
+class Detector(Protocol):
+    """Turns one frame into labelled, clickable elements."""
+
+    name: str
+
+    def detect(self, image: Any) -> tuple[Element, ...]: ...
+
+
+@dataclass
+class OcrDetector:
+    """Indexes on-screen text with RapidOCR: ONNX, CPU, offline after install.
+
+    ``max_label_chars`` is the important knob. Screen text includes paragraphs and
+    titles that are not widgets, and a model shown them will happily click prose -
+    which is exactly what a 0.5B planner did when dialog text sat in its menu. Short
+    strings are much more likely to be buttons and menu items, so the default is
+    deliberately tight.
+    """
+
+    min_score: float = 0.6
+    max_label_chars: int = 24
+    min_label_chars: int = 2
+    name: str = "ocr"
+    _engine: Any = field(default=None, init=False, repr=False)
+
+    def _loaded(self) -> Any:
+        if self._engine is None:
+            try:
+                from rapidocr_onnxruntime import RapidOCR
+            except ImportError as exc:  # pragma: no cover - depends on the machine
+                raise DetectorUnavailable(
+                    "no OCR engine available: pip install 'tiny-computer-agent[ocr]'"
+                ) from exc
+            self._engine = RapidOCR()
+        return self._engine
+
+    def detect(self, image: Any) -> tuple[Element, ...]:
+        import numpy as np
+
+        array = np.asarray(image.convert("RGB") if hasattr(image, "convert") else image)
+        result, _ = self._loaded()(array)
+        elements: list[Element] = []
+        for box, text, score in result or []:
+            element = self._to_element(box, text, float(score))
+            if element is not None:
+                elements.append(element)
+        return tuple(elements)
+
+    def _to_element(self, box: Any, text: Any, score: float) -> Element | None:
+        """Convert one OCR box to an Element, or None if it is not a usable label."""
+        label = " ".join(str(text).split())
+        if score < self.min_score:
+            return None
+        if not self.min_label_chars <= len(label) <= self.max_label_chars:
+            return None
+        if not any(character.isalnum() for character in label):
+            return None
+        xs = [float(point[0]) for point in box]
+        ys = [float(point[1]) for point in box]
+        x, y = int(min(xs)), int(min(ys))
+        return Element(
+            label=label, kind="text",
+            x=x, y=y,
+            width=int(max(xs)) - x, height=int(max(ys)) - y,
+        )
+
+
+@dataclass
+class IndexedScreen:
+    """Composes pixel capture with a detector, so ``capture()`` yields an index.
+
+    This is the wiring that turns the harness from a scripted demo into something
+    that can act on a real desktop.
+    """
+
+    source: ScreenSource
+    detector: Detector
+    fill_text: bool = True
+
+    def capture(self) -> Observation:
+        observation = self.source.capture()
+        try:
+            elements = self.detector.detect(observation.frame)
+        except DetectorUnavailable:
+            raise
+        except Exception as exc:
+            # A perception failure must not kill a run that is otherwise fine.
+            elements = ()
+            observation = Observation(
+                index=observation.index, frame=observation.frame,
+                elements=(), text=f"perception error: {type(exc).__name__}: {exc}",
+            )
+            return observation
+        mapper = getattr(self.source, "map_to_screen", None)
+        if mapper is not None:
+            # Frame pixels -> absolute screen pixels, so a click cannot drift.
+            elements = tuple(mapper(element) for element in elements)
+        text = observation.text
+        if self.fill_text and not text:
+            # Grid labels are coordinates, not visible text, so they are not page text.
+            text = " ".join(element.label for element in elements if element.kind != "grid")
+        return Observation(
+            index=observation.index, frame=observation.frame, elements=elements, text=text,
+        )
+
+
+@dataclass
+class GridDetector:
+    """A numbered grid over the frame: click targets that need nothing installed.
+
+    The guarantee that matters. OCR finds text, the accessibility tree finds widgets,
+    but a canvas app, a game, a photo, a custom-drawn toolbar - anything with no text
+    - leaves a text index empty, and an agent with an empty index has *no legal
+    moves*. A grid always has moves, and a numbered menu is something even a 256M
+    model can choose from reliably.
+    """
+
+    rows: int = 8
+    cols: int = 12
+    name: str = "grid"
+
+    def __post_init__(self) -> None:
+        if self.rows < 1 or self.cols < 1:
+            raise ValueError("rows and cols must be at least 1")
+
+    def detect(self, image: Any) -> tuple[Element, ...]:
+        width, height = image.size
+        cell_w = width / self.cols
+        cell_h = height / self.rows
+        elements = []
+        for row in range(self.rows):
+            for col in range(self.cols):
+                elements.append(Element(
+                    label=f"cell {row * self.cols + col + 1}",
+                    kind="grid",
+                    x=round(col * cell_w), y=round(row * cell_h),
+                    width=max(1, round(cell_w)), height=max(1, round(cell_h)),
+                ))
+        return tuple(elements)
+
+
+@dataclass
+class AutoDetector:
+    """OCR when it finds text, numbered cells when it does not.
+
+    This is the default, so the harness is useful on a machine with nothing installed
+    beyond Pillow - and gets better automatically once an OCR engine appears.
+    """
+
+    name: str = "auto"
+    ocr: Detector = field(default_factory=OcrDetector)
+    grid: Detector = field(default_factory=GridDetector)
+
+    def detect(self, image: Any) -> tuple[Element, ...]:
+        try:
+            elements = self.ocr.detect(image)
+        except DetectorUnavailable:
+            return self.grid.detect(image)
+        return elements or self.grid.detect(image)
+
+
+def _ocr_available() -> tuple[bool, str]:
+    try:
+        import rapidocr_onnxruntime  # noqa: F401
+    except ImportError:
+        return False, "rapidocr-onnxruntime is not installed (pip install 'tiny-computer-agent[ocr]')"
+    try:
+        import onnxruntime  # noqa: F401
+    except ImportError:  # pragma: no cover
+        return False, "onnxruntime is not installed"
+    return True, "rapidocr-onnxruntime ready (CPU, offline)"
+
+
+def detector_available(name: str = "auto") -> tuple[bool, str]:
+    """(available, detail) for preflight and the ``doctor`` command."""
+    if name == "grid":
+        return True, "numbered grid cells (no dependencies)"
+    if name == "auto":
+        available, detail = _ocr_available()
+        if available:
+            return True, f"OCR when it finds text, grid otherwise ({detail})"
+        return True, "grid only: no OCR engine installed (pip install 'tiny-computer-agent[ocr]')"
+    if name != "ocr":
+        return False, f"unknown detector {name!r}"
+    return _ocr_available()
+
+
+def make_detector(name: str = "auto", **kwargs: Any) -> Detector:
+    """Build a detector by name, with an actionable error when it cannot run."""
+    available, detail = detector_available(name)
+    if not available:
+        raise DetectorUnavailable(detail)
+    if name == "grid":
+        return GridDetector(**kwargs)
+    if name == "auto":
+        if _ocr_available()[0]:
+            return AutoDetector(**kwargs)
+        return GridDetector(**kwargs)
+    return OcrDetector(**kwargs)

tiny_computer_agent/safety.py

@@ -50,6 +50,11 @@ class Policy:
     deny_keys          : key chords that are never pressed
     kill_switch_path   : if this file exists, the episode aborts before the next
                          action — create it to stop a live run
+    allow_grid_clicks  : numbered grid cells are *blind* clicks: they carry no
+                         meaning, only coordinates. Allowed while planning (a dry
+                         run touches nothing), refused with real input unless this
+                         is set, because "click cell 47" is a guess about your
+                         desktop, not a decision about it.
     """
 
     dry_run: bool = True
@@ -58,6 +63,7 @@ class Policy:
     deny_labels: tuple[str, ...] = DEFAULT_DENY_LABELS
     deny_keys: tuple[str, ...] = DEFAULT_DENY_KEYS
     kill_switch_path: str | None = None
+    allow_grid_clicks: bool = False
     _compiled: tuple[Pattern[str], ...] = field(default=(), init=False, repr=False)
 
     def __post_init__(self) -> None:
@@ -85,6 +91,15 @@ class Policy:
 
         if action.kind == "click" and action.label:
             element = observation.find(action.label) if observation is not None else None
+            if not self.dry_run and not self.allow_grid_clicks:
+                is_grid = (element is not None and element.kind == "grid") or (
+                    element is None and action.label.strip().lower().startswith("cell ")
+                )
+                if is_grid:
+                    raise PolicyViolation(
+                        "clicking a numbered grid cell is a blind click; install an OCR "
+                        "engine or pass --allow-grid if you really mean it"
+                    )
             haystack = (element.label if element is not None else action.label).lower()
             for denied in self.deny_labels:
                 if denied.lower() in haystack: