Playback: capture model reasoning + per-turn goal tracker + viewer

- providers: ChatReply.reasoning; extract reasoning_content/reasoning
from the raw reply; _wire_messages() strips playback-only keys so
reasoning never posts back (strict vLLM-safe), _reply_from_data pure
- agent: persist reply.reasoning on the assistant history turn
- goal_tracker: leaf_progress (win-cond predicate progress) +
reward_vector (normalized cumulative, scenario-agnostic) + turn_goal
bundling both side by side
- playback.record_turn records goal; eval_core wires turn_goal
- playback_view.load_episode (tested data layer) + Streamlit viewer
(scripts/view_playback.py): minimap, briefing, reasoning, tool_calls,
signals, objective bars + reward vector
- tests: test_playback_completeness.py (7) — 299 passed, 1 skipped

openra_bench/agent.py

@@ -421,6 +421,10 @@ class ModelAgent:
             {
                 "role": "assistant",
                 "content": reply.text or "",
+                # Playback-only: the wire layer (providers._wire_messages)
+                # strips this before posting, so it never goes back to
+                # the model but is preserved in messages.json.
+                "reasoning": reply.reasoning or "",
                 "tool_calls": [
                     {
                         "id": f"c{i}",

openra_bench/eval_core.py

@@ -217,8 +217,12 @@ def run_level(
                     _png = _render_minimap_b64(rs)
                 except Exception:  # noqa: BLE001 — playback never breaks a run
                     pass
+                from .goal_tracker import turn_goal
+
                 playback.record_turn(
-                    turns, rs, cmds, adapter.signals, _png, interrupt=interrupt
+                    turns, rs, cmds, adapter.signals, _png,
+                    interrupt=interrupt,
+                    goal=turn_goal(compiled.win_condition, ctx),
                 )
             trace.append(
                 {

openra_bench/goal_tracker.py

@@ -0,0 +1,131 @@
+"""Per-turn goal tracking for playback.
+
+Two complementary views of "how is the agent doing toward the
+objective", recorded every turn so a replay shows the trajectory of
+intent, not just the final verdict:
+
+* ``leaf_progress`` — walks the scenario's declarative win-condition
+  tree and, for every leaf predicate, reports current vs. target and a
+  0..1 ratio. Scenario-specific, directly tied to *this* objective.
+* ``reward_vector`` — a fixed, normalized, monotone-cumulative vector
+  (economy / military / territory / scouting / objective) that is
+  scenario-agnostic and therefore comparable across runs and on the
+  leaderboard. Mirrors the training rollout's per-turn reward_vector.
+
+``turn_goal`` bundles both side by side plus a scalar
+``objective_progress`` and the boolean ``won``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .scenarios.win_conditions import LEAF_KEYS, WinContext, evaluate
+
+# Current-value extractors for the numeric leaf predicates. Anything not
+# listed is treated as a pure boolean gate (ratio = 1.0 iff satisfied).
+_CURRENT: dict[str, Any] = {
+    "explored_pct_gte": lambda c: c.signals.explored_percent,
+    "enemies_discovered_gte": lambda c: len(c.signals.enemies_seen_ids),
+    "buildings_discovered_gte": lambda c: len(c.signals.enemy_buildings_seen_ids),
+    "units_killed_gte": lambda c: c.signals.units_killed,
+    "units_lost_lte": lambda c: c.signals.units_lost,
+    "within_ticks": lambda c: c.signals.game_tick,
+    "after_ticks": lambda c: c.signals.game_tick,
+    "own_units_gte": lambda c: len(c.render_state.get("units_summary", []) or []),
+    "cash_gte": lambda c: c.signals.cash,
+    "resources_gte": lambda c: c.signals.resources,
+    "economy_value_gte": lambda c: c.signals.cash + c.signals.resources,
+    "power_surplus_gte": lambda c: c.signals.power_provided
+    - c.signals.power_drained,
+    "buildings_owned_gte": lambda c: len(c.signals.own_building_types),
+    "building_total_gte": lambda c: len(c.signals.own_buildings),
+}
+# `_lte` predicates are satisfied while *below* the bound (constraints).
+_LTE = frozenset({"units_lost_lte", "within_ticks"})
+
+
+def _clamp(x: float) -> float:
+    return 0.0 if x < 0.0 else 1.0 if x > 1.0 else float(x)
+
+
+def _leaves(node: Any) -> list[dict]:
+    """Flatten a win-condition tree to its leaf predicate dicts."""
+    if node is None:
+        return []
+    if not isinstance(node, dict):
+        node = dict(getattr(node, "__pydantic_extra__", {}) or {})
+    out: list[dict] = []
+    for k, v in node.items():
+        if k in ("all_of", "any_of"):
+            for child in v:
+                out.extend(_leaves(child))
+        elif k == "not":
+            out.extend(_leaves(v))
+        elif k in LEAF_KEYS:
+            out.append({k: v})
+    return out
+
+
+def leaf_progress(win_condition: Any, ctx: WinContext) -> list[dict]:
+    rows: list[dict] = []
+    for leaf in _leaves(win_condition):
+        (name, target), = leaf.items()
+        satisfied = bool(evaluate({name: target}, ctx))
+        cur_fn = _CURRENT.get(name)
+        if cur_fn is None or not isinstance(target, (int, float)):
+            rows.append({
+                "name": name, "target": target, "current": None,
+                "ratio": 1.0 if satisfied else 0.0, "satisfied": satisfied,
+            })
+            continue
+        cur = cur_fn(ctx)
+        tgt = float(target)
+        if name in _LTE:
+            ratio = 1.0 if satisfied else _clamp(tgt / cur) if cur else 1.0
+        else:
+            ratio = 1.0 if satisfied else (_clamp(cur / tgt) if tgt else 0.0)
+        rows.append({
+            "name": name, "target": target, "current": cur,
+            "ratio": round(ratio, 4), "satisfied": satisfied,
+        })
+    return rows
+
+
+def reward_vector(signals: Any) -> dict[str, float]:
+    """Normalized, scenario-agnostic cumulative progress vector.
+
+    All dimensions are 0..1 and (because the underlying signals are
+    monotone over an episode) non-decreasing turn over turn — a true
+    cumulative tracker, comparable across scenarios.
+    """
+    cash = getattr(signals, "cash", 0) + getattr(signals, "resources", 0)
+    kills = getattr(signals, "units_killed", 0)
+    seen = len(getattr(signals, "enemies_seen_ids", []) or []) + len(
+        getattr(signals, "enemy_buildings_seen_ids", []) or []
+    )
+    return {
+        "economy": round(_clamp(cash / 10000.0), 4),
+        "military": round(_clamp(kills / 10.0), 4),
+        "territory": round(_clamp(
+            getattr(signals, "explored_percent", 0.0) / 100.0), 4),
+        "scouting": round(_clamp(seen / 10.0), 4),
+        "objective": 0.0,  # filled by turn_goal once the win cond is known
+    }
+
+
+def turn_goal(win_condition: Any, ctx: WinContext) -> dict:
+    leaves = leaf_progress(win_condition, ctx)
+    won = bool(evaluate(win_condition, ctx))
+    obj = 1.0 if won else (
+        round(sum(r["ratio"] for r in leaves) / len(leaves), 4)
+        if leaves else 0.0
+    )
+    rv = reward_vector(ctx.signals)
+    rv["objective"] = obj
+    return {
+        "leaves": leaves,
+        "reward_vector": rv,
+        "objective_progress": obj,
+        "won": won,
+    }

openra_bench/playback.py

@@ -58,6 +58,7 @@ class Playback:
         signals: Any,
         minimap_png_b64: str | None = None,
         interrupt: str | None = None,
+        goal: dict | None = None,
     ) -> None:
         self._n += 1
         if minimap_png_b64:
@@ -88,6 +89,9 @@ class Playback:
             ),
             "units": render_state.get("units_summary", []),
             "enemies": render_state.get("enemy_summary", []),
+            # Per-turn goal tracker: win-condition leaf progress AND the
+            # normalized cumulative reward vector, side by side.
+            "goal": goal or {},
         }
         self._turns_fh.write(json.dumps(_jsonable(rec)) + "\n")
         self._turns_fh.flush()

openra_bench/playback_view.py

@@ -0,0 +1,116 @@
+"""Playback viewer (pipeline step 7, read side).
+
+`load_episode` is the pure data layer — it reassembles a saved episode
+(manifest + turns + transcript + per-turn minimap PNG paths) into one
+dict, and is unit-tested. `render_streamlit` is a thin optional UI on
+top of it (mirrors the training repo's Streamlit pipeline viewer):
+
+    pip install streamlit
+    streamlit run scripts/view_playback.py -- <playback_root>
+
+Per turn it shows: the minimap, the user briefing, the model's
+reasoning + tool calls, the signal snapshot, and the goal tracker
+(win-condition leaf bars + the cumulative reward vector).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+def load_episode(ep_dir: str | Path) -> dict:
+    """Reassemble one ``seed<N>`` episode folder. Tolerant of a still-
+    running episode (missing files become empty)."""
+    d = Path(ep_dir)
+    manifest = _read_json(d / "manifest.json", {})
+    messages = _read_json(d / "messages.json", [])
+    turns = []
+    tj = d / "turns.jsonl"
+    if tj.exists():
+        for line in tj.read_text().splitlines():
+            line = line.strip()
+            if line:
+                rec = json.loads(line)
+                png = d / f"minimap_turn{rec.get('turn', 0):03d}.png"
+                rec["minimap_png"] = str(png) if png.exists() else None
+                turns.append(rec)
+    return {"dir": str(d), "manifest": manifest, "turns": turns,
+            "messages": messages}
+
+
+def find_episodes(root: str | Path) -> list[Path]:
+    """All ``.../seed<N>`` episode dirs under a playback root."""
+    return sorted(
+        p.parent for p in Path(root).glob("**/manifest.json")
+    ) or sorted(Path(root).glob("**/seed*"))
+
+
+def _read_json(p: Path, default):
+    try:
+        return json.loads(p.read_text())
+    except Exception:  # noqa: BLE001 — partial/running episode
+        return default
+
+
+def _assistant_turns(messages: list[dict]) -> list[dict]:
+    return [m for m in messages if m.get("role") == "assistant"]
+
+
+def render_streamlit(root: str) -> None:  # pragma: no cover - UI glue
+    import streamlit as st
+
+    st.set_page_config(page_title="OpenRA-Bench playback", layout="wide")
+    eps = find_episodes(root)
+    if not eps:
+        st.error(f"no episodes under {root}")
+        return
+    pick = st.sidebar.selectbox(
+        "episode", eps, format_func=lambda p: f"{p.parent.name}/{p.name}"
+    )
+    ep = load_episode(pick)
+    m = ep["manifest"]
+    st.title(f"{m.get('scenario', pick)} — {m.get('outcome', '?')}")
+    st.caption(
+        f"turns {m.get('turns')}/{m.get('max_turns')} · "
+        f"capability {m.get('capability')} · seed {m.get('seed')}"
+    )
+
+    asst = _assistant_turns(ep["messages"])
+    for i, t in enumerate(ep["turns"]):
+        with st.expander(
+            f"turn {t['turn']} · tick {t.get('tick')}"
+            + (f" · ⚡{t['interrupt']}" if t.get("interrupt") else ""),
+            expanded=(i == 0),
+        ):
+            left, right = st.columns([1, 1])
+            with left:
+                if t.get("minimap_png"):
+                    st.image(t["minimap_png"], caption="minimap")
+                elif t.get("ascii_minimap"):
+                    st.code(t["ascii_minimap"])
+                st.json(t.get("signals", {}))
+            with right:
+                a = asst[i] if i < len(asst) else {}
+                if a.get("reasoning"):
+                    st.markdown("**reasoning**")
+                    st.write(a["reasoning"])
+                st.markdown("**commands**")
+                st.code("\n".join(t.get("commands", [])) or "(none)")
+                g = t.get("goal") or {}
+                if g:
+                    st.markdown(
+                        f"**objective progress: "
+                        f"{g.get('objective_progress', 0):.0%}**"
+                        + ("  ✅ won" if g.get("won") else "")
+                    )
+                    for leaf in g.get("leaves", []):
+                        st.progress(
+                            min(1.0, float(leaf.get("ratio", 0.0))),
+                            text=f"{leaf['name']} "
+                            f"{leaf.get('current')}/{leaf.get('target')}",
+                        )
+                    rv = g.get("reward_vector", {})
+                    st.caption("reward vector: " + "  ".join(
+                        f"{k}={v:.2f}" for k, v in rv.items()
+                    ))

openra_bench/providers.py

@@ -81,6 +81,7 @@ class ChatReply:
 
     text: str
     tool_calls: list[dict]  # [{"name": str, "arguments": dict}]
+    reasoning: str = ""  # chain-of-thought, when the model/provider emits it
     raw: dict = field(default_factory=dict)
 
 
@@ -105,7 +106,7 @@ class OpenAICompatibleProvider(ChatProvider):
         }
         body: dict[str, Any] = {
             "model": cfg.model,
-            "messages": messages,
+            "messages": self._wire_messages(messages),
             "temperature": cfg.temperature,
             "max_tokens": cfg.max_tokens,
         }
@@ -118,7 +119,28 @@ class OpenAICompatibleProvider(ChatProvider):
             json=body,
         )
         resp.raise_for_status()
-        data = resp.json()
+        return self._reply_from_data(resp.json())
+
+    # Keys the OpenAI Chat Completions wire format accepts per message.
+    # `history` carries extra playback-only keys (notably "reasoning");
+    # those must never be posted back or strict servers (vLLM) 400.
+    _WIRE_KEYS = frozenset(
+        {"role", "content", "name", "tool_calls", "tool_call_id"}
+    )
+
+    @staticmethod
+    def _wire_messages(messages: list[dict]) -> list[dict]:
+        """Pure: project each message onto OpenAI-legal keys only."""
+        return [
+            {k: v for k, v in m.items() if k in OpenAICompatibleProvider._WIRE_KEYS}
+            for m in messages
+        ]
+
+    @staticmethod
+    def _reply_from_data(data: dict) -> ChatReply:
+        """Pure parse of a Chat Completions response, including the
+        provider-specific reasoning channel (vLLM/DeepSeek emit
+        `reasoning_content`; OpenRouter/others a flat `reasoning`)."""
         msg = data["choices"][0]["message"]
         calls: list[dict] = []
         for tc in msg.get("tool_calls") or []:
@@ -130,7 +152,17 @@ class OpenAICompatibleProvider(ChatProvider):
                 except json.JSONDecodeError:
                     args = {}
             calls.append({"name": fn.get("name", ""), "arguments": args})
-        return ChatReply(text=msg.get("content") or "", tool_calls=calls, raw=data)
+        rc = msg.get("reasoning_content") or msg.get("reasoning") or ""
+        if isinstance(rc, list):  # some providers chunk it
+            rc = "".join(
+                p.get("text", "") if isinstance(p, dict) else str(p) for p in rc
+            )
+        return ChatReply(
+            text=msg.get("content") or "",
+            tool_calls=calls,
+            reasoning=str(rc),
+            raw=data,
+        )
 
     def close(self) -> None:
         self._client.close()

scripts/view_playback.py

@@ -0,0 +1,11 @@
+"""Streamlit entrypoint for the playback viewer.
+
+    pip install streamlit
+    streamlit run scripts/view_playback.py -- <playback_root>
+"""
+import sys
+
+from openra_bench.playback_view import render_streamlit
+
+if __name__ == "__main__":
+    render_streamlit(sys.argv[1] if len(sys.argv) > 1 else "playback")

tests/test_playback_completeness.py

@@ -0,0 +1,164 @@
+"""Playback must capture everything needed to *replay reasoning*.
+
+Three gaps closed here, each asserted end to end:
+
+1. Model reasoning/thinking is extracted from the provider raw reply
+   and persisted on the assistant turn (so messages.json contains the
+   chain-of-thought, like the training traces) — while the *wire*
+   messages stay OpenAI-clean (no non-standard keys leak back).
+2. Every turn records a goal tracker: per-win-condition-leaf progress
+   AND a normalized cumulative reward vector, side by side.
+3. The saved episode is loadable by the viewer's data layer.
+"""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from openra_bench.goal_tracker import leaf_progress, reward_vector, turn_goal
+from openra_bench.providers import ChatReply, OpenAICompatibleProvider, ProviderConfig
+from openra_bench.scenarios.win_conditions import WinContext
+
+
+# ---- 1. reasoning capture --------------------------------------------------
+
+
+class _Sig:
+    explored_percent = 40.0
+    enemies_seen_ids = ["e1", "e2"]
+    enemy_buildings_seen_ids: list[str] = []
+    units_killed = 3
+    units_lost = 1
+    game_tick = 1200
+    cash = 1400
+    resources = 600
+    power_provided = 100
+    power_drained = 60
+    own_building_types = {"powr"}
+    own_buildings = [("powr", 5, 5)]
+
+
+def test_provider_extracts_reasoning_content():
+    cfg = ProviderConfig(provider="vllm")
+    prov = OpenAICompatibleProvider(cfg)
+    data = {
+        "choices": [
+            {
+                "message": {
+                    "content": "moving out",
+                    "reasoning_content": "enemy is east, scout first",
+                    "tool_calls": [],
+                }
+            }
+        ]
+    }
+    reply = prov._reply_from_data(data)  # pure parse, no network
+    assert isinstance(reply, ChatReply)
+    assert reply.reasoning == "enemy is east, scout first"
+    # also accepts the OpenRouter-style flat "reasoning" key
+    data["choices"][0]["message"] = {"content": "x", "reasoning": "because"}
+    assert prov._reply_from_data(data).reasoning == "because"
+
+
+def test_wire_messages_strip_nonstandard_keys():
+    # history may carry reasoning for playback; it must NOT be posted.
+    msgs = [
+        {"role": "system", "content": "s"},
+        {"role": "assistant", "content": "a", "reasoning": "secret",
+         "tool_calls": [{"id": "c0", "type": "function",
+                         "function": {"name": "observe", "arguments": {}}}]},
+        {"role": "tool", "tool_call_id": "c0", "content": "ok"},
+    ]
+    wire = OpenAICompatibleProvider._wire_messages(msgs)
+    assert all("reasoning" not in m for m in wire)
+    # structural keys survive
+    assert wire[1]["tool_calls"][0]["function"]["name"] == "observe"
+    assert wire[2]["tool_call_id"] == "c0"
+    # original untouched (pure)
+    assert msgs[1]["reasoning"] == "secret"
+
+
+def test_agent_records_reasoning_on_assistant_turn():
+    from openra_bench.agent import ModelAgent
+
+    class P:
+        def complete(self, messages, tools):
+            return ChatReply(text="go", tool_calls=[], reasoning="I think east")
+
+    a = ModelAgent(ProviderConfig(vision=False), allowed_tools=["observe"],
+                   provider=P())
+    a.agent_fn({"minimap": "", "units_summary": [], "enemy_summary": []},
+               type("C", (), {"observe": staticmethod(lambda: "obs")}))
+    asst = [m for m in a.history if m["role"] == "assistant"][-1]
+    assert asst["reasoning"] == "I think east"
+
+
+# ---- 2. per-turn goal tracker ---------------------------------------------
+
+
+def _ctx():
+    return WinContext(signals=_Sig(), render_state={"units_summary": []})
+
+
+def test_leaf_progress_reports_ratio_per_predicate():
+    wc = {"all_of": [{"units_killed_gte": 5}, {"explored_pct_gte": 80},
+                     {"cash_gte": 1000}]}
+    prog = leaf_progress(wc, _ctx())
+    by = {p["name"]: p for p in prog}
+    assert by["units_killed_gte"]["current"] == 3
+    assert by["units_killed_gte"]["target"] == 5
+    assert by["units_killed_gte"]["ratio"] == pytest.approx(0.6)
+    assert by["cash_gte"]["satisfied"] is True  # 1400 >= 1000
+    assert by["explored_pct_gte"]["satisfied"] is False
+
+
+def test_reward_vector_is_normalized_and_cumulative():
+    rv = reward_vector(_Sig())
+    for k in ("economy", "military", "territory", "scouting", "objective"):
+        assert k in rv and 0.0 <= rv[k] <= 1.0
+    assert rv["territory"] == pytest.approx(0.40)  # 40% explored
+
+
+def test_turn_goal_bundles_predicates_and_vector_side_by_side():
+    g = turn_goal({"units_killed_gte": 5}, _ctx())
+    assert "leaves" in g and "reward_vector" in g
+    assert "objective_progress" in g and 0.0 <= g["objective_progress"] <= 1.0
+    assert g["won"] is False  # 3 < 5
+
+
+# ---- 3. end-to-end persisted + loadable -----------------------------------
+
+
+def test_playback_round_trip_has_reasoning_and_goal(tmp_path):
+    from openra_bench.playback import Playback
+    from openra_bench.playback_view import load_episode
+
+    pb = Playback(tmp_path, "pack:easy", 7)
+
+    class S(_Sig):
+        pass
+
+    pb.record_turn(
+        1, {"minimap": "..", "units_summary": [], "enemy_summary": []},
+        ["Command::Observe"], S(), None,
+        goal=turn_goal({"units_killed_gte": 5}, _ctx()),
+    )
+    pb.write_messages([
+        {"role": "system", "content": "s"},
+        {"role": "user", "content": "briefing"},
+        {"role": "assistant", "content": "go", "reasoning": "scout east",
+         "tool_calls": []},
+    ])
+    pb.finalize({"scenario": "pack:easy", "outcome": "loss", "turns": 1})
+
+    ep = load_episode(pb.dir)
+    assert ep["manifest"]["outcome"] == "loss"
+    assert ep["turns"][0]["goal"]["won"] is False
+    assert ep["turns"][0]["goal"]["leaves"][0]["name"] == "units_killed_gte"
+    asst = [m for m in ep["messages"] if m["role"] == "assistant"][0]
+    assert asst["reasoning"] == "scout east"
+    # turns.jsonl is valid JSONL
+    raw = (pb.dir / "turns.jsonl").read_text().strip().splitlines()
+    assert json.loads(raw[0])["goal"]["reward_vector"]["territory"] >= 0.0