pounteradds

server/viraltest_environment.py

@@ -387,6 +387,8 @@ class ViraltestEnvironment(Environment):
         self._hours_since_sleep = 2
         self._sleep_debt = 0.0
 
+        self._reward_mode = "combined"
+
     def _load_competitors(self) -> List[CompetitorState]:
         archetypes = _COMPETITORS_DATA.get("archetypes", [])
         return [
@@ -1136,6 +1138,8 @@ class ViraltestEnvironment(Environment):
 
         self._shift_label = kwargs.get("shift_label")
         self._chain_id = kwargs.get("episode_chain_id")
+        mode = kwargs.get("reward_mode", "combined")
+        self._reward_mode = mode if mode in ("timing", "content", "combined") else "combined"
 
         if self._chain_id and self._chain_id in _BRAND_STORE:
             brand = _BRAND_STORE[self._chain_id]
@@ -1439,20 +1443,29 @@ class ViraltestEnvironment(Environment):
     # ----- reward -----
 
     def _compute_hourly_reward(self, sa: ScheduledAction, engagement: float) -> float:
-        eng_component = min(1.0, engagement / 2.0) * 0.3
+        if self._reward_mode == "timing":
+            return self._compute_timing_reward(sa, engagement)
+        if self._reward_mode == "content":
+            return self._compute_content_reward(sa, engagement)
+        return self._compute_combined_reward(sa, engagement)
 
+    def _energy_component(self) -> float:
         prev_energy = self._energy_history[-2] if len(self._energy_history) >= 2 else 1.0
         energy_delta = self._energy - prev_energy
-        energy_component = max(0.0, min(1.0, (energy_delta + 0.3) / 0.6)) * 0.15
+        return max(0.0, min(1.0, (energy_delta + 0.3) / 0.6))
 
+    def _consistency_score(self) -> float:
         day_posts = self._posts_per_day.get(self._day, 0)
         if 1 <= day_posts <= 2:
-            consistency = 1.0
-        elif day_posts == 0 or day_posts == 3:
-            consistency = 0.5
-        else:
-            consistency = 0.0
-        consistency_component = consistency * 0.15
+            return 1.0
+        if day_posts == 0 or day_posts == 3:
+            return 0.5
+        return 0.0
+
+    def _compute_combined_reward(self, sa: ScheduledAction, engagement: float) -> float:
+        eng_component = min(1.0, engagement / 2.0) * 0.3
+        energy_component = self._energy_component() * 0.15
+        consistency_component = self._consistency_score() * 0.15
 
         tag_component = 0.0
         if sa.action_type == "post" and sa.tags:
@@ -1474,22 +1487,54 @@ class ViraltestEnvironment(Environment):
         )
         return max(0.0, min(1.0, raw))
 
-    def _compute_rest_reward(self) -> float:
-        prev_energy = self._energy_history[-2] if len(self._energy_history) >= 2 else 1.0
-        energy_delta = self._energy - prev_energy
-        energy_component = max(0.0, min(1.0, (energy_delta + 0.3) / 0.6)) * 0.15
+    def _compute_timing_reward(self, sa: ScheduledAction, engagement: float) -> float:
+        is_post = sa.action_type == "post"
+        peak_hour_mult = 1.3 if is_post and self._get_hour_multiplier() >= 1.2 else 1.0
+        trending_topic_mult = 1.5 if is_post and self._is_topic_trending(sa.topic) else 1.0
+        eng_component = min(1.0, engagement / 2.0) * 0.40 * trending_topic_mult * peak_hour_mult
 
-        day_posts = self._posts_per_day.get(self._day, 0)
-        if 1 <= day_posts <= 2:
-            consistency = 1.0
-        elif day_posts == 0 or day_posts == 3:
-            consistency = 0.5
-        else:
-            consistency = 0.0
-        consistency_component = consistency * 0.15
+        peak_bonus = min(1.0, self._get_hour_multiplier() / 1.3) if is_post else 0.0
+        peak_component = peak_bonus * 0.20
+
+        energy_component = self._energy_component() * 0.20
+        consistency_component = self._consistency_score() * 0.20
+        burnout_penalty = 0.1 if self._energy < 0.2 else 0.0
 
+        raw = eng_component + peak_component + energy_component + consistency_component - burnout_penalty
+        return max(0.0, min(1.0, raw))
+
+    def _compute_content_reward(self, sa: ScheduledAction, engagement: float) -> float:
+        is_post = sa.action_type == "post"
+        trending_topic_mult = 1.5 if is_post and self._is_topic_trending(sa.topic) else 1.0
+        eng_component = min(1.0, engagement / 2.0) * 0.20 * trending_topic_mult
+
+        tag_component = 0.0
+        if is_post and sa.tags:
+            trending_match = sum(1 for t in sa.tags if t.lower() in self._trending_tags) / 5.0
+            tag_component = min(1.0, trending_match + 0.3) * 0.25
+
+        comp_component = 0.0
+        if is_post:
+            diff = self._calc_competitor_diff(sa.topic)
+            comp_component = min(1.0, diff / 1.3) * 0.25
+
+        variety_component = 0.0
+        intent_component = 0.0
+        if is_post:
+            variety_component = min(1.0, len(self._unique_content_types) / 4.0) * 0.15
+            intent_component = (0.15 if sa.intent in INTENT_MULTIPLIER else 0.0)
+
+        burnout_penalty = 0.05 if self._energy < 0.2 else 0.0
+        raw = eng_component + tag_component + comp_component + variety_component + intent_component - burnout_penalty
+        return max(0.0, min(1.0, raw))
+
+    def _compute_rest_reward(self) -> float:
+        energy_component = self._energy_component() * 0.15
+        consistency_component = self._consistency_score() * 0.15
         burnout_penalty = 0.1 if self._energy < 0.2 else 0.0
         raw = energy_component + consistency_component - burnout_penalty
+        if self._reward_mode == "content":
+            raw *= 0.5
         return max(0.0, min(1.0, raw))
 
     def _advance_time(self) -> None:
@@ -1700,6 +1745,13 @@ class ViraltestEnvironment(Environment):
         return max(0.0, min(1.0, raw))
 
 
+def get_peak_hours(day_of_week: int, top_k: int = 2) -> List[int]:
+    row = _HEATMAP_GRID.get(day_of_week % 7, [])
+    if not row:
+        return []
+    return sorted(range(len(row)), key=lambda h: row[h], reverse=True)[:top_k]
+
+
 def _topic_overlap(topic_a: str, topic_b: str) -> bool:
     words_a = set(topic_a.split())
     words_b = set(topic_b.split())

training/train_grpo.ipynb

@@ -25,9 +25,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 1: Install dependencies (quote versions — zsh treats `>` as redirect otherwise)\n",
         "!pip install -q torch torchvision torchaudio\n",
@@ -36,13 +34,13 @@
         "!pip install -q \"typing_extensions>=4.13.0\" pydantic httpx\n",
         "!pip install -q \"openenv-core[core]>=0.2.2\"\n",
         "!pip install -q flash-attn --no-build-isolation || echo \"flash-attn install skipped; will use sdpa\""
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 2: Resolve repo path (Colab: fresh clone. Local: auto-detect project root)\n",
         "import os\n",
@@ -118,13 +116,13 @@
         "print(f\"Branch: {REPO_BRANCH}\")\n",
         "print(f\"Commit: {commit}\")\n",
         "print(f\"Plots dir: {PLOTS_DIR}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 3: Imports (with runtime validation)\n",
         "import json, random, time, textwrap, copy, os, sys\n",
@@ -156,7 +154,7 @@
         "from models import ScheduledAction, ToolCall, ViraltestAction\n",
         "from server.viraltest_environment import (\n",
         "    ViraltestEnvironment, TAG_POOL, TASK_HORIZON,\n",
-        "    TOPIC_CATEGORIES,\n",
+        "    TOPIC_CATEGORIES, get_peak_hours,\n",
         ")\n",
         "\n",
         "ALL_TOPICS = [t for topics in TOPIC_CATEGORIES.values() for t in topics]\n",
@@ -178,7 +176,9 @@
         "import ast\n",
         "ast.parse(\"def _t(x: int) -> str: return f'{x}'\")\n",
         "print(\"OK: ast.parse (syntax check)\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -191,9 +191,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 4: Define heuristic agents + episode runner\n",
         "_rng = random.Random(42)\n",
@@ -269,13 +267,13 @@
         "            \"rewards\": rewards, \"energies\": energies}\n",
         "\n",
         "print(\"Agents and episode runner defined.\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 5: Run baselines (safe)\n",
         "print(\"Running heuristic baselines (5 agents × 3 tasks)...\")\n",
@@ -310,13 +308,13 @@
         "for name in BASELINE_AGENTS:\n",
         "    scores = [baseline_results[name][t][\"grader_score\"] for t in TASKS]\n",
         "    print(f\"{name:<14s} {scores[0]:>10.4f} {scores[1]:>12.4f} {scores[2]:>14.4f} {sum(scores)/3:>8.4f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 6: Baseline plots\n",
         "fig, axes = plt.subplots(1, 3, figsize=(16, 5), sharey=True)\n",
@@ -334,7 +332,9 @@
         "fig.tight_layout()\n",
         "fig.savefig(f\"{PLOTS_DIR}/baseline_leaderboard.png\", dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -347,9 +347,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 7: Load model (Qwen2.5-3B bf16 on CUDA + flash-attn-2; fp16/fp32 fallback)\n",
         "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
@@ -393,13 +391,13 @@
         "print(f\"Model loaded. dtype={next(model.parameters()).dtype} device={next(model.parameters()).device}\")\n",
         "if torch.cuda.is_available():\n",
         "    print(f\"CUDA memory: {torch.cuda.memory_allocated()/1e9:.2f} GB\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 8: LLM agent functions\n",
         "_SYSTEM_BASE = textwrap.dedent(\"\"\"\\\n",
@@ -454,6 +452,16 @@
         "SYSTEM_PROMPT_EVAL = SYSTEM_PROMPT\n",
         "SYSTEM_PROMPT_TRAIN = SYSTEM_PROMPT\n",
         "\n",
+        "SYSTEM_PROMPT_TIMING = SYSTEM_PROMPT + textwrap.dedent(\"\"\"\n",
+        "\n",
+        "FOCUS: optimise WHEN to post. Identify peak hours for the audience (use query_audience / query_trends).\n",
+        "2 posts/day at peak hours beats 4 posts at random hours.\"\"\")\n",
+        "\n",
+        "SYSTEM_PROMPT_CONTENT = SYSTEM_PROMPT + textwrap.dedent(\"\"\"\n",
+        "\n",
+        "FOCUS: optimise WHAT to post. Vary content_type and intent across the week,\n",
+        "pick differentiated topics, exploit trending tags.\"\"\")\n",
+        "\n",
         "\n",
         "_DAY_NAMES = [\"Mon\", \"Tue\", \"Wed\", \"Thu\", \"Fri\", \"Sat\", \"Sun\"]\n",
         "\n",
@@ -472,7 +480,7 @@
         "    return out\n",
         "\n",
         "\n",
-        "def format_obs(obs, history=None):\n",
+        "def format_obs(obs, history=None, extra_hint=None):\n",
         "    day_name = _DAY_NAMES[obs.day_of_week] if 0 <= obs.day_of_week < 7 else \"?\"\n",
         "    signals_str = \"\"\n",
         "    signals = getattr(obs, \"engagement_signals\", None)\n",
@@ -486,12 +494,14 @@
         "            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",
+        "    hint_str = f\"Coach hint: today's peak hours are {extra_hint}.\\n\" if extra_hint else \"\"\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\"{_format_history(history)}\"\n",
         "            f\"Tool results:\\n{tool_str}\"\n",
+        "            f\"{hint_str}\"\n",
         "            f\"Plan today's actions (JSON only):\")\n",
         "\n",
         "\n",
@@ -615,12 +625,13 @@
         "    return out\n",
         "\n",
         "\n",
-        "def run_llm_episodes_batched(mdl, tok, tasks_seeds, verbose=True, eval=False, system=None, log_tag=None):\n",
+        "def run_llm_episodes_batched(mdl, tok, tasks_seeds, verbose=True, eval=False, system=None,\n",
+        "                              log_tag=None, hint_peak_hours=False, reward_mode=\"combined\"):\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",
+        "    obss = [envs[i].reset(task=t, seed=s, reward_mode=reward_mode) 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",
@@ -641,7 +652,12 @@
         "\n",
         "        actions_by_idx = {i: rest_action for i in rest}\n",
         "        if active:\n",
-        "            base_prompts = [format_obs(obss[i], histories[i]) for i in active]\n",
+        "            def _hint_for(i):\n",
+        "                if not hint_peak_hours:\n",
+        "                    return None\n",
+        "                hrs = get_peak_hours(obss[i].day_of_week, top_k=2)\n",
+        "                return \", \".join(f\"{h:02d}:00\" for h in hrs) if hrs else None\n",
+        "            base_prompts = [format_obs(obss[i], histories[i], extra_hint=_hint_for(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",
@@ -716,7 +732,9 @@
         "\n",
         "\n",
         "print(\"LLM agent functions defined (batched).\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -729,9 +747,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 9: Run untrained model (batched: all 3 tasks in parallel envs)\n",
         "print(\"Running UNTRAINED base model on all tasks (batched)...\")\n",
@@ -745,7 +761,9 @@
         "print(f\"BEFORE TRAINING (took {time.time()-t0:.1f}s):\")\n",
         "for t in TASKS:\n",
         "    print(f\"  {t}: grader={before_results[t]['grader_score']:.4f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -764,9 +782,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 10: Attach LoRA adapter\n",
         "from peft import LoraConfig, get_peft_model, TaskType\n",
@@ -780,118 +796,144 @@
         "model.enable_input_require_grads()\n",
         "peft_model = get_peft_model(model, lora_config)\n",
         "peft_model.print_trainable_parameters()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
-        "# Cell 11: Training loop\n",
+        "# Cell 11: Two-phase training loop (timing -> content)\n",
+        "# Each phase: 3 rounds (round 0 = hardcoded peak-hours hint, rounds 1-2 = normal prompt).\n",
+        "# Adapter persisted to ./checkpoints/phaseN_adapter/ between phases.\n",
         "from trl import SFTTrainer, SFTConfig\n",
         "from datasets import Dataset\n",
         "\n",
-        "NUM_ROUNDS = 2\n",
         "EPISODES_PER_ROUND = 6\n",
-        "QUALITY_FLOOR = 0.0  # 0 = always run SFT on positive-advantage samples\n",
+        "ROUNDS_PER_PHASE = 3\n",
+        "QUALITY_FLOOR = 0.0\n",
+        "\n",
+        "PHASES = [\n",
+        "    {\"name\": \"phase1_timing\",  \"reward_mode\": \"timing\",  \"system\": SYSTEM_PROMPT_TIMING},\n",
+        "    {\"name\": \"phase2_content\", \"reward_mode\": \"content\", \"system\": SYSTEM_PROMPT_CONTENT},\n",
+        "]\n",
         "\n",
         "training_log = {\n",
-        "    \"round\": [], \"avg_episode_reward\": [], \"max_episode_reward\": [],\n",
-        "    \"min_episode_reward\": [], \"avg_grader\": [], \"max_grader\": [],\n",
+        "    \"phase\": [], \"round\": [], \"global_step\": [], \"use_hint\": [],\n",
+        "    \"avg_episode_reward\": [], \"max_episode_reward\": [], \"min_episode_reward\": [],\n",
+        "    \"avg_grader\": [], \"max_grader\": [],\n",
         "    \"n_training_samples\": [], \"train_loss\": [],\n",
         "}\n",
         "\n",
         "t_start = time.time()\n",
-        "\n",
-        "for round_idx in range(1, NUM_ROUNDS + 1):\n",
-        "    print(f\"\\n{'=' * 60}\")\n",
-        "    print(f\"TRAINING ROUND {round_idx}/{NUM_ROUNDS}\")\n",
-        "    print(f\"{'=' * 60}\")\n",
-        "\n",
-        "    peft_model.eval()\n",
-        "    tasks_seeds = [(TASKS[ep % len(TASKS)], 42 + (round_idx - 1) * 100 + ep) for ep in range(EPISODES_PER_ROUND)]\n",
-        "    t_roll = time.time()\n",
-        "    results = run_llm_episodes_batched(peft_model, tokenizer, tasks_seeds, verbose=False,\n",
-        "                                       eval=False, system=SYSTEM_PROMPT_TRAIN,\n",
-        "                                       log_tag=f\"train_round{round_idx}\")\n",
-        "    print(f\"  Rollouts: {len(results)} eps × {TASK_HORIZON} days in {time.time()-t_roll:.1f}s\")\n",
-        "\n",
-        "    all_pairs, episode_rewards, episode_graders = [], [], []\n",
-        "    for ep, result in enumerate(results):\n",
-        "        ep_reward = result[\"total_reward\"] + 2.0 * result[\"grader_score\"]\n",
-        "        episode_rewards.append(ep_reward)\n",
-        "        episode_graders.append(result[\"grader_score\"])\n",
-        "        kept = 0\n",
-        "        for pr in result[\"pairs\"]:\n",
-        "            if not is_well_formed_response(pr[\"response\"]):\n",
-        "                continue\n",
-        "            text = (f\"<|im_start|>system\\n{SYSTEM_PROMPT_TRAIN}<|im_end|>\\n\"\n",
-        "                    f\"<|im_start|>user\\n{pr['prompt']}<|im_end|>\\n\"\n",
-        "                    f\"<|im_start|>assistant\\n{pr['response']}<|im_end|>\")\n",
-        "            all_pairs.append({\"text\": text, \"reward\": pr[\"return\"]})\n",
-        "            kept += 1\n",
-        "        print(f\"  ep {ep+1}/{EPISODES_PER_ROUND}: {result['task'].split('_')[-1]:>11s} \"\n",
-        "              f\"grader={result['grader_score']:.4f} reward={ep_reward:.3f} kept={kept}/{len(result['pairs'])}\")\n",
-        "\n",
-        "    avg_r = float(np.mean(episode_rewards))\n",
-        "    avg_g = float(np.mean(episode_graders))\n",
-        "    max_g = float(max(episode_graders))\n",
-        "    print(f\"  Avg reward={avg_r:.3f} Avg grader={avg_g:.4f} max_grader={max_g:.4f} | pairs={len(all_pairs)}\")\n",
-        "    if not all_pairs:\n",
-        "        print(\"  WARNING: 0 well-formed pairs collected; skipping SFT.\")\n",
-        "        continue\n",
-        "    if max_g < QUALITY_FLOOR:\n",
-        "        print(f\"  SKIP SFT: no episode beat quality_floor={QUALITY_FLOOR:.2f}\")\n",
-        "        continue\n",
-        "\n",
-        "    rets = np.array([p[\"reward\"] for p in all_pairs], dtype=float)\n",
-        "    adv = (rets - rets.mean()) / (rets.std() + 1e-6)\n",
-        "    filtered = [p for p, a in zip(all_pairs, adv) if a > 0.0]\n",
-        "    if not filtered:\n",
-        "        print(\"  SKIP SFT: zero positive-advantage samples\")\n",
-        "        continue\n",
-        "    print(f\"  Kept {len(filtered)}/{len(all_pairs)} positive-advantage samples\")\n",
-        "\n",
-        "    dataset = Dataset.from_list([{\"text\": p[\"text\"]} for p in filtered])\n",
-        "\n",
-        "    # SFT training (real gradient updates)\n",
-        "    sft_config = SFTConfig(\n",
-        "        output_dir=f\"./checkpoints/round_{round_idx}\",\n",
-        "        num_train_epochs=1,\n",
-        "        per_device_train_batch_size=2,\n",
-        "        gradient_accumulation_steps=4,\n",
-        "        learning_rate=5e-6,\n",
-        "        warmup_steps=5,\n",
-        "        logging_steps=1,\n",
-        "        save_strategy=\"no\",\n",
-        "        max_length=2048,\n",
-        "        bf16=True,\n",
-        "        report_to=\"none\",\n",
-        "    )\n",
-        "\n",
-        "    peft_model.train()\n",
-        "    trainer = SFTTrainer(\n",
-        "        model=peft_model, processing_class=tokenizer,\n",
-        "        train_dataset=dataset, args=sft_config,\n",
-        "    )\n",
-        "    train_result = trainer.train()\n",
-        "    loss = train_result.training_loss\n",
-        "    print(f\"  Training loss: {loss:.4f}\")\n",
-        "\n",
-        "    training_log[\"round\"].append(round_idx)\n",
-        "    training_log[\"avg_episode_reward\"].append(round(float(avg_r), 3))\n",
-        "    training_log[\"max_episode_reward\"].append(round(float(max(episode_rewards)), 3))\n",
-        "    training_log[\"min_episode_reward\"].append(round(float(min(episode_rewards)), 3))\n",
-        "    training_log[\"avg_grader\"].append(round(float(avg_g), 4))\n",
-        "    training_log[\"max_grader\"].append(round(float(max(episode_graders)), 4))\n",
-        "    training_log[\"n_training_samples\"].append(len(filtered))\n",
-        "    training_log[\"train_loss\"].append(round(loss, 4))\n",
+        "global_step = 0\n",
+        "\n",
+        "for phase in PHASES:\n",
+        "    phase_name = phase[\"name\"]\n",
+        "    sys_prompt = phase[\"system\"]\n",
+        "    reward_mode = phase[\"reward_mode\"]\n",
+        "    print(f\"\\n{'#' * 60}\\n# PHASE {phase_name} (reward_mode={reward_mode})\\n{'#' * 60}\")\n",
+        "\n",
+        "    for round_idx in range(ROUNDS_PER_PHASE):\n",
+        "        use_hint = (round_idx == 0)\n",
+        "        print(f\"\\n{'=' * 60}\\n{phase_name} | ROUND {round_idx+1}/{ROUNDS_PER_PHASE} | hint={use_hint}\\n{'=' * 60}\")\n",
+        "\n",
+        "        peft_model.eval()\n",
+        "        tasks_seeds = [(TASKS[ep % len(TASKS)], 42 + ep + round_idx * 10) for ep in range(EPISODES_PER_ROUND)]\n",
+        "        t_roll = time.time()\n",
+        "        results = run_llm_episodes_batched(\n",
+        "            peft_model, tokenizer, tasks_seeds, verbose=False, eval=False,\n",
+        "            system=sys_prompt, hint_peak_hours=use_hint, reward_mode=reward_mode,\n",
+        "            log_tag=f\"{phase_name}_r{round_idx}\",\n",
+        "        )\n",
+        "        print(f\"  Rollouts: {len(results)} eps × {TASK_HORIZON} days in {time.time()-t_roll:.1f}s\")\n",
+        "\n",
+        "        all_pairs, episode_rewards, episode_graders = [], [], []\n",
+        "        for ep, result in enumerate(results):\n",
+        "            ep_reward = result[\"total_reward\"] + 2.0 * result[\"grader_score\"]\n",
+        "            episode_rewards.append(ep_reward)\n",
+        "            episode_graders.append(result[\"grader_score\"])\n",
+        "            kept = 0\n",
+        "            for pr in result[\"pairs\"]:\n",
+        "                if not is_well_formed_response(pr[\"response\"]):\n",
+        "                    continue\n",
+        "                text = (f\"<|im_start|>system\\n{sys_prompt}<|im_end|>\\n\"\n",
+        "                        f\"<|im_start|>user\\n{pr['prompt']}<|im_end|>\\n\"\n",
+        "                        f\"<|im_start|>assistant\\n{pr['response']}<|im_end|>\")\n",
+        "                all_pairs.append({\"text\": text, \"reward\": pr[\"return\"]})\n",
+        "                kept += 1\n",
+        "            print(f\"  ep {ep+1}/{EPISODES_PER_ROUND}: {result['task'].split('_')[-1]:>11s} \"\n",
+        "                  f\"grader={result['grader_score']:.4f} reward={ep_reward:.3f} kept={kept}/{len(result['pairs'])}\")\n",
+        "\n",
+        "        avg_r = float(np.mean(episode_rewards))\n",
+        "        avg_g = float(np.mean(episode_graders))\n",
+        "        max_g = float(max(episode_graders))\n",
+        "        print(f\"  Avg reward={avg_r:.3f} Avg grader={avg_g:.4f} max_grader={max_g:.4f} | pairs={len(all_pairs)}\")\n",
+        "\n",
+        "        loss = float(\"nan\")\n",
+        "        n_filtered = 0\n",
+        "        if not all_pairs:\n",
+        "            print(\"  WARNING: 0 well-formed pairs collected; skipping SFT.\")\n",
+        "        elif max_g < QUALITY_FLOOR:\n",
+        "            print(f\"  SKIP SFT: no episode beat quality_floor={QUALITY_FLOOR:.2f}\")\n",
+        "        else:\n",
+        "            rets = np.array([p[\"reward\"] for p in all_pairs], dtype=float)\n",
+        "            adv = (rets - rets.mean()) / (rets.std() + 1e-6)\n",
+        "            filtered = [p for p, a in zip(all_pairs, adv) if a > 0.0]\n",
+        "            if not filtered:\n",
+        "                print(\"  SKIP SFT: zero positive-advantage samples\")\n",
+        "            else:\n",
+        "                n_filtered = len(filtered)\n",
+        "                print(f\"  Kept {n_filtered}/{len(all_pairs)} positive-advantage samples\")\n",
+        "                dataset = Dataset.from_list([{\"text\": p[\"text\"]} for p in filtered])\n",
+        "                sft_config = SFTConfig(\n",
+        "                    output_dir=f\"./checkpoints/{phase_name}_r{round_idx}\",\n",
+        "                    num_train_epochs=1,\n",
+        "                    per_device_train_batch_size=2,\n",
+        "                    gradient_accumulation_steps=4,\n",
+        "                    learning_rate=5e-6,\n",
+        "                    warmup_steps=5,\n",
+        "                    logging_steps=1,\n",
+        "                    save_strategy=\"no\",\n",
+        "                    max_length=2048,\n",
+        "                    bf16=True,\n",
+        "                    report_to=\"none\",\n",
+        "                )\n",
+        "                peft_model.train()\n",
+        "                trainer = SFTTrainer(\n",
+        "                    model=peft_model, processing_class=tokenizer,\n",
+        "                    train_dataset=dataset, args=sft_config,\n",
+        "                )\n",
+        "                train_result = trainer.train()\n",
+        "                loss = float(train_result.training_loss)\n",
+        "                print(f\"  Training loss: {loss:.4f}\")\n",
+        "\n",
+        "        global_step += 1\n",
+        "        training_log[\"phase\"].append(phase_name)\n",
+        "        training_log[\"round\"].append(round_idx + 1)\n",
+        "        training_log[\"global_step\"].append(global_step)\n",
+        "        training_log[\"use_hint\"].append(use_hint)\n",
+        "        training_log[\"avg_episode_reward\"].append(round(float(avg_r), 3))\n",
+        "        training_log[\"max_episode_reward\"].append(round(float(max(episode_rewards)), 3))\n",
+        "        training_log[\"min_episode_reward\"].append(round(float(min(episode_rewards)), 3))\n",
+        "        training_log[\"avg_grader\"].append(round(float(avg_g), 4))\n",
+        "        training_log[\"max_grader\"].append(round(float(max(episode_graders)), 4))\n",
+        "        training_log[\"n_training_samples\"].append(n_filtered)\n",
+        "        training_log[\"train_loss\"].append(round(loss, 4) if loss == loss else float(\"nan\"))\n",
+        "\n",
+        "    save_dir = f\"./checkpoints/{phase_name}_adapter\"\n",
+        "    os.makedirs(save_dir, exist_ok=True)\n",
+        "    peft_model.save_pretrained(save_dir)\n",
+        "    tokenizer.save_pretrained(save_dir)\n",
+        "    print(f\"\\n  Saved {phase_name} adapter -> {save_dir}\")\n",
         "\n",
         "elapsed = time.time() - t_start\n",
-        "print(f\"\\nTraining complete in {elapsed/60:.1f} min\")\n",
+        "print(f\"\\nTwo-phase training complete in {elapsed/60:.1f} min\")\n",
         "print(pd.DataFrame(training_log).to_string(index=False))"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -904,9 +946,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 12: Run trained model (batched)\n",
         "print(\"Running TRAINED model on all tasks (batched)...\")\n",
@@ -921,7 +961,9 @@
         "print(f\"AFTER TRAINING (took {time.time()-t0:.1f}s):\")\n",
         "for t in TASKS:\n",
         "    print(f\"  {t}: grader={after_results[t]['grader_score']:.4f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -932,37 +974,41 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
-        "# Cell 13: Training curves\n",
+        "# Cell 13: Training curves (two-phase)\n",
         "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
-        "rounds = training_log[\"round\"]\n",
+        "steps = training_log[\"global_step\"]\n",
+        "phases = training_log[\"phase\"]\n",
+        "phase1_end = max([s for s, p in zip(steps, phases) if p == \"phase1_timing\"], default=0)\n",
         "\n",
-        "axes[0].plot(rounds, training_log[\"avg_grader\"], 'o-', color='#2196F3', lw=2, label='Avg grader')\n",
-        "axes[0].fill_between(rounds, training_log[\"avg_grader\"],\n",
+        "axes[0].plot(steps, training_log[\"avg_grader\"], 'o-', color='#2196F3', lw=2, label='Avg grader')\n",
+        "axes[0].fill_between(steps, training_log[\"avg_grader\"],\n",
         "                     training_log[\"max_grader\"], alpha=0.2, color='#2196F3')\n",
-        "axes[0].set_xlabel('Round'); axes[0].set_ylabel('Grader Score')\n",
-        "axes[0].set_title('Grader Score Over Rounds', fontweight='bold')\n",
+        "if phase1_end > 0:\n",
+        "    axes[0].axvline(phase1_end + 0.5, color='gray', ls='--', alpha=0.6, label='phase split')\n",
+        "axes[0].set_xlabel('Global step'); axes[0].set_ylabel('Grader Score')\n",
+        "axes[0].set_title('Grader Score (timing -> content)', fontweight='bold')\n",
         "axes[0].legend(); axes[0].grid(True, alpha=0.3)\n",
         "\n",
-        "axes[1].plot(rounds, training_log[\"train_loss\"], 's-', color='#E53935', lw=2)\n",
-        "axes[1].set_xlabel('Round'); axes[1].set_ylabel('Loss')\n",
+        "axes[1].plot(steps, training_log[\"train_loss\"], 's-', color='#E53935', lw=2)\n",
+        "if phase1_end > 0:\n",
+        "    axes[1].axvline(phase1_end + 0.5, color='gray', ls='--', alpha=0.6)\n",
+        "axes[1].set_xlabel('Global step'); axes[1].set_ylabel('Loss')\n",
         "axes[1].set_title('Training Loss', fontweight='bold')\n",
         "axes[1].grid(True, alpha=0.3)\n",
         "\n",
-        "fig.suptitle('Viraltest v2 — LoRA Training Progress (Qwen 1.5B)', fontsize=14, fontweight='bold')\n",
+        "fig.suptitle('Viraltest v2 — Two-Phase LoRA Training (timing -> content)', fontsize=14, fontweight='bold')\n",
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/reward_curve.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 14: Before vs After\n",
         "task_labels = [t.replace('monthly_', '').title() for t in TASKS]\n",
@@ -992,13 +1038,13 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/before_after.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 15: Trajectory comparison\n",
         "fig, axes = plt.subplots(2, 3, figsize=(16, 8))\n",
@@ -1022,7 +1068,9 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/training_trajectories.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -1033,9 +1081,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 16: Final summary\n",
         "print(\"=\" * 67)\n",
@@ -1057,8 +1103,10 @@
         "\n",
         "summary = {\n",
         "    \"model\": MODEL_NAME,\n",
-        "    \"training\": \"LoRA SFT (real weight updates)\",\n",
-        "    \"rounds\": NUM_ROUNDS, \"episodes_per_round\": EPISODES_PER_ROUND,\n",
+        "    \"training\": \"Two-phase LoRA SFT (timing -> content) with hardcoded peak-hours hint on round 1 of each phase\",\n",
+        "    \"phases\": [p[\"name\"] for p in PHASES],\n",
+        "    \"rounds_per_phase\": ROUNDS_PER_PHASE,\n",
+        "    \"episodes_per_round\": EPISODES_PER_ROUND,\n",
         "    \"before\": {t: before_results[t][\"grader_score\"] for t in TASKS},\n",
         "    \"after\": {t: after_results[t][\"grader_score\"] for t in TASKS},\n",
         "    \"smart_heuristic\": {t: baseline_results[\"smart\"][t][\"grader_score\"] for t in TASKS},\n",
@@ -1072,13 +1120,13 @@
         "\n",
         "print(f\"\\nSaved to {PLOTS_DIR}/\")\n",
         "print(\"All results are from real LoRA weight updates on real environment runs.\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 17: Save adapter\n",
         "save_path = \"./viraltest_trained_adapter\"\n",
@@ -1086,7 +1134,9 @@
         "tokenizer.save_pretrained(save_path)\n",
         "print(f\"LoRA adapter saved to {save_path}\")\n",
         "print(\"Load with: PeftModel.from_pretrained(base_model, save_path)\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     }
   ],
   "metadata": {
@@ -1112,4 +1162,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 4
-}
+}