poimen-memory/tasks/M5.5-verl-training-loop.md

M5.5 — verl training loop

Field	Value
Phase	M5 — Post-training
Size	L — 3+ days
Status	⬜ Not started
Flags	—
Spec	inlined below
Blocks	M5.3, M5.4

Goal

Train the LoRA that makes the update and exit gates better than prompting alone.

Facts (inlined — no spec read needed)

Python, using verl (github.com/volcengine/verl), which is what the paper used. Lives outside the Rust workspace; the JSONL corpus is the only interface.

Paper hyperparameters, Table 3 — start here rather than guessing:

chunk size            5000
max prompt length     8192
max response length   2048
clip ratio            0.20
learning rate         1e-6
temperature (train)   1.0      top_p 1.0
temperature (val)     1.0      top_p 0.7
train batch size      128
rollout N             16
mini batch size       128
LR warmup steps       20

Advantage, paper §3.2.2 — the part most likely to be implemented wrong:

Â_traj = r_traj_g − mean over the GROUP of trajectories
Â_turn = r_update_{g,t} − mean over turns AT STEP t across groups
      = α·Â_traj + (1−α)·Â_turn        α = 0.9

Two distinct baselines. Â_turn is normalised across groups at the same t, and the group size at step t can differ from the trajectory group size, because trajectories that exited early have fewer turns.

α=0.9 is the paper's default and the ablation (Figure 8) explains why: at α=1 there is no update-gate reward and accuracy on evidence-free chunks collapses — the model updates indiscriminately, which is exactly the failure this whole system exists to avoid.

Expect instability. The paper's own limitations section says the extra rewards "reduce training stability, requiring a smaller off-policy degree and longer convergence time."

Steps

1. training/ directory, Python, uv-managed. System Python is 3.9.6; this needs 3.11+.
2. Corpus loader for M5.3's format.
3. Configure verl for LoRA on Qwen2.5-3B-Instruct, rank 16–32.
4. Implement the three rewards and the two-baseline advantage, α configurable.
5. Log per step: update accuracy split by evidence-present and evidence-free, exact-exit ratio, format correctness, mean response length, validation reward.
6. Hold out a validation split by project, not by trajectory — same-project trajectories share vocabulary and leak.
7. Export the adapter, version it, publish where M5.4 can mount it.

Acceptance

• Training runs to convergence on the validation reward.
• Both advantage terms are computed with their own baselines.
• Update accuracy on evidence-free chunks does not collapse.
• The adapter loads in M5.4's server.

Verify

Harness: pytest for reward and advantage maths; a short training run for the loop itself.

Integration test — training/tests/test_rewards.py:

1. a1_r_update_signs — matching label +1, mismatching −1.
2. a2_r_exit_bands — exact 0, early −0.75, late −0.5.
3. a3_r_format_strict — any unparsed turn zeroes the whole trajectory.
4. a4_traj_baseline — Â_traj uses the group mean; hand-computed fixture.
5. a5_turn_baseline_at_step_t — Â_turn normalises across groups at the same t; a fixture with unequal trajectory lengths must not misalign. This is the assertion that catches the most likely implementation error.
6. a6_alpha_mix — α=1 yields pure trajectory advantage; α=0 pure turn.
7. a7_alpha_1_degenerates — train 50 steps at α=1 on a fixture; assert evidence-free accuracy drops relative to α=0.9, reproducing the paper's Figure 8b.
8. a8_validation_split_by_project — assert no project appears in both splits.
9. a9_adapter_loads — export, mount in M5.4, assert a completion returns 200.

Command: uv run pytest training/tests -v

False pass:

• Normalising Â_turn over the whole batch rather than per step t. It trains, loss goes down, and the turn-level signal is diluted into noise — assertion 5 is the only thing that catches it.
• Splitting validation by trajectory. Same-project trajectories share phrasing and file paths, so validation reward looks excellent and generalisation is untested.
• Skipping assertion 7 as "too slow". It is the only end-to-end evidence that the update reward is wired to anything.

Traps

• Tuning α before the rewards are verified. Every α is wrong if r_update's sign is flipped, and the symptom looks identical.
• Training on labels whose κ was never measured. The policy learns the labeler, and there is no held-out signal that would reveal it — M5.2 exists for this and M5.3 refuses to export without it.

---

Background: DESIGN.md — P6 · paper §3.2.2, Table 3, Fig 8