Update drp/simulate.py

drp/simulate.py

@@ -1,13 +1,12 @@
 import os
 import random
 import uuid
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 from .bundle import write_bundle_zip
 
 
 def _env_fingerprint() -> Dict[str, Any]:
-    # Keep lightweight; users can expand this in real exporters
     return {
         "python": os.environ.get("PYTHON_VERSION") or "unknown",
         "space": os.environ.get("SPACE_ID") or os.environ.get("HF_SPACE_ID") or "unknown",
@@ -22,6 +21,7 @@ def make_demo_bundle_zip(out_path: str, *, seed: int, chaos: float, label: str)
     """
     Creates a synthetic agent timeline with controlled randomness.
     'chaos' increases divergence probability.
+    Also includes reward_total in state snapshots to demonstrate reward delta.
     """
     rng = random.Random(seed)
     run_id = f"demo-{label}-{uuid.uuid4().hex[:8]}"
@@ -30,37 +30,50 @@ def make_demo_bundle_zip(out_path: str, *, seed: int, chaos: float, label: str)
 
     events: List[Dict[str, Any]] = []
     memory: Dict[str, Any] = {"goal": "reach_target", "notes": []}
+    reward_total = 0.0
+
+    # Optional replay link template (demo)
+    replay = {
+        "base_url": "https://example.com/replay",
+        "pattern": "/?run_id={run_id}&i={i}",
+    }
 
     for i in range(40):
-        # planning
         action = rng.choice(["scan", "move", "ask_tool", "write_memory"])
         if rng.random() < chaos:
             action = rng.choice(["scan", "move", "ask_tool", "write_memory", "panic"])
 
-        events.append(_mk_event("plan_step", f"t{i}.plan", {"action": action, "score": rng.random()}))
+        score = rng.random()
+        events.append(_mk_event("plan_step", f"t{i}.plan", {"action": action, "score": score}))
 
         if action == "ask_tool":
             q = rng.choice(["price", "status", "latency", "risk"])
             events.append(_mk_event("tool_call", f"t{i}.tool_call", {"tool": "mock_api", "query": q}))
-            # tool sometimes flakes
             if rng.random() < (0.15 + chaos * 0.2):
                 events.append(_mk_event("tool_result", f"t{i}.tool_result", {"ok": False, "error": "timeout"}))
+                reward_total -= 0.5
             else:
                 val = rng.randint(1, 100)
                 events.append(_mk_event("tool_result", f"t{i}.tool_result", {"ok": True, "value": val}))
+                reward_total += 0.2
         elif action == "write_memory":
             note = rng.choice(["cached", "retry", "validated", "unsafe", "needs_review"])
             memory["notes"].append(note)
             events.append(_mk_event("memory_write", f"t{i}.mem", {"write": {"notes": list(memory["notes"])}}))
+            reward_total += 0.05
         elif action == "panic":
             events.append(_mk_event("guardrail", f"t{i}.guardrail", {"blocked": True, "reason": "anomaly"}))
-            events.append(_mk_event("state_snapshot", f"t{i}.state", {"memory": memory, "mode": "halt"}))
+            events.append(
+                _mk_event("state_snapshot", f"t{i}.state", {"memory": memory, "mode": "halt", "reward_total": reward_total})
+            )
             break
         else:
-            # move / scan influences a synthetic "state"
-            events.append(_mk_event("state_snapshot", f"t{i}.state", {"x": rng.randint(0, 9), "y": rng.randint(0, 9), "memory": memory}))
+            # scan/move
+            x = rng.randint(0, 9)
+            y = rng.randint(0, 9)
+            reward_total += 0.01
+            events.append(_mk_event("state_snapshot", f"t{i}.state", {"x": x, "y": y, "memory": memory, "reward_total": reward_total}))
 
-        # llm sample (synthetic text)
         txt = rng.choice(
             [
                 "Proceed with caution.",
@@ -88,6 +101,7 @@ def make_demo_bundle_zip(out_path: str, *, seed: int, chaos: float, label: str)
         model_id=model_id,
         env_fingerprint=_env_fingerprint(),
         events_payloads=events,
+        replay=replay,
     )
 
 
@@ -101,8 +115,7 @@ def fork_patch_bundle(
     patch_payload_json: Optional[Dict[str, Any]] = None,
 ) -> str:
     """
-    Simple “what-if” fork: take an existing bundle and patch a single event
-    (kind/step/payload) then re-hash-chain and re-emit as a new run.
+    Counterfactual workflow: patch an event at index N, re-hash-chain into a new bundle.
     """
     from .bundle import load_bundle, write_bundle_zip
 
@@ -136,4 +149,6 @@ def fork_patch_bundle(
         model_id=b.manifest.get("model_id", "unknown"),
         env_fingerprint=b.manifest.get("env", {}),
         events_payloads=payloads,
+        replay=b.manifest.get("replay"),
+        run_url=b.manifest.get("run_url"),
     )