ReAct two-pass per day so model sees current-day tool results

Phase A (discovery): model emits tool_calls only; we dispatch them via
env._dispatch_tool (read-only) and capture results.
Phase B (planning): same observation + fresh tool results, model emits
scheduled_actions. env.step then runs those actions only (tool_calls
already executed, not double-dispatched).

Both phases logged separately to plots/io_log.jsonl with /A and /B tags
so we can verify discovery actually fires.

Made-with: Cursor

training/train_grpo.ipynb

@@ -443,19 +443,17 @@
         "\n",
         "SYSTEM_PROMPT = _SYSTEM_BASE + textwrap.dedent(\"\"\"\n",
         "\n",
-        "TOOL POLICY (tool_calls cost nothing — call them aggressively):\n",
-        "- The observation tells you ONLY your account stats. Audience peak hours, segment\n",
-        "  affinities, trending topics/tags and competitor schedules are NOT given. You must\n",
-        "  discover them via tool_calls and read them from `Tool results` next turn.\n",
-        "- days_elapsed == 0  -> call EVERY discovery tool you might need, e.g.:\n",
-        "    {\"name\": \"query_trends\",        \"arguments\": {\"niche\": \"<TOPIC_CATEGORIES key>\"}}\n",
-        "    {\"name\": \"query_audience\",      \"arguments\": {\"segment_id\": \"young_professionals\"}}\n",
-        "    {\"name\": \"query_audience\",      \"arguments\": {\"segment_id\": \"students\"}}\n",
-        "    {\"name\": \"query_creator_pool\",  \"arguments\": {}}\n",
-        "    {\"name\": \"query_competitor\",    \"arguments\": {\"competitor_id\": \"niche_expert\", \"window_days\": 7}}\n",
-        "- days_elapsed >= 1  -> before scheduling posts, call:\n",
-        "    {\"name\": \"predict_engagement\",  \"arguments\": {\"scheduled_actions\": [...]}}\n",
-        "  and any query_* whose result is missing from `Tool results`.\"\"\")\n",
+        "TWO-PHASE FLOW (each day has two turns — same observation, two responses):\n",
+        "PHASE A — DISCOVERY: respond with {\"tool_calls\": [...]} only. Tools cost nothing,\n",
+        "  call as many query_* / predict_engagement / draft_review as useful. Their results\n",
+        "  are dispatched immediately and shown to you in PHASE B of the SAME day.\n",
+        "PHASE B — PLANNING: respond with {\"scheduled_actions\": [...], \"notes\": \"...\"}\n",
+        "  using the freshly returned Tool results.\n",
+        "Audience peak hours, segment affinities, trends, competitor schedules are NOT in\n",
+        "the observation — discover them in PHASE A. Useful PHASE-A starter set:\n",
+        "  query_trends(niche), query_audience(segment_id), query_creator_pool(),\n",
+        "  query_competitor(competitor_id, window_days), and on later days also\n",
+        "  predict_engagement(scheduled_actions=[...candidate plan...]).\"\"\")\n",
         "SYSTEM_PROMPT_EVAL = SYSTEM_PROMPT\n",
         "SYSTEM_PROMPT_TRAIN = SYSTEM_PROMPT\n",
         "\n",
@@ -469,18 +467,18 @@
         "        signals_str = (f\"Signals: watch={signals.watch_time:.3f} \"\n",
         "                       f\"sends={signals.sends_per_reach:.3f} \"\n",
         "                       f\"saves={signals.saves:.3f}\\n\")\n",
-    "    tool_str = \"\"\n",
-    "    for tr in getattr(obs, \"tool_results\", []):\n",
-    "        if tr.success:\n",
-    "            tool_str += f\"  {tr.name}: {json.dumps(tr.data)}\\n\"\n",
-    "    if not tool_str:\n",
-    "        tool_str = \"  (none — call query_* tools to discover)\\n\"\n",
-    "    return (f\"Day: {day_name} | days_elapsed={obs.days_elapsed}\\n\"\n",
-    "            f\"Energy: {obs.creator_energy:.2f} | Followers: {obs.follower_count}\\n\"\n",
-    "            f\"Engagement: {obs.engagement_rate:.3f} | Queue: {obs.content_queue_size}\\n\"\n",
-    "            f\"{signals_str}\"\n",
-    "            f\"Tool results:\\n{tool_str}\"\n",
-    "            f\"Plan today's actions (JSON only):\")\n",
+        "    tool_str = \"\"\n",
+        "    for tr in getattr(obs, \"tool_results\", []):\n",
+        "        if tr.success:\n",
+        "            tool_str += f\"  {tr.name}: {json.dumps(tr.data)}\\n\"\n",
+        "    if not tool_str:\n",
+        "        tool_str = \"  (none — call query_* tools to discover)\\n\"\n",
+        "    return (f\"Day: {day_name} | days_elapsed={obs.days_elapsed}\\n\"\n",
+        "            f\"Energy: {obs.creator_energy:.2f} | Followers: {obs.follower_count}\\n\"\n",
+        "            f\"Engagement: {obs.engagement_rate:.3f} | Queue: {obs.content_queue_size}\\n\"\n",
+        "            f\"{signals_str}\"\n",
+        "            f\"Tool results:\\n{tool_str}\"\n",
+        "            f\"Plan today's actions (JSON only):\")\n",
         "\n",
         "\n",
         "def is_well_formed_response(text):\n",
@@ -578,49 +576,96 @@
         "        f.write(json.dumps(rec) + \"\\n\")\n",
         "\n",
         "\n",
-        "def run_llm_episodes_batched(mdl, tok, tasks_seeds, verbose=True, eval=False, system=None, log_tag=None):\n",
-        "    \"\"\"Run N episodes in parallel. tasks_seeds: list of (task, seed). One batched generate per day.\"\"\"\n",
-        "    sys_prompt = system or (SYSTEM_PROMPT_EVAL if eval else SYSTEM_PROMPT_TRAIN)\n",
-        "    n = len(tasks_seeds)\n",
-        "    envs = [ViraltestEnvironment() for _ in range(n)]\n",
-        "    obss = [envs[i].reset(task=t, seed=s) for i, (t, s) in enumerate(tasks_seeds)]\n",
-        "    rewards = [[] for _ in range(n)]\n",
-        "    energies = [[obs.creator_energy] for obs in obss]\n",
-        "    pairs = [[] for _ in range(n)]\n",
-        "    done_mask = [obs.done for obs in obss]\n",
-        "    rest_action = ViraltestAction(scheduled_actions=[])\n",
-        "\n",
-        "    for day in range(1, TASK_HORIZON + 1):\n",
-        "        active = [i for i in range(n) if not done_mask[i] and obss[i].creator_energy > 0.25]\n",
-        "        rest = [i for i in range(n) if not done_mask[i] and obss[i].creator_energy <= 0.25]\n",
-        "        if not active and not rest:\n",
-        "            break\n",
-        "\n",
-        "        actions_by_idx = {i: rest_action for i in rest}\n",
-        "        if active:\n",
-        "            prompts = [format_obs(obss[i]) for i in active]\n",
-        "            chats = [_build_chat(sys_prompt, p) for p in prompts]\n",
-        "            texts = [tok.apply_chat_template(c, tokenize=False, add_generation_prompt=True) for c in chats]\n",
-        "            resps, ptok = _batched_generate(mdl, tok, texts, eval=eval)\n",
-        "            if verbose:\n",
-        "                print(f\"  D{day:2d}: batch={len(active)} rest={len(rest)} prompt_tok={ptok}\")\n",
-        "            for j, i in enumerate(active):\n",
-        "                actions_by_idx[i] = parse_model_output(resps[j])\n",
-        "                pairs[i].append({\"prompt\": prompts[j], \"response\": resps[j],\n",
-        "                                 \"step\": len(rewards[i])})\n",
-        "                if log_tag is not None:\n",
-        "                    t, s = tasks_seeds[i]\n",
-        "                    _log_io(log_tag, i, day, t, s, prompts[j], resps[j])\n",
-        "\n",
-        "        for i in range(n):\n",
-        "            if done_mask[i] or i not in actions_by_idx:\n",
-        "                continue\n",
-        "            obss[i] = envs[i].step(actions_by_idx[i])\n",
-        "            r = obss[i].reward or 0.0\n",
-        "            rewards[i].append(r)\n",
-        "            energies[i].append(obss[i].creator_energy)\n",
-        "            if obss[i].done:\n",
-        "                done_mask[i] = True\n",
+    "DISCOVERY_SUFFIX = \"\\n\\nPHASE A (DISCOVERY): respond with JSON {\\\"tool_calls\\\": [...]} only.\"\n",
+    "PLANNING_SUFFIX = \"\\n\\nPHASE B (PLANNING): respond with JSON {\\\"scheduled_actions\\\": [...], \\\"notes\\\": \\\"...\\\"} using the fresh Tool results above.\"\n",
+    "\n",
+    "\n",
+    "def _parse_tool_calls_only(text):\n",
+    "    return parse_model_output(text).tool_calls\n",
+    "\n",
+    "\n",
+    "def _parse_actions_only(text):\n",
+    "    a = parse_model_output(text)\n",
+    "    return ViraltestAction(tool_calls=[], scheduled_actions=a.scheduled_actions, notes=a.notes)\n",
+    "\n",
+    "\n",
+    "def _format_fresh_results(fresh):\n",
+    "    if not fresh:\n",
+    "        return \"\"\n",
+    "    out = \"Fresh tool results (PHASE A):\\n\"\n",
+    "    for tr in fresh:\n",
+    "        if tr.success:\n",
+    "            out += f\"  {tr.name}: {json.dumps(tr.data)}\\n\"\n",
+    "        else:\n",
+    "            out += f\"  {tr.name}: ERROR {tr.error}\\n\"\n",
+    "    return out\n",
+    "\n",
+    "\n",
+    "def run_llm_episodes_batched(mdl, tok, tasks_seeds, verbose=True, eval=False, system=None, log_tag=None):\n",
+    "    \"\"\"Run N episodes in parallel. ReAct two-pass: discovery -> dispatch -> planning.\"\"\"\n",
+    "    sys_prompt = system or (SYSTEM_PROMPT_EVAL if eval else SYSTEM_PROMPT_TRAIN)\n",
+    "    n = len(tasks_seeds)\n",
+    "    envs = [ViraltestEnvironment() for _ in range(n)]\n",
+    "    obss = [envs[i].reset(task=t, seed=s) for i, (t, s) in enumerate(tasks_seeds)]\n",
+    "    rewards = [[] for _ in range(n)]\n",
+    "    energies = [[obs.creator_energy] for obs in obss]\n",
+    "    pairs = [[] for _ in range(n)]\n",
+    "    done_mask = [obs.done for obs in obss]\n",
+    "    rest_action = ViraltestAction(scheduled_actions=[])\n",
+    "\n",
+    "    def _gen(prompts):\n",
+    "        chats = [_build_chat(sys_prompt, p) for p in prompts]\n",
+    "        texts = [tok.apply_chat_template(c, tokenize=False, add_generation_prompt=True) for c in chats]\n",
+    "        return _batched_generate(mdl, tok, texts, eval=eval)\n",
+    "\n",
+    "    for day in range(1, TASK_HORIZON + 1):\n",
+    "        active = [i for i in range(n) if not done_mask[i] and obss[i].creator_energy > 0.25]\n",
+    "        rest = [i for i in range(n) if not done_mask[i] and obss[i].creator_energy <= 0.25]\n",
+    "        if not active and not rest:\n",
+    "            break\n",
+    "\n",
+    "        actions_by_idx = {i: rest_action for i in rest}\n",
+    "        if active:\n",
+    "            base_prompts = [format_obs(obss[i]) for i in active]\n",
+    "\n",
+    "            disc_prompts = [p + DISCOVERY_SUFFIX for p in base_prompts]\n",
+    "            disc_resps, ptok = _gen(disc_prompts)\n",
+    "            if verbose:\n",
+    "                print(f\"  D{day:2d}A: batch={len(active)} rest={len(rest)} prompt_tok={ptok}\")\n",
+    "\n",
+    "            fresh_per_active = []\n",
+    "            for j, i in enumerate(active):\n",
+    "                tcs = _parse_tool_calls_only(disc_resps[j])\n",
+    "                fresh_per_active.append([envs[i]._dispatch_tool(tc) for tc in tcs])\n",
+    "                pairs[i].append({\"prompt\": disc_prompts[j], \"response\": disc_resps[j],\n",
+    "                                 \"step\": len(rewards[i]), \"phase\": \"A\"})\n",
+    "                if log_tag is not None:\n",
+    "                    t, s = tasks_seeds[i]\n",
+    "                    _log_io(f\"{log_tag}/A\", i, day, t, s, disc_prompts[j], disc_resps[j])\n",
+    "\n",
+    "            plan_prompts = [base_prompts[j] + \"\\n\" + _format_fresh_results(fresh_per_active[j]) + PLANNING_SUFFIX\n",
+    "                            for j in range(len(active))]\n",
+    "            plan_resps, ptok2 = _gen(plan_prompts)\n",
+    "            if verbose:\n",
+    "                print(f\"  D{day:2d}B: batch={len(active)} prompt_tok={ptok2}\")\n",
+    "\n",
+    "            for j, i in enumerate(active):\n",
+    "                actions_by_idx[i] = _parse_actions_only(plan_resps[j])\n",
+    "                pairs[i].append({\"prompt\": plan_prompts[j], \"response\": plan_resps[j],\n",
+    "                                 \"step\": len(rewards[i]), \"phase\": \"B\"})\n",
+    "                if log_tag is not None:\n",
+    "                    t, s = tasks_seeds[i]\n",
+    "                    _log_io(f\"{log_tag}/B\", i, day, t, s, plan_prompts[j], plan_resps[j])\n",
+    "\n",
+    "        for i in range(n):\n",
+    "            if done_mask[i] or i not in actions_by_idx:\n",
+    "                continue\n",
+    "            obss[i] = envs[i].step(actions_by_idx[i])\n",
+    "            r = obss[i].reward or 0.0\n",
+    "            rewards[i].append(r)\n",
+    "            energies[i].append(obss[i].creator_energy)\n",
+    "            if obss[i].done:\n",
+    "                done_mask[i] = True\n",
         "\n",
         "    GAMMA, TERMINAL_W = 0.95, 5.0\n",
         "    results = []\n",