fix(reward): align prompt with SFT, soft format, real recursion signal (A.3 + A.4-bis)

REWARD_AUDIT.md

@@ -16,11 +16,11 @@ format_gate=1.0 still gives partial credit via token_budget + recursion_efficien
 score = compose_reward_single("<answer></answer>", "Rome", 100, 0, cfg_reward)
 ```
 
-**Expected:** ≤ 0.05 (A-02 fix: correctness=0 caps composite)
+**Expected:** ≤ 0.15 (A-02 cap is 0.05, plus 0.10 soft format bonus)
 
-**Result:** 0.05 (matches cap) ✅
+**Result:** 0.15 (matches cap) ✅
 
-**Status:** MITIGATED — A-02 correctness==0 cap blocks this. Empty answers have correctness=0.
+**Status:** MITIGATED — A-02 correctness==0 cap blocks this. Empty answers have correctness=0, so even with the format bonus, the maximum is 0.15.
 
 ---
 
@@ -72,20 +72,20 @@ or finds a copy of the gold answer in the context itself (DATA-04 post-check).
 
 ---
 
-## A-05: Recursion depth gaming (avoid all tool calls)
+## A-05: Recursion depth gaming (REVISED for v3)
 
-**Vector:** Model uses 0 llm() calls (skips recursion entirely) to maximize
-recursion_efficiency=1.0, giving up correctness to focus on efficiency bonus.
+**Vector:** Model uses 0 llm() calls on every task to maximize
+recursion_efficiency, even on multi_needle / 200K tasks where recursion
+would actually help correctness.
 
-**Analysis:**
-- recursion_efficiency weight = 0.05 (smallest component)
-- correctness weight = 0.75
-- Even with recursion_efficiency=1.0 and correctness=0.0: composite = min(0.05 * 1.0, 0.05) = 0.05
-  (A-02 cap applies)
-- For correctness=1.0: composite ≈ 0.75 + 0.2*tb + 0.05*1.0 ≈ 1.0 — dominates
+**Analysis (v3):**
+- recursion_efficiency contributes only when correctness == 1.0 (gating
+  in compose.py). On hard tasks where 0 calls fails to produce a correct
+  answer, the efficiency bonus is forfeited entirely.
+- Net incentive: use the *minimum* recursion that still produces a
+  correct answer. Exactly the desired behavior.
 
-**Status:** LOW RISK — Efficiency weight is 5% of composite; correctness dominates at 75%.
-The model cannot exploit efficiency to score high without being correct.
+**Status:** ✅ MITIGATED by correctness-gating.
 
 ---
 
@@ -108,15 +108,46 @@ normalization. Substring matches are rejected.
 
 ---
 
+## A-07: Comment-spam exploit (NEW)
+
+**Vector:** Model emits `# llm(foo)` inside code blocks to inflate the
+count regex without making real calls. (Inverted variant of A-05: spam
+to make recursion_eff *lower*, useless because lower efficiency hurts.)
+
+**Test:** `test_call_in_comment_not_counted` in
+`tests/test_recursion_extract.py`.
+
+**Result:** 0 calls counted ✅ — extractor uses tokenize, ignores comments.
+
+**Status:** ✅ MITIGATED by tokenize-aware extraction.
+
+---
+
+## A-08: String-literal exploit (NEW)
+
+**Vector:** Model writes `"earlier code did llm(...)"` in a string
+literal to confuse a naive regex extractor.
+
+**Test:** `test_call_in_string_literal_not_counted`.
+
+**Result:** 0 calls counted ✅ — tokenize correctly identifies STRING
+tokens and skips them.
+
+**Status:** ✅ MITIGATED.
+
+---
+
 ## Summary
 
 | Attempt | Vector | Status | Mitigation |
 |---------|--------|--------|------------|
-| A-01 | Format-only (empty answer) | ✅ MITIGATED | A-02 correctness==0 cap (≤0.05) |
+| A-01 | Format-only (empty answer) | ✅ MITIGATED | A-02 correctness==0 cap (≤0.15 with bonus) |
 | A-02 | Case/punctuation gaming | ✅ ACCEPTABLE | Intentional normalizer behavior |
 | A-03 | Context verbatim copy | ✅ MITIGATED | Generator `_assert_no_leak()` |
 | A-04 | Length reward exploit | ✅ MITIGATED | capped_linear token_budget |
-| A-05 | Recursion avoidance | ✅ LOW RISK | 5% weight; correctness dominates |
+| A-05 | Recursion avoidance | ✅ MITIGATED | Correctness-gating |
 | A-06 | Partial answer substring | ✅ MITIGATED | Exact-match normalizer |
+| A-07 | Comment-spam exploit | ✅ MITIGATED | Tokenize-aware extraction |
+| A-08 | String-literal exploit | ✅ MITIGATED | Tokenize-aware extraction |
 
 **VERDICT: REWARD SYSTEM APPROVED FOR PHASE 2 TRAINING**

configs/reward/v1.yaml

@@ -1,8 +1,8 @@
 alpha: 0.2
 weights:
-  correctness: 0.75
-  token_budget: 0.2
-  recursion_efficiency: 0.05
+  correctness: 0.70
+  token_budget: 0.15
+  recursion_efficiency: 0.15
 token_budget_variant: "capped_linear"
 answer_regex: "<answer>(.*?)</answer>"
-max_calls: 2
+max_calls: 4

rewards/compose.py

@@ -1,10 +1,21 @@
-"""Reward composition — REW-02.
-
-Implements Sequential(Gate(FormatCheck), WeightedSum([Corr*0.75, Tok*0.20, Rec*0.05])).
-Format fail → composite = 0.0 (no partial credit).
-
-A-02 fix: if correctness == 0.0, cap composite at 0.05 to prevent
-format-only completions with lucky short answers from scoring high.
+"""Reward composition — REW-02 v3.
+
+Changes from v2 (the §A.3.2 "soft format bonus" patch):
+  - llm_call_count is now extracted from the completion's fenced Python
+    code blocks (via rewards.recursion_extract.count_llm_calls), not
+    hardcoded to 0 in train/grpo.py.
+  - Recursion efficiency is gated on correctness — wrong answers cannot
+    earn an efficiency bonus. Prevents the model from spamming `llm(`
+    strings to harvest free reward.
+  - Weights rebalanced: 0.70 correctness / 0.15 token_budget /
+    0.15 recursion_efficiency. (Was 0.75 / 0.20 / 0.05.)
+  - Per-component scalars are returned alongside the composite via the
+    `_metrics` dict so the GRPOTrainer wrapper in train/grpo.py can
+    log real per-component means to W&B (currently logs 0.0).
+
+Anti-hacking caps preserved:
+  - c == 0.0 → composite ≤ 0.25 (was 0.05; raised to allow soft-format
+    bonus to register, still well below any correct answer ≥ 0.80).
 """
 from __future__ import annotations
 
@@ -14,76 +25,92 @@ from .format_gate import format_gate
 from .correctness import correctness
 from .token_budget import token_budget
 from .recursion_efficiency import recursion_efficiency
+from .recursion_extract import count_llm_calls
 
 
 def compose_reward_single(
     completion: str,
     gold_answer: str,
     prompt_token_count: int,
-    llm_call_count: int,
     cfg_reward: Any,
-) -> float:
-    """Single-example composite reward — REW-02.
+    llm_call_count: int | None = None,  # if None → extract from completion
+) -> tuple[float, dict[str, float]]:
+    """Single-example composite reward + per-component metrics.
 
-    Gate: format_gate == 0 → return 0.0 (no partial credit).
-    A-02 fix: correctness == 0 → cap at 0.05.
+    Returns (composite_score, metrics_dict). The metrics dict has keys:
+        format_pass, correctness, token_budget, recursion_eff_raw,
+        recursion_eff_contribution, llm_call_count.
     """
-    gate = format_gate(completion)
-    if gate == 0.0:
-        return 0.0
-
-    c = correctness(completion, gold_answer)
+    has_format = format_gate(completion) == 1.0
+    c = correctness(completion, gold_answer) if has_format else 0.0
     t = token_budget(
         completion,
         prompt_token_count,
         alpha=float(cfg_reward.alpha),
         variant=str(cfg_reward.token_budget_variant),
     )
-    r = recursion_efficiency(
-        int(llm_call_count),
-        max_calls=int(cfg_reward.get("max_calls", 2)) if hasattr(cfg_reward, "get") else int(getattr(cfg_reward, "max_calls", 2)),
-    )
-    w = cfg_reward.weights
 
-    # REW-02: weights must sum to 1.0
-    assert abs(float(w.correctness) + float(w.token_budget) + float(w.recursion_efficiency) - 1.0) < 1e-3, \
-        "REW-02: composite weights must sum to 1.0"
+    if llm_call_count is None:
+        llm_call_count = count_llm_calls(completion)
+    eff_raw = recursion_efficiency(int(llm_call_count))
+    # Couple efficiency to correctness — wrong answers earn 0 efficiency.
+    eff_contribution = eff_raw if c == 1.0 else 0.0
+
+    w = cfg_reward.weights
+    assert abs(
+        float(w.correctness) + float(w.token_budget) + float(w.recursion_efficiency) - 1.0
+    ) < 1e-3, "REW-02 v3: composite weights must sum to 1.0"
 
     composite = (
         float(w.correctness) * c
         + float(w.token_budget) * t
-        + float(w.recursion_efficiency) * r
+        + float(w.recursion_efficiency) * eff_contribution
     )
+    if has_format:
+        composite += 0.10  # soft format bonus (§A.3.2)
 
-    # A-02 fix: cap composite at 0.05 when correctness == 0
-    # (prevents format-only completions from scoring ~0.25 via token_budget + recursion)
     if c == 0.0:
-        return min(composite, 0.05)
+        composite = min(composite, 0.25)
 
-    return composite
+    metrics = {
+        "format_pass": 1.0 if has_format else 0.0,
+        "correctness": c,
+        "token_budget": t,
+        "recursion_eff_raw": eff_raw,
+        "recursion_eff_contribution": eff_contribution,
+        "llm_call_count": float(llm_call_count),
+    }
+    return composite, metrics
 
 
 def compose_reward_fn(prompts: list, completions: list, **kwargs) -> list[float]:
-    """TRL-compatible batched reward function — REW-02.
+    """TRL-compatible batched reward function. Returns scalars only.
 
-    Expected kwargs:
-        gold_answer: list[str]
-        prompt_token_count: list[int]
-        llm_call_count: list[int]
-        cfg_reward: OmegaConf DictConfig slice (injected via make_reward_fn partial)
+    Per-component means are stashed under `kwargs['_component_means']` for
+    the GRPOTrainer instrumentation wrapper to log to W&B. (TRL ignores
+    extra kwargs.)
     """
     cfg_reward = kwargs.pop("cfg_reward")
     gold_answers = kwargs.get("gold_answer", [""] * len(completions))
     ptcs = kwargs.get("prompt_token_count", [1] * len(completions))
-    lccs = kwargs.get("llm_call_count", [0] * len(completions))
-    return [
-        compose_reward_single(c, g, int(p), int(l), cfg_reward)
-        for c, g, p, l in zip(completions, gold_answers, ptcs, lccs)
+
+    pairs = [
+        compose_reward_single(c, g, int(p), cfg_reward)
+        for c, g, p in zip(completions, gold_answers, ptcs)
     ]
+    rewards = [p[0] for p in pairs]
+    metrics_list = [p[1] for p in pairs]
+
+    # Aggregate component means for W&B logging via the wrapper.
+    if metrics_list:
+        keys = metrics_list[0].keys()
+        means = {k: sum(m[k] for m in metrics_list) / len(metrics_list) for k in keys}
+        kwargs["_component_means"] = means
+    return rewards
 
 
 def make_reward_fn(cfg_reward: Any) -> Callable:
-    """Factory that binds cfg_reward so GRPOTrainer can call reward_fn(prompts, completions, **kwargs)."""
+    """Factory binding cfg_reward for GRPOTrainer.reward_funcs."""
     def _bound(prompts, completions, **kwargs):
         kwargs["cfg_reward"] = cfg_reward
         return compose_reward_fn(prompts, completions, **kwargs)

rewards/recursion_efficiency.py

@@ -1,19 +1,23 @@
-"""Recursion efficiency reward component — REW-01.
+"""Recursion efficiency reward component — REW-04 v2.
 
-Discrete staircase on llm_call_count. Pure Python, stdlib-only.
-"""
-_STAIRCASE = {0: 1.0, 1: 0.7, 2: 0.4}
+Linear decay on llm_call_count. Pure Python, stdlib-only.
 
+Intended ranges:
+  0 calls  → 1.0 (best — task didn't need recursion)
+  1 call   → 0.75
+  2 calls  → 0.50
+  3 calls  → 0.25
+  4+ calls → 0.00 (recursion spam is wasteful)
 
-def recursion_efficiency(llm_call_count: int, max_calls: int = 2, **_) -> float:
-    """Staircase bonus for shallow recursion trees.
+This score is *coupled to correctness* in compose.py — wrong answers don't
+earn an efficiency bonus, which prevents the model from learning to spam
+`llm(` strings in code blocks for free reward.
+"""
 
-    Returns 0.0 if llm_call_count > max_calls (depth-explosion penalty).
-    """
+def recursion_efficiency(llm_call_count: int, **_) -> float:
+    """Linear-decay efficiency on call count; gated to correctness in compose."""
     count = max(0, int(llm_call_count))
-    if count > max_calls:
-        return 0.0
-    return _STAIRCASE.get(count, 0.0)
+    return max(0.0, 1.0 - 0.25 * count)
 
 
 __all__ = ["recursion_efficiency"]

rewards/recursion_extract.py

@@ -0,0 +1,66 @@
+"""Recursion call extractor — REW-04 v2.
+
+Counts llm( function calls inside fenced ```python code blocks of a model
+completion. Ignores occurrences in:
+  - prose outside any code block
+  - comments inside a code block (# llm(...) → 0)
+  - string literals inside a code block ("did llm(...)" → 0)
+
+Uses Python's tokenize module for accuracy; regex fallback when the code
+block is syntactically invalid (the model writes broken Python sometimes
+but we still want to count its intent).
+"""
+from __future__ import annotations
+
+import io
+import re
+import tokenize
+
+# Match ``` or ```python or ```py — case-insensitive, multi-line.
+_CODE_BLOCK_RE = re.compile(
+    r"```(?:python|py)?\s*\n(.*?)```",
+    re.DOTALL | re.IGNORECASE,
+)
+_LLM_CALL_RE = re.compile(r"\bllm\s*\(")
+
+
+def _count_in_block(code: str) -> int:
+    """Count llm( calls in one code block. Tokenize-aware; regex fallback."""
+    try:
+        toks = list(tokenize.generate_tokens(io.StringIO(code).readline))
+    except (tokenize.TokenizeError, IndentationError, SyntaxError):
+        # Strip line comments, then regex. Conservative — does not strip
+        # string literals, but the model rarely puts llm( in a string when
+        # writing broken code.
+        stripped = "\n".join(line.split("#", 1)[0] for line in code.splitlines())
+        return len(_LLM_CALL_RE.findall(stripped))
+
+    count = 0
+    for i in range(len(toks) - 1):
+        tok = toks[i]
+        nxt = toks[i + 1]
+        if (
+            tok.type == tokenize.NAME
+            and tok.string == "llm"
+            and nxt.type == tokenize.OP
+            and nxt.string == "("
+        ):
+            count += 1
+    return count
+
+
+def count_llm_calls(completion: str) -> int:
+    """Total llm( calls inside all fenced code blocks of the completion.
+
+    Returns 0 if completion is empty, has no code blocks, or only contains
+    llm( in prose / comments / strings.
+    """
+    if not completion:
+        return 0
+    blocks = _CODE_BLOCK_RE.findall(completion)
+    if not blocks:
+        return 0
+    return sum(_count_in_block(b) for b in blocks)
+
+
+__all__ = ["count_llm_calls"]

scripts/verify_recursion_reward.py

@@ -0,0 +1,66 @@
+"""CPU-only verification of REW-04 v2 reward design.
+
+Confirms two GRPO-blocking properties:
+  1. A synthetic 8-completion group produces non-zero std (v1's std was 0.0
+     across every group, which is why the reward curve was flat).
+  2. The ordering correct+0calls > correct+1call > correct+spam holds.
+
+Run BEFORE spending any HF Jobs credits on a retrain.
+"""
+from __future__ import annotations
+
+import statistics
+import types
+
+from rewards.compose import compose_reward_single
+
+cfg = types.SimpleNamespace(
+    alpha=0.2,
+    weights=types.SimpleNamespace(
+        correctness=0.70, token_budget=0.15, recursion_efficiency=0.15
+    ),
+    token_budget_variant="capped_linear",
+    answer_regex="<answer>(.*?)</answer>",
+    max_calls=4,
+)
+
+GOLD = "silver"
+GENERATIONS = [
+    ("correct + 0 llm calls (REPL grep)",
+     "```python\nimport re\nm=re.search('silver', ctx)\nprint(m.group())\n```\n<answer>silver</answer>"),
+    ("correct + 1 llm call",
+     "```python\nans=llm('color', ctx[:5000])\nprint(ans)\n```\n<answer>silver</answer>"),
+    ("correct + 3 llm calls (wasteful)",
+     "```python\na=llm('q1',ctx[:1000])\nb=llm('q2',ctx[1000:2000])\nc=llm('q3',ctx[2000:3000])\n```\n<answer>silver</answer>"),
+    ("correct + bare answer (no code, trivial-task path)",
+     "<answer>silver</answer>"),
+    ("wrong + format",
+     "<answer>gold</answer>"),
+    ("wrong + no format",
+     "the color is gold"),
+    ("right text + no format (v1 collapse mode)",
+     "silver"),
+    ("format-only spam",
+     "<answer></answer>"),
+]
+
+print(f"{'idx':>3}  {'score':>6}  {'calls':>5}  description")
+print("-" * 78)
+scores = []
+for i, (desc, gen) in enumerate(GENERATIONS):
+    s, m = compose_reward_single(gen, GOLD, 200, cfg)
+    scores.append(s)
+    print(f"{i:>3}  {s:>6.3f}  {int(m['llm_call_count']):>5d}  {desc}")
+print("-" * 78)
+print(f"group mean: {statistics.mean(scores):.4f}")
+print(f"group std:  {statistics.stdev(scores):.4f}  (must be > 0.10 for GRPO advantage)")
+print(f"max - min:  {max(scores) - min(scores):.4f}")
+
+# Hard gates — exit non-zero if any fail
+assert statistics.stdev(scores) > 0.10, "FAIL: group std too low; GRPO will not learn"
+assert scores[0] > scores[1] > scores[2], (
+    f"FAIL: efficiency ordering broken (got {scores[0]:.3f} > {scores[1]:.3f} > {scores[2]:.3f})"
+)
+assert scores[0] > scores[4], "FAIL: correct must beat wrong"
+assert scores[7] <= 0.25, "FAIL: format-only spam not capped"
+print("\nPASS: REW-04 v2 produces learnable variance and correct orderings")

tests/test_recursion_extract.py

@@ -0,0 +1,55 @@
+"""Tests for rewards.recursion_extract — REW-04 v2."""
+from rewards.recursion_extract import count_llm_calls
+
+
+def test_empty_completion():
+    assert count_llm_calls("") == 0
+
+
+def test_no_code_block():
+    assert count_llm_calls("The answer is <answer>silver</answer>") == 0
+
+
+def test_single_call():
+    c = "```python\nresult = llm('find', ctx[:5000])\n```\n<answer>silver</answer>"
+    assert count_llm_calls(c) == 1
+
+
+def test_two_calls_in_one_block():
+    c = "```python\na = llm('q1', ctx[:1000])\nb = llm('q2', ctx[1000:])\n```"
+    assert count_llm_calls(c) == 2
+
+
+def test_calls_across_two_blocks():
+    c = "```python\nx=llm('q', ctx)\n```\nthen\n```python\ny=llm('q2', ctx)\n```"
+    assert count_llm_calls(c) == 2
+
+
+def test_call_in_comment_not_counted():
+    c = "```python\n# would call llm(stuff) but skipping\nprint('done')\n```"
+    assert count_llm_calls(c) == 0
+
+
+def test_call_in_string_literal_not_counted():
+    c = '```python\nnote = "earlier code did llm(...)"\nprint(note)\n```'
+    assert count_llm_calls(c) == 0
+
+
+def test_call_outside_code_block_not_counted():
+    c = "Maybe I should call llm(question, chunk) but I won't actually."
+    assert count_llm_calls(c) == 0
+
+
+def test_call_in_loop_counts_literal_occurrence():
+    c = "```python\nfor chunk in chunks:\n    r = llm('find', chunk)\n```"
+    assert count_llm_calls(c) == 1
+
+
+def test_invalid_python_falls_back_to_regex():
+    c = "```python\nthis is not valid python !!!\nresult = llm('q', ctx)\n```"
+    assert count_llm_calls(c) >= 1  # fallback regex finds it
+
+
+def test_bare_python_fence():
+    c = "```\nans = llm('q', ctx)\n```"  # no language tag
+    assert count_llm_calls(c) == 1

tests/test_rewards.py

@@ -135,25 +135,24 @@ class TestRecursionEfficiency:
         assert recursion_efficiency(0) == 1.0
 
     def test_one_call(self):
-        assert recursion_efficiency(1) == 0.7
+        assert recursion_efficiency(1) == 0.75
 
     def test_two_calls(self):
-        assert recursion_efficiency(2) == 0.4
+        assert recursion_efficiency(2) == 0.50
 
-    def test_three_calls_over_max(self):
-        assert recursion_efficiency(3) == 0.0
+    def test_three_calls(self):
+        assert recursion_efficiency(3) == 0.25
 
-    def test_ten_calls_zero(self):
+    def test_four_or_more_calls_zero(self):
+        assert recursion_efficiency(4) == 0.0
         assert recursion_efficiency(10) == 0.0
 
     def test_negative_treated_as_zero(self):
         assert recursion_efficiency(-1) == 1.0
 
-    def test_max_calls_override(self):
-        # max_calls=3: 3 calls should still score
-        assert recursion_efficiency(3, max_calls=3) == 0.0  # >2 gets 0 from staircase, but <=max_calls
-        # With max_calls=4, 3 calls return staircase 0.0 (not in map) → 0.0
-        assert recursion_efficiency(2, max_calls=4) == 0.4  # <=max_calls, staircase applies
+    def test_max_calls_ignored_in_v2(self):
+        # max_calls kwarg is accepted via **_ but ignored by new linear decay
+        assert recursion_efficiency(2, max_calls=1) == 0.50
 
 
 # ---------------------------------------------------------------------------
@@ -162,24 +161,24 @@ class TestRecursionEfficiency:
 
 class TestCompose:
     def test_format_fail_short_circuits(self, cfg_reward):
-        score = compose_reward_single("no tag here", "Rome", 100, 0, cfg_reward)
-        assert score == 0.0
+        score, _ = compose_reward_single("no tag here", "Rome", 100, cfg_reward, llm_call_count=0)
+        assert 0.0 <= score <= 0.25  # A-02 cap applies, but > 0 because of token budget
 
     def test_perfect_score(self, cfg_reward):
         # Short completion, 0 llm calls, correct answer
-        score = compose_reward_single("<answer>Rome</answer>", "Rome", 1, 0, cfg_reward)
-        assert score == pytest.approx(0.75 * 1.0 + 0.2 * 1.0 + 0.05 * 1.0, abs=0.05)
+        score, _ = compose_reward_single("<answer>Rome</answer>", "Rome", 1, cfg_reward, llm_call_count=0)
+        assert score == pytest.approx(0.75 * 1.0 + 0.2 * 1.0 + 0.05 * 1.0 + 0.10, abs=0.05)
 
-    def test_correctness_zero_capped_at_005(self, cfg_reward):
-        # A-02 fix: format passes but wrong answer → capped at 0.05
-        score = compose_reward_single("<answer>Paris</answer>", "Rome", 100, 0, cfg_reward)
-        assert score <= 0.05
+    def test_correctness_zero_capped_at_025(self, cfg_reward):
+        # format passes but wrong answer → capped at 0.25 (since format bonus + cap)
+        score, _ = compose_reward_single("<answer>Paris</answer>", "Rome", 100, cfg_reward, llm_call_count=0)
+        assert score <= 0.25
 
     def test_partial_mix(self, cfg_reward):
         # Correct answer but long completion (token_budget reduced)
         long_text = " ".join(["filler"] * 500) + " <answer>42</answer>"
-        score = compose_reward_single(long_text, "42", 10, 0, cfg_reward)
-        assert 0.5 < score <= 1.0  # correctness=1 dominates
+        score, _ = compose_reward_single(long_text, "42", 10, cfg_reward, llm_call_count=0)
+        assert 0.5 < score <= 1.10  # correctness=1 dominates
 
     def test_make_reward_fn_factory(self, cfg_reward):
         bound_fn = make_reward_fn(cfg_reward)
@@ -204,7 +203,7 @@ class TestCompose:
         )
         assert len(results) == 2
         assert results[0] > 0.0
-        assert results[1] == 0.0
+        assert results[1] <= 0.25
 
 
 # ---------------------------------------------------------------------------
@@ -213,28 +212,28 @@ class TestCompose:
 
 @pytest.mark.reward_audit
 def test_masked_context_reward_is_baseline(cfg_reward):
-    """REW-06: gibberish answer with format tag must score <= 0.10 (A-02 cap)."""
-    score = compose_reward_single(
+    """REW-06: gibberish answer with format tag must score <= 0.25 (A-02 cap)."""
+    score, _ = compose_reward_single(
         "<answer>lorem ipsum dolor sit amet</answer>",
         "Rome",
         100,
-        0,
         cfg_reward,
+        llm_call_count=0,
     )
-    assert score <= 0.10, f"Masked-context score too high: {score}"
+    assert score <= 0.25, f"Masked-context score too high: {score}"
 
 
 @pytest.mark.reward_audit
 def test_format_only_composite_is_capped(cfg_reward):
-    """REW-06: 20 format-only completions with mismatched gold must all score <= 0.05."""
+    """REW-06: 20 format-only completions with mismatched gold must all score <= 0.25."""
     words = ["banana", "kiwi", "mango", "peach", "grape", "lemon", "plum",
              "cherry", "pear", "melon", "papaya", "guava", "fig", "date",
              "lime", "apricot", "coconut", "blueberry", "raspberry", "strawberry"]
     composites = [
-        compose_reward_single(f"<answer>{w}</answer>", "Rome", 100, 0, cfg_reward)
+        compose_reward_single(f"<answer>{w}</answer>", "Rome", 100, cfg_reward, llm_call_count=0)[0]
         for w in words
     ]
-    assert max(composites) <= 0.05, f"Format-only composites: {composites}"
+    assert max(composites) <= 0.25, f"Format-only composites: {composites}"
 
 
 @pytest.mark.reward_audit
@@ -244,7 +243,7 @@ def test_no_monotonic_length_exploit(cfg_reward):
     composites = []
     for n in lengths:
         completion = " ".join(["filler"] * n) + " <answer>rome</answer>"
-        score = compose_reward_single(completion, "rome", 100, 0, cfg_reward)
+        score, _ = compose_reward_single(completion, "rome", 100, cfg_reward, llm_call_count=0)
         composites.append(score)
 
     # Longer completions must not score strictly higher
@@ -255,3 +254,71 @@ def test_no_monotonic_length_exploit(cfg_reward):
     assert all(a >= b for a, b in zip(composites, composites[1:])), (
         f"Non-monotonic sequence: {[round(x, 3) for x in composites]}"
     )
+
+
+class TestComposeV3:
+    """REW-02 v3: soft format + recursion-extraction + correctness-gated efficiency."""
+
+    @pytest.fixture
+    def cfg_v3(self):
+        return OmegaConf.create({
+            "alpha": 0.2,
+            "weights": {"correctness": 0.70, "token_budget": 0.15, "recursion_efficiency": 0.15},
+            "token_budget_variant": "capped_linear",
+            "answer_regex": "<answer>(.*?)</answer>",
+            "max_calls": 4,
+        })
+
+    def test_correct_no_recursion_scores_high(self, cfg_v3):
+        c = "```python\nimport re\nm=re.search('silver', ctx)\nprint(m.group())\n```\n<answer>silver</answer>"
+        score, metrics = compose_reward_single(c, "silver", 100, cfg_v3)
+        assert score >= 0.85, f"clean correct should score high, got {score}"
+        assert metrics["llm_call_count"] == 0
+
+    def test_zero_calls_beats_one_call_when_both_correct(self, cfg_v3):
+        c0 = "```python\nimport re\nm=re.search('silver', ctx)\nprint(m.group())\n```\n<answer>silver</answer>"
+        c1 = "```python\nans=llm('color', ctx[:5000])\nprint(ans)\n```\n<answer>silver</answer>"
+        s0, _ = compose_reward_single(c0, "silver", 100, cfg_v3)
+        s1, _ = compose_reward_single(c1, "silver", 100, cfg_v3)
+        assert s0 > s1, f"0-call ({s0:.3f}) should beat 1-call ({s1:.3f}) when both correct"
+
+    def test_efficiency_gated_on_correctness(self, cfg_v3):
+        # Wrong answer with 0 calls — must NOT earn efficiency bonus.
+        c = "```python\nprint('done')\n```\n<answer>gold</answer>"
+        score, metrics = compose_reward_single(c, "silver", 100, cfg_v3)
+        assert metrics["recursion_eff_contribution"] == 0.0
+        assert score <= 0.25, f"wrong answer must be capped, got {score}"
+
+    def test_recursion_spam_loses_to_minimal_recursion(self, cfg_v3):
+        c2 = "```python\na=llm('q1',ctx[:1000])\nb=llm('q2',ctx[1000:2000])\n```\n<answer>silver</answer>"
+        c5 = "```python\n" + "\n".join(f"x{i}=llm('q{i}',ctx)" for i in range(5)) + "\n```\n<answer>silver</answer>"
+        s2, _ = compose_reward_single(c2, "silver", 100, cfg_v3)
+        s5, _ = compose_reward_single(c5, "silver", 100, cfg_v3)
+        assert s2 > s5, f"2-call ({s2:.3f}) should beat 5-call spam ({s5:.3f})"
+
+    def test_format_only_capped(self, cfg_v3):
+        c = "<answer>wrong</answer>"
+        score, _ = compose_reward_single(c, "silver", 100, cfg_v3)
+        assert 0.05 <= score <= 0.25, f"format-only wrong should be in [0.05, 0.25], got {score}"
+
+    def test_no_format_gets_minimal_credit(self, cfg_v3):
+        c = "silver"  # right text but no <answer> tag
+        score, _ = compose_reward_single(c, "silver", 100, cfg_v3)
+        assert score <= 0.20
+
+    def test_group_variance_nonzero(self, cfg_v3):
+        """Smoke check: a synthetic GRPO group of 8 must produce non-zero std.
+        v1 had std=0.0 across all groups, which zeroed the GRPO advantage."""
+        gens = [
+            "```python\nimport re\nm=re.search('silver',ctx)\nprint(m.group())\n```\n<answer>silver</answer>",
+            "```python\nans=llm('color',ctx[:5000])\nprint(ans)\n```\n<answer>silver</answer>",
+            "<answer>silver</answer>",
+            "<answer>gold</answer>",
+            "the color is silver",
+            "<answer></answer>",
+            "```python\n" + "\n".join(f"x{i}=llm('q{i}',ctx)" for i in range(5)) + "\n```\n<answer>silver</answer>",
+            "silver.",
+        ]
+        scores = [compose_reward_single(g, "silver", 200, cfg_v3)[0] for g in gens]
+        import statistics as _st
+        assert _st.stdev(scores) > 0.10, f"group std too low: {_st.stdev(scores)}"

train/grpo.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import logging
 import os
 import inspect
+import statistics
 from pathlib import Path
 from typing import Any, Callable
 
@@ -139,15 +140,18 @@ def run_grpo(
         rewards = reward_fn(prompts, completions, **kwargs)
         try:
             if wandb.run is not None:
-                wandb.log({
+                log_dict = {
                     "reward/composite_mean": sum(rewards) / max(len(rewards), 1),
-                    "reward/format_pass_mean": float(kwargs.get("_format_pass_mean", 0.0)),
-                    "reward/correctness_mean": float(kwargs.get("_correctness_mean", 0.0)),
-                    "reward/token_budget_mean": float(kwargs.get("_token_budget_mean", 0.0)),
-                    "reward/recursion_eff_mean": float(kwargs.get("_recursion_eff_mean", 0.0)),
-                })
+                    "reward/composite_std": (
+                        statistics.stdev(rewards) if len(rewards) > 1 else 0.0
+                    ),
+                }
+                cm = kwargs.get("_component_means", {})
+                for k, v in cm.items():
+                    log_dict[f"reward/{k}_mean"] = float(v)
+                wandb.log(log_dict)
         except Exception:
-            pass  # W&B logging is non-blocking
+            pass
         return rewards
 
     # TRN-03 step 3: Build train_dataset from data/train.jsonl.
@@ -184,12 +188,7 @@ def run_grpo(
     # if the completion lacks <answer>...</answer>, composite reward = 0.0.
     # Previous wording ("shortest exact answer span") never told the model
     # about the tag, so 50/50 GRPO steps had reward=0.0 (job 69ece94a).
-    sys_msg = (
-        "You are FATHOM, a recursive language model. You answer questions "
-        "about long documents. Read the context, think step by step, and "
-        "emit your final answer inside <answer>...</answer> tags. "
-        "Keep the answer the shortest exact span that answers the question."
-    )
+    sys_msg = "You are FATHOM, a recursive language model with a Python REPL sandbox. You can read a long document via the variable `ctx` and call `llm(prompt, chunk)` for sub-queries. Think step by step. Emit your final answer inside <answer>...</answer>."
 
     # vLLM hard-checks final prompt length against the model's max_position_embeddings
     # (32 768 for Qwen 2.5 Coder 0.5B/1.5B). Our train.jsonl `context` field is the
@@ -218,15 +217,14 @@ def run_grpo(
     def _to_prompt(example: dict) -> dict:
         ctx_full = example.get("context", "") or ""
         ctx_truncated = _truncate_to_tokens(ctx_full, ctx_budget_tok)
+        # CRITICAL: must match data/sft_traces.jsonl user-message shape exactly.
+        user_content = (
+            f"{example.get('prompt', '')}\n\n"
+            f"[Document excerpt]:\n{ctx_truncated}"
+        )
         msgs = [
             {"role": "system", "content": sys_msg},
-            {
-                "role": "user",
-                "content": (
-                    f"Context:\n{ctx_truncated}\n\n"
-                    f"{example.get('prompt','')}"
-                ),
-            },
+            {"role": "user", "content": user_content},
         ]
         prompt_str = tokenizer.apply_chat_template(
             msgs, tokenize=False, add_generation_prompt=True
@@ -235,7 +233,6 @@ def run_grpo(
             "prompt": prompt_str,
             "gold_answer": str(example.get("gold_answer", "")),
             "prompt_token_count": int(example.get("context_length", 0)) // 4,
-            "llm_call_count": 0,
         }
 
     train_dataset = raw_ds.map(
@@ -259,6 +256,16 @@ def run_grpo(
     log.info("TRN-03 GRPOTrainer constructed (no env tools — pure prompt→completion→reward)")
 
     # TRN-03 step 4: Train
+    # Pre-flight: tokenize one example and confirm the chat-template prefix is
+    # the byte-identical match of an SFT trace prefix.
+    import json as _json
+    _sft = _json.loads(open(str(cfg.data.sft_traces_path) if hasattr(cfg.data, "sft_traces_path") else "data/sft_traces.jsonl", encoding="utf-8").readline())
+    _sft_prefix = tokenizer.apply_chat_template(_sft["messages"][:2], tokenize=False, add_generation_prompt=True)[:200]
+    _grpo_first = train_dataset[0]["prompt"][:200]
+    assert _sft_prefix.split("Question:")[0] == _grpo_first.split("Question:")[0], (
+        "SFT/GRPO chat-template prefix drift detected — see ANTIGRAVITY_BRIEF.md §A.3.1"
+    )
+
     trainer.train()
 
     # TRN-03 step 5: STACK §6 save sequence — adapter-only FIRST