B28 composite reward framework

status: scaffolding-shipped-2026-06-13 (training-loop integration pending)
owner: unassigned
created: 2026-06-13
parent_plan: docs/PLAN.md §3 Tier B (B28)
parent_learn: docs/learn/castform-rl-finetune.md (Steal #1)
related_prds: 5.1-reasoning-on-22M.md (GRPO consumer), B22-trajectory-recorder.md (the rollouts each Reward scores)

PRD — Composite reward with named dimensions

Goal

Ship a typed CompositeReward abstraction: a reward is a bag of named scalar dimensions + weights → total. Every dimension is independently logged, plotted, and inspectable. DPO / ES / GRPO (5.1) all consume the same struct.

Stolen from Castform’s reward composition pattern (see docs/learn/castform-rl-finetune.md §1). Different from existing ES/DPO single-scalar paths.

Why now

• A1 specialist’s eventual DPO-refinement pass benefits immediately. A specialist’s reward is rarely “one thing” — it’s “tool-call is right” + “spoken_text is concise” + “didn’t hallucinate”. Hand- rolling that combination per recipe is the friction this removes.
• 5.1 GRPO PRD calls out RLVR for math + code with verifiable outcome rewards. Once that lands, “verifiable outcome” is one dimension; “did the rollout use the cheap tool first?” is another. Same composite abstraction supports both.

Scope — in

• Sources/TinyGPTModel/CompositeReward.swift — types + aggregation. Already shipped in this PR:
  • RewardDimension (name, score, weight)
  • CompositeReward (dimensions, total computed property, logging)
  • JSONL round-trip for per-rollout dashboards
• Training-loop integration (the remaining work):
  • posttrainllm dpo --reward-fn <path> — load a user-defined reward fn from a Swift-shaped specification file (JSON with weights; individual scorers are still Swift code per recipe). The shipped .swift recipe files do the actual scoring; the --reward-fn arg just routes the aggregation.
  • ES same flag.
  • Sources/TinyGPT/Grpo.swift (5.1) consumes natively.
• Train-time logging: per-step JSONL row gains composite_reward: {dim_a: …, dim_b: …, total: …}.

Scope — out

• Pluggable Python reward fns — Castform’s BaseEnv.compute_reward is in Python. We keep scoring in Swift for type safety + perf; Python harness scoring (existing eval_pace_unhappy.py style) emits its own JSONL that the Swift training loop reads as input.
• Auto-weighting (learn the weights from data). Manual weights for V1.
• Reward shaping over time (decay one dimension’s weight as training progresses). V2.

Files (already shipped in this PR)

Files still to touch (remaining B28 work)

Acceptance criteria

Scaffolding (shipped this PR)

• CompositeReward aggregates weighted dimensions to a total matching a hand-checked formula.
• JSON round-trip preserves dimension names + weights + scores.
• Unit tests pass under xcrun swift test --filter CompositeRewardTests.

Full B28 ship (remaining)

• posttrainllm dpo --reward-fn weights.json ... runs end-to-end on a fixture dataset; per-step history JSONL carries the composite reward block.
• Train viewer (C10) renders per-dimension reward curves.
• One published recipe in docs/recipes/composite-reward.md.

Reference patterns

• Sources/TinyGPT/Es.swift lines 187–227 — the existing centred-reward pattern; CompositeReward stacks above it.
• Sources/TinyGPT/Dpo.swift SimPO/IPO loss paths — the margin-reward shape composite replaces with named dims.
• Castform site — the public reward- composition copy.

Open questions

• Whether to require weights to sum to 1.0 (normalize) vs allow arbitrary weights. Recommendation: don’t normalize — the user knows what they mean; the total is what’s plotted.