| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Sequence |
|
|
| import numpy as np |
|
|
| from engine.brain import BrainSignals, ReplayBrain |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| DEFAULT_MODAL_DUMP = ROOT / "eval" / "probe_dump_modal.npz" |
| DEFAULT_AGENTS = ("investor_a", "investor_b") |
|
|
|
|
| @dataclass(frozen=True) |
| class TraceStep: |
| signals_by_agent: dict[str, BrainSignals] |
| floor_holder: str = "human" |
| note: str = "" |
|
|
|
|
| @dataclass(frozen=True) |
| class SignalTrace: |
| name: str |
| steps: list[TraceStep] |
| expected: dict[str, object] = field(default_factory=dict) |
|
|
| @property |
| def agent_ids(self) -> list[str]: |
| return list(self.steps[0].signals_by_agent.keys()) if self.steps else [] |
|
|
|
|
| def synthetic_traces() -> list[SignalTrace]: |
| return [ |
| clean_monologue_trace(), |
| surprising_claim_trace(), |
| topic_shift_trace(), |
| rambling_pause_trace(), |
| ] |
|
|
|
|
| def load_modal_trace(path: str | Path = DEFAULT_MODAL_DUMP, agent_ids: Sequence[str] = DEFAULT_AGENTS) -> SignalTrace: |
| path = Path(path) |
| replay = ReplayBrain.from_npz(path, name="modal_probe_replay") |
| samples = list(replay) |
| steps: list[TraceStep] = [] |
| for index, sample in enumerate(samples): |
| signals: dict[str, BrainSignals] = {} |
| for offset, agent_id in enumerate(agent_ids): |
| readiness = max(0.0, sample.readiness - 0.08 * offset) |
| signals[agent_id] = BrainSignals( |
| surprise=sample.surprise, |
| hidden=sample.hidden, |
| readiness=readiness, |
| p_end=sample.p_end, |
| ) |
| steps.append(TraceStep(signals_by_agent=signals, note=f"modal step {index + 1}")) |
| return SignalTrace( |
| name="modal_probe_replay", |
| steps=steps, |
| expected={"source": str(path), "n_steps": len(steps)}, |
| ) |
|
|
|
|
| def all_traces() -> list[SignalTrace]: |
| traces = synthetic_traces() |
| if DEFAULT_MODAL_DUMP.exists(): |
| traces.append(load_modal_trace(DEFAULT_MODAL_DUMP)) |
| return traces |
|
|
|
|
| def clean_monologue_trace() -> SignalTrace: |
| base = _unit([1, 0, 0, 0, 0, 0, 0, 0]) |
| steps = _trace_from_series( |
| name="clean_monologue_take_floor", |
| hidden_series=[_nudge(base, index, 0.01) for index in range(6)], |
| surprise=[2.0, 2.1, 2.0, 2.2, 2.1, 2.45], |
| readiness_a=[0.25, 0.30, 0.35, 0.40, 0.45, 0.78], |
| readiness_b=[0.20, 0.25, 0.30, 0.32, 0.35, 0.52], |
| p_end=[0.02, 0.03, 0.04, 0.05, 0.10, 0.95], |
| notes=["setup", "details", "still talking", "more context", "closing", "turn complete"], |
| ) |
| return SignalTrace(name="clean_monologue_take_floor", steps=steps, expected={"take_floor_step": 6}) |
|
|
|
|
| def surprising_claim_trace() -> SignalTrace: |
| base = _unit([1, 0, 0, 0, 0, 0, 0, 0]) |
| steps = _trace_from_series( |
| name="surprising_claim_interrupt", |
| hidden_series=[_nudge(base, index, 0.01) for index in range(6)], |
| surprise=[2.0, 2.1, 7.3, 3.0, 2.4, 2.2], |
| readiness_a=[0.30, 0.42, 0.92, 0.82, 0.55, 0.45], |
| readiness_b=[0.25, 0.32, 0.35, 0.36, 0.38, 0.40], |
| p_end=[0.02, 0.04, 0.18, 0.20, 0.30, 0.55], |
| notes=["setup", "build", "wild claim", "continues", "settles", "handoff"], |
| ) |
| return SignalTrace(name="surprising_claim_interrupt", steps=steps, expected={"interrupt_step": 3}) |
|
|
|
|
| def topic_shift_trace() -> SignalTrace: |
| topic_a = _unit([1, 0, 0, 0, 0, 0, 0, 0]) |
| topic_b = _unit([0, 1, 0, 0, 0, 0, 0, 0]) |
| steps = _trace_from_series( |
| name="topic_shift_backchannel", |
| hidden_series=[ |
| _nudge(topic_a, 0, 0.01), |
| _nudge(topic_a, 1, 0.01), |
| _nudge(topic_a, 2, 0.01), |
| _nudge(topic_a, 3, 0.01), |
| _nudge(topic_b, 4, 0.01), |
| _nudge(topic_b, 5, 0.01), |
| _nudge(topic_b, 6, 0.01), |
| ], |
| surprise=[2.0, 2.1, 2.0, 2.1, 2.12, 2.08, 2.06], |
| readiness_a=[0.25, 0.30, 0.35, 0.38, 0.50, 0.45, 0.40], |
| readiness_b=[0.20, 0.25, 0.28, 0.30, 0.35, 0.34, 0.34], |
| p_end=[0.02, 0.04, 0.06, 0.08, 0.22, 0.25, 0.28], |
| notes=["topic a", "topic a", "topic a", "topic a", "topic b shift", "topic b", "topic b"], |
| ) |
| return SignalTrace(name="topic_shift_backchannel", steps=steps, expected={"backchannel_step": 5}) |
|
|
|
|
| def rambling_pause_trace() -> SignalTrace: |
| base = _unit([0, 0, 1, 0, 0, 0, 0, 0]) |
| steps = _trace_from_series( |
| name="rambling_pause_take_floor", |
| hidden_series=[_nudge(base, index, 0.015) for index in range(7)], |
| surprise=[2.1, 2.0, 2.2, 2.1, 2.0, 2.1, 2.45], |
| readiness_a=[0.20, 0.25, 0.30, 0.38, 0.42, 0.42, 0.80], |
| readiness_b=[0.18, 0.22, 0.25, 0.30, 0.32, 0.35, 0.50], |
| p_end=[0.02, 0.05, 0.10, 0.38, 0.42, 0.44, 0.96], |
| notes=["start", "ramble", "ramble", "awkward pause", "holds", "still unsure", "complete"], |
| ) |
| return SignalTrace(name="rambling_pause_take_floor", steps=steps, expected={"hold_step": 4, "take_floor_step": 7}) |
|
|
|
|
| def _trace_from_series( |
| *, |
| name: str, |
| hidden_series: Sequence[np.ndarray], |
| surprise: Sequence[float], |
| readiness_a: Sequence[float], |
| readiness_b: Sequence[float], |
| p_end: Sequence[float], |
| notes: Sequence[str], |
| ) -> list[TraceStep]: |
| steps: list[TraceStep] = [] |
| for index, hidden in enumerate(hidden_series): |
| signals = { |
| "investor_a": BrainSignals( |
| surprise=surprise[index], |
| hidden=hidden, |
| readiness=readiness_a[index], |
| p_end=p_end[index], |
| ), |
| "investor_b": BrainSignals( |
| surprise=max(0.0, surprise[index] - 0.2), |
| hidden=hidden, |
| readiness=readiness_b[index], |
| p_end=p_end[index], |
| ), |
| } |
| steps.append(TraceStep(signals_by_agent=signals, note=notes[index])) |
| return steps |
|
|
|
|
| def _unit(values: Sequence[float]) -> np.ndarray: |
| vector = np.asarray(values, dtype=np.float32) |
| norm = np.linalg.norm(vector) |
| if norm == 0: |
| return vector |
| return vector / norm |
|
|
|
|
| def _nudge(base: np.ndarray, index: int, scale: float) -> np.ndarray: |
| vector = base.astype(np.float32, copy=True) |
| vector[(index % (len(vector) - 1)) + 1] += scale |
| return _unit(vector) |
|
|
