Publish tiny_computer_agent/ from marimo

tiny_computer_agent/__init__.py

@@ -14,6 +14,15 @@ call at run time, and the core imports nothing outside the standard library.
     loop = AgentLoop(screen=screen, brain=brain)             # dry run by default
     episode = loop.run("save the document")
     print(episode.render())
+
+An episode can be saved as a portable *trace*, replayed offline, and scored, so
+"the model looked plausible on that screenshot" becomes a number:
+
+    from tiny_computer_agent import Trace, benchmark, run_scenario
+
+    trace = run_scenario(Scenario("click 'Save'", frames=[...], expected="Save"), brain)
+    trace.save("traces/save.json")                           # diffable, no pixels
+    print(benchmark(scenarios, {"planner": lambda s: PlannerBrain()}).render())
 """
 
 from tiny_computer_agent.act import (
@@ -74,8 +83,22 @@ from tiny_computer_agent.perception import (
     make_detector,
 )
 from tiny_computer_agent.safety import Policy, PolicyViolation
+from tiny_computer_agent.trace import (
+    Scenario,
+    Scoreboard,
+    Trace,
+    TraceStep,
+    Verdict,
+    benchmark,
+    benchmark_traces,
+    load_traces,
+    run_scenario,
+    score,
+    score_replay,
+    score_trace,
+)
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 __all__ = [
     "ACTION_KINDS", "ACTION_MENU", "Action", "ActionError", "ActionResult",
@@ -84,9 +107,11 @@ __all__ = [
     "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__",
+    "PolicyViolation", "PyAutoGuiExecutor", "Scenario", "Scoreboard", "ScreenSource",
+    "ScreenUnavailable", "Step", "TextPlannerBrain", "Trace", "TraceStep", "VLMBrain",
+    "Verdict", "benchmark", "benchmark_traces", "detector_available",
+    "extract_json_object", "extract_json_objects", "load_traces", "make_detector",
+    "make_model_brain", "model_cached", "model_dependencies", "normalise_label",
+    "parse_action", "parse_decision", "resolve_device", "resolve_point",
+    "run_scenario", "score", "score_replay", "score_trace", "__version__",
 ]

tiny_computer_agent/__main__.py

@@ -32,6 +32,7 @@ from tiny_computer_agent.observe import (
 )
 from tiny_computer_agent.perception import IndexedScreen, detector_available, make_detector
 from tiny_computer_agent.safety import Policy
+from tiny_computer_agent.trace import Trace, score_trace
 
 ALLOW_REAL_INPUT_ENV = "TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT"
 EXAMPLE_IMAGE = "example"
@@ -88,6 +89,12 @@ def build_parser() -> argparse.ArgumentParser:
                         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")
+    parser.add_argument("--trace", default=None, metavar="PATH",
+                        help="write a portable trace of the run (frames, no pixels) so it "
+                             "can be replayed and scored offline")
+    parser.add_argument("--expected", default=None, metavar="LABEL",
+                        help="the element label that is the right answer for this goal; "
+                             "recorded in the trace so the run can be scored")
     return parser
 
 
@@ -247,6 +254,14 @@ def main(argv: list[str] | None = None) -> int:
 
     print()
     print(episode.render())
+
+    if args.trace:
+        trace = Trace.record(episode, expected=args.expected, brain=args.brain,
+                             scenario=args.goal)
+        print(f"trace: wrote {trace.save(args.trace)} ({trace.frames} frames)")
+        if args.expected:
+            print("score: " + score_trace(trace).line().strip())
+            print(f"       replay it later: Trace.load({str(args.trace)!r}).render()")
     return 0 if episode.success else 1
 
 

tiny_computer_agent/trace.py

@@ -0,0 +1,539 @@
+"""Traces: an episode you can save, replay and score.
+
+An :class:`~tiny_computer_agent.loop.Episode` is the record of one run; a **trace**
+is that record made portable. A trace keeps each frame's *element index* - the
+labels and their geometry - but never the pixels, so it replays on a machine with
+no display, no OCR engine and no model weights: hand the recorded frames back to a
+brain and count how often it picks the element the trace says was right.
+
+That count is the whole point of this module. It is what turns "the model looked
+plausible on that one screenshot" into an **accuracy** you can compare between
+brains, put in a table, and regression-test in CI - the harness's own answer to the
+question the results table keeps asking: *which brain is actually right, and how
+often?*
+
+Two judgements are kept apart on purpose. The score is about the brain's
+**decision** - what it chose when it had to commit. Whether the harness *allowed*
+that decision to run is a separate axis, reported by :attr:`Verdict.blocked`, so a
+policy refusal can never masquerade as a clean success (the harness owns success)
+nor hide a brain that chose correctly.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Any, Callable, Mapping, Sequence
+
+from tiny_computer_agent.actions import Action, parse_action
+from tiny_computer_agent.brain import Brain
+from tiny_computer_agent.loop import AgentLoop, Episode
+from tiny_computer_agent.observe import (
+    Element,
+    FakeScreen,
+    Observation,
+    normalise_label,
+)
+from tiny_computer_agent.safety import Policy
+
+TRACE_VERSION = 1
+DONE = "done"
+NOTHING = "nothing"
+
+# The fields of Element, in one place, so serialisation and Element stay in step.
+_ELEMENT_FIELDS = ("label", "kind", "x", "y", "width", "height")
+
+
+def _element_to_dict(element: Element) -> dict[str, Any]:
+    return {name: getattr(element, name) for name in _ELEMENT_FIELDS}
+
+
+def _element_from_dict(data: Mapping[str, Any]) -> Element:
+    return Element(**{name: data[name] for name in _ELEMENT_FIELDS if name in data})
+
+
+@dataclass(frozen=True)
+class TraceStep:
+    """One recorded frame: what was on screen, and what the brain did about it."""
+
+    index: int
+    elements: tuple[Element, ...] = ()
+    text: str = ""
+    action: Action | None = None
+    error: str | None = None
+    outcome: str = ""
+
+    def observation(self) -> Observation:
+        """The frame as the brain saw it - pixels deliberately omitted."""
+        return Observation(index=self.index, elements=self.elements, text=self.text)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "index": self.index,
+            "elements": [_element_to_dict(e) for e in self.elements],
+            "text": self.text,
+            "action": self.action.to_dict() if self.action is not None else None,
+            "error": self.error,
+            "outcome": self.outcome,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> TraceStep:
+        raw_action = data.get("action")
+        return cls(
+            index=int(data.get("index", 0)),
+            elements=tuple(_element_from_dict(e) for e in data.get("elements", ())),
+            text=data.get("text", ""),
+            action=parse_action(raw_action) if raw_action else None,
+            error=data.get("error"),
+            outcome=data.get("outcome", ""),
+        )
+
+
+@dataclass
+class Trace:
+    """A saved run: the frames, the decisions, and the answer that was right.
+
+    ``expected`` is the ground truth a scorer compares against. It lives on the
+    trace rather than on a step because a task has one correct *choice*, not one
+    correct action per turn - the waits and retries on the way there are not
+    decisions about what to do.
+    """
+
+    goal: str
+    expected: str | None = None
+    brain: str = ""
+    scenario: str = ""
+    steps: list[TraceStep] = field(default_factory=list)
+    success: bool = False
+    stop_reason: str = ""
+
+    # -- recording ---------------------------------------------------------
+
+    @classmethod
+    def record(
+        cls,
+        episode: Episode,
+        *,
+        expected: str | None = None,
+        brain: str = "",
+        scenario: str = "",
+    ) -> Trace:
+        """Copy an episode into a portable trace (no pixels, no live objects)."""
+        return cls(
+            goal=episode.goal,
+            expected=expected,
+            brain=brain,
+            scenario=scenario,
+            steps=[
+                TraceStep(
+                    index=step.index,
+                    elements=tuple(step.observation.elements),
+                    text=step.observation.text,
+                    action=step.action,
+                    error=step.error,
+                    outcome=step.result.summary() if step.result is not None else "",
+                )
+                for step in episode.steps
+            ],
+            success=episode.success,
+            stop_reason=episode.stop_reason,
+        )
+
+    # -- reading -----------------------------------------------------------
+
+    @property
+    def name(self) -> str:
+        return self.scenario or self.goal
+
+    @property
+    def frames(self) -> int:
+        return len(self.steps)
+
+    def decision_step(self) -> TraceStep | None:
+        """The step where the brain committed: first labelled click, else first ``done``.
+
+        Waits are skipped - waiting is patience, not a choice about what to do.
+        """
+        for step in self.steps:
+            action = step.action
+            if action is None:
+                continue
+            if action.kind == "click" and action.label:
+                return step
+            if action.kind == "done":
+                return step
+        return None
+
+    @property
+    def choice(self) -> str | None:
+        """The label the brain chose, or None when it committed to no element."""
+        step = self.decision_step()
+        if step is not None and step.action is not None and step.action.kind == "click":
+            return step.action.label
+        return None
+
+    @property
+    def decided(self) -> str:
+        """What the brain committed to: a label, the bare word ``done``, or nothing.
+
+        Kept deliberately blunt so a mis-scored run shows up in the table: a brain
+        that answered "done" before ever looking has not chosen the right element.
+        """
+        step = self.decision_step()
+        if step is None or step.action is None:
+            return NOTHING
+        return step.action.label if step.action.kind == "click" else DONE
+
+    @property
+    def blocked(self) -> bool:
+        """True when the harness stopped the deciding action (e.g. a policy refusal)."""
+        step = self.decision_step()
+        return step is not None and step.error is not None
+
+    def observations(self) -> list[Observation]:
+        return [step.observation() for step in self.steps]
+
+    def render(self) -> str:
+        """A compact, replayable summary - the trace's own version of Episode.render()."""
+        header = f"trace: {self.name}"
+        if self.brain:
+            header += f"  (brain: {self.brain})"
+        lines = [
+            header,
+            f"goal: {self.goal}",
+            f"expected: {self.expected!r}   chose: {self.decided!r}"
+            + ("   (the harness refused it)" if self.blocked else ""),
+            f"result: {'success' if self.success else 'failure'} - {self.stop_reason}",
+        ]
+        for step in self.steps:
+            if step.action is not None:
+                body = step.action.describe()
+            else:
+                body = step.error or "no action"
+            lines.append(f"  {step.index} | {body}")
+        return "\n".join(lines)
+
+    # -- replay ------------------------------------------------------------
+
+    def replay(self, brain: Brain) -> Trace:
+        """Answer the recorded frames again with a (usually different) brain.
+
+        Nothing is captured and nothing is clicked: this is the offline half of
+        "replay a trace", and it is why a trace can be scored on any machine.
+        """
+        steps: list[TraceStep] = []
+        for step in self.steps:
+            try:
+                action = brain.decide(self.goal, step.observation())
+                steps.append(replace(step, action=action, error=None, outcome="replayed"))
+            except Exception as exc:  # a brain that fails is a scored outcome, not a crash
+                steps.append(
+                    replace(step, action=None, error=f"{type(exc).__name__}: {exc}", outcome="")
+                )
+        return Trace(
+            goal=self.goal,
+            expected=self.expected,
+            brain=getattr(brain, "name", type(brain).__name__),
+            scenario=self.scenario,
+            steps=steps,
+            success=False,
+            stop_reason="replayed",
+        )
+
+    # -- persistence -------------------------------------------------------
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "trace_version": TRACE_VERSION,
+            "goal": self.goal,
+            "expected": self.expected,
+            "brain": self.brain,
+            "scenario": self.scenario,
+            "success": self.success,
+            "stop_reason": self.stop_reason,
+            "steps": [step.to_dict() for step in self.steps],
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> Trace:
+        version = int(data.get("trace_version", TRACE_VERSION))
+        if version > TRACE_VERSION:
+            raise ValueError(
+                f"trace_version {version} is newer than this build understands "
+                f"({TRACE_VERSION}); upgrade tiny-computer-agent to read it"
+            )
+        return cls(
+            goal=data.get("goal", ""),
+            expected=data.get("expected"),
+            brain=data.get("brain", ""),
+            scenario=data.get("scenario", ""),
+            steps=[TraceStep.from_dict(step) for step in data.get("steps", ())],
+            success=bool(data.get("success", False)),
+            stop_reason=data.get("stop_reason", ""),
+        )
+
+    def to_json(self, *, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent, ensure_ascii=False)
+
+    @classmethod
+    def from_json(cls, text: str) -> Trace:
+        return cls.from_dict(json.loads(text))
+
+    def save(self, path: str | Path) -> Path:
+        target = Path(path)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(self.to_json() + "\n", encoding="utf-8")
+        return target
+
+    @classmethod
+    def load(cls, path: str | Path) -> Trace:
+        return cls.from_json(Path(path).read_text(encoding="utf-8"))
+
+
+def load_traces(path: str | Path) -> list[Trace]:
+    """Every trace in a file or a directory of ``*.json`` traces, in name order."""
+    target = Path(path)
+    files = sorted(target.glob("*.json")) if target.is_dir() else [target]
+    return [Trace.load(file) for file in files]
+
+
+# -- scoring ---------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Verdict:
+    """One trace, judged: did the brain choose what the trace says was right?"""
+
+    scenario: str
+    expected: str | None
+    chosen: str
+    correct: bool
+    brain: str = ""
+    blocked: bool = False
+    detail: str = ""
+
+    def line(self) -> str:
+        mark = "correct  " if self.correct else "incorrect"
+        line = f"  {mark} {self.scenario}: expected {self.expected!r}, chose {self.chosen!r}"
+        if self.blocked:
+            line += " (harness refused it)"
+        elif self.detail:
+            line += f" ({self.detail})"
+        return line
+
+
+def score_trace(trace: Trace) -> Verdict:
+    """Judge one recorded choice against the trace's ground truth.
+
+    Comparison is on :func:`~tiny_computer_agent.observe.normalise_label`, so a
+    choice the OCR rendered as ``"Save ."`` is still credited for ``"Save"`` - the
+    same forgiveness the element index uses everywhere else. A decision the policy
+    refused is still scored on the *choice* (that is what the benchmark measures),
+    and flagged by :attr:`Verdict.blocked` rather than counted as a clean success.
+    """
+    chosen = trace.decided
+    expected = trace.expected
+    if expected is None:
+        return Verdict(
+            scenario=trace.name, expected=None, chosen=chosen, correct=False,
+            brain=trace.brain, blocked=trace.blocked,
+            detail="trace records no expected answer to score against",
+        )
+    correct = normalise_label(chosen) == normalise_label(expected)
+    detail = "" if correct else (
+        f"chose {chosen!r}, not {expected!r}" if chosen != NOTHING
+        else "never committed to a choice"
+    )
+    return Verdict(
+        scenario=trace.name, expected=expected, chosen=chosen, correct=correct,
+        brain=trace.brain, blocked=trace.blocked, detail=detail,
+    )
+
+
+def score_replay(trace: Trace, brain: Brain) -> Verdict:
+    """"Replay a trace, count how often the choice was right", for a single trace."""
+    return score_trace(trace.replay(brain))
+
+
+@dataclass
+class Scoreboard:
+    """Verdicts collected into an accuracy - the number this whole module exists for."""
+
+    verdicts: list[Verdict] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.verdicts)
+
+    def __iter__(self):
+        return iter(self.verdicts)
+
+    def add(self, verdict: Verdict) -> Scoreboard:
+        self.verdicts.append(verdict)
+        return self
+
+    @property
+    def total(self) -> int:
+        return len(self.verdicts)
+
+    @property
+    def earned(self) -> int:
+        return sum(1 for verdict in self.verdicts if verdict.correct)
+
+    @property
+    def accuracy(self) -> float:
+        """Share of decisions that matched the right answer, 0.0 when empty."""
+        return self.earned / self.total if self.total else 0.0
+
+    @property
+    def blocked(self) -> int:
+        """Decisions the harness stopped, whatever the score - a separate axis."""
+        return sum(1 for verdict in self.verdicts if verdict.blocked)
+
+    def brains(self) -> list[str]:
+        """Every brain that appears, in first-seen order, so tables stay stable."""
+        seen: list[str] = []
+        for verdict in self.verdicts:
+            if verdict.brain not in seen:
+                seen.append(verdict.brain)
+        return seen
+
+    def accuracy_for(self, brain: str) -> float:
+        rows = [v for v in self.verdicts if v.brain == brain]
+        return sum(1 for v in rows if v.correct) / len(rows) if rows else 0.0
+
+    def by_brain(self) -> dict[str, float]:
+        return {brain: self.accuracy_for(brain) for brain in self.brains()}
+
+    def to_rows(self) -> list[dict[str, Any]]:
+        """Rows for a dataframe / notebook table. ``correct`` is a plain bool."""
+        return [
+            {
+                "brain": verdict.brain or "-",
+                "scenario": verdict.scenario,
+                "expected": verdict.expected,
+                "chose": verdict.chosen,
+                "correct": verdict.correct,
+            }
+            for verdict in self.verdicts
+        ]
+
+    def confusion(self) -> dict[tuple[str, str], int]:
+        """(expected, chose) -> count, for the cases a scalar accuracy hides."""
+        counts: dict[tuple[str, str], int] = {}
+        for verdict in self.verdicts:
+            key = (verdict.expected or "-", verdict.chosen)
+            counts[key] = counts.get(key, 0) + 1
+        return counts
+
+    def render(self) -> str:
+        lines = [f"scoreboard: {self.earned}/{self.total} correct ({self.accuracy:.0%})"]
+        lines += [f"  {brain or '-'}: {self.accuracy_for(brain):.0%}" for brain in self.brains()]
+        lines += [verdict.line() for verdict in self.verdicts]
+        if self.blocked:
+            lines.append(f"  ({self.blocked} decision(s) stopped by the harness)")
+        return "\n".join(lines)
+
+
+# -- a suite of tasks, for benchmarking -----------------------------------
+
+
+@dataclass(frozen=True)
+class Scenario:
+    """A task with one element that is the right answer, and the frames to show."""
+
+    goal: str
+    frames: tuple[tuple[Element, ...], ...]
+    expected: str
+    text: str = ""
+    name: str = ""
+
+    def __post_init__(self) -> None:
+        if not self.frames:
+            raise ValueError("a scenario needs at least one frame")
+        # freeze the frames so a scenario can be replayed and compared safely
+        object.__setattr__(self, "frames", tuple(tuple(frame) for frame in self.frames))
+
+    @property
+    def label(self) -> str:
+        return self.name or self.goal
+
+    def screen(self) -> FakeScreen:
+        """A fresh screen each time, so one scenario can be run against many brains."""
+        return FakeScreen(list(self.frames), text=self.text)
+
+
+def run_scenario(
+    scenario: Scenario, brain: Brain, *, brain_name: str = "", max_steps: int = 6
+) -> Trace:
+    """Run one scenario and return its trace, ground truth included."""
+    episode = AgentLoop(
+        screen=scenario.screen(),
+        brain=brain,
+        policy=Policy(step_budget=max_steps),
+    ).run(scenario.goal)
+    return Trace.record(
+        episode,
+        expected=scenario.expected,
+        brain=brain_name or getattr(brain, "name", type(brain).__name__),
+        scenario=scenario.label,
+    )
+
+
+def score(scenario: Scenario, brain: Brain, *, brain_name: str = "", max_steps: int = 6) -> Verdict:
+    return score_trace(run_scenario(scenario, brain, brain_name=brain_name, max_steps=max_steps))
+
+
+def benchmark(
+    scenarios: Sequence[Scenario],
+    brains: Mapping[str, Callable[[Scenario], Brain]],
+    *,
+    max_steps: int = 6,
+) -> Scoreboard:
+    """Run every scenario through every brain, scoring the choice each one made.
+
+    ``brains`` maps a display name to a *factory* taking the scenario, because some
+    brains (a scripted :class:`~tiny_computer_agent.brain.MockBrain`, say) need to see
+    the task before they can answer it, and because a fresh brain per task keeps state
+    like "already clicked this" from leaking between scenarios.
+    """
+    board = Scoreboard()
+    for scenario in scenarios:
+        for name, factory in brains.items():
+            board.add(score(scenario, factory(scenario), brain_name=name, max_steps=max_steps))
+    return board
+
+
+def benchmark_traces(traces: Sequence[Trace], brain: Brain, *, name: str = "") -> Scoreboard:
+    """Replay recorded traces against one brain and count how often it was right.
+
+    The literal roadmap step: the frames are already recorded, nothing new is
+    captured, and the only variable is the brain under test.
+    """
+    label = name or getattr(brain, "name", type(brain).__name__)
+    board = Scoreboard()
+    for trace in traces:
+        replayed = trace.replay(brain)
+        replayed.brain = label
+        board.add(score_trace(replayed))
+    return board
+
+
+__all__ = [
+    "DONE",
+    "NOTHING",
+    "TRACE_VERSION",
+    "Scenario",
+    "Scoreboard",
+    "Trace",
+    "TraceStep",
+    "Verdict",
+    "benchmark",
+    "benchmark_traces",
+    "load_traces",
+    "run_scenario",
+    "score",
+    "score_replay",
+    "score_trace",
+]