Skip to content

OVERVIEW.md

Latest commit

 

History

History
421 lines (339 loc) · 23.4 KB

OVERVIEW.md

File metadata and controls

421 lines (339 loc) · 23.4 KB
Raw

Fireworks RFT on τ²-bench — Full Training Instructions

How to run reinforcement fine-tuning (RFT) of a small open model on our τ²-bench banking_knowledge domain, using Fireworks RFT driven by the Eval Protocol (EP).

This is the operational doc — the exact reward, MCP server, dataset, and CLI commands. For the why (research grounding, reward-densification rationale, the <10%-pass problem), read these first; this doc assumes them:

⚠️ The EP/Fireworks SDK moves fast. Every flag name, decorator argument, and import path below is verified against the docs as of 2026-06, but you MUST re-check against the installed version (pip show eval-protocol, eval-protocol create rft --help) before launching a paid job. Where the SDK has likely drifted, it's flagged inline.

0. TL;DR — the whole loop in one picture

                      Eval Protocol (EP) orchestrates this loop
   ┌──────────────────────────────────────────────────────────────────────┐
   │                                                                        │
   │   policy model  ──►  ROLLOUT          ──►  EVALUATOR        ──► reward │
   │  (Fireworks      (MCP-Gym steps our    (@evaluation_test:    (scalar   │
   │   base model)     banking env: agent    our partial-credit   0..1)     │
   │                   ↔ tools ↔ sim user)   reward fn)              │       │
   │        ▲                                                       │       │
   │        └───────────────  TRAINER updates weights (GRPO)  ◄─────┘       │
   └──────────────────────────────────────────────────────────────────────┘

EP gives one scalar reward in [0,1] per full rollout. There is no native per-turn / per-token reward channel. Everything in our reward design exists to make that single scalar dense in value so a small model that passes <10% still produces a usable gradient. See §4.

1. Prerequisites

pip install eval-protocol            # the EP Python SDK + CLI
eval-protocol --version              # confirm it installed
export FIREWORKS_API_KEY="fw_..."    # from app.fireworks.ai → API keys

You also need:

  • A Fireworks account with RFT access. Models ≤ ~16B are the free/eligible tier to tune — confirm eligibility for our Qwen variants before launching.
  • The base-model path for whatever you're tuning, e.g. accounts/fireworks/models/qwen3-8b (verify the exact slug in the Model Library at app.fireworks.ai/models — see workspace/training_plan.md §1 for our two Qwen deployments).

Repo assets we reuse (do not rebuild these):

  • Domain/env: src/tau2/domains/banking_knowledge/ (environment, tools.py, retrieval).
  • Data: data/tau2/domains/banking_knowledge/ (db.json, ~700 knowledge docs, ~100 tasks).
  • Reward: reward_functions/recompute_partial_credit.py (gated partial-credit — §4).
  • Difficulty split: sort-tasks/task_difficulty.csv (Easy/Medium/Hard/super_hard buckets).

2. The Eval Protocol model (what you actually write)

Two pieces of code, plus a dataset:

  1. An MCP server that exposes our banking_knowledge environment as a tool-calling env the rollout processor can step (McpGym subclass). See §3.
  2. An evaluator — a pytest-style function decorated with @evaluation_test that receives the finished trajectory and returns a scalar reward. See §4.
  3. A dataset — one JSONL row per task (id + scenario seed + evaluation criteria). See §5.

Canonical evaluator shape (verified against the τ²-bench tutorial)

from eval_protocol.models import EvaluationRow, EvaluateResult
from eval_protocol.pytest import evaluation_test
from eval_protocol.mcp import MCPGymRolloutProcessor   # verify import path in your version

@evaluation_test(
    input_dataset=["FIREWORKS_TRAINING/data/banking_easy.jsonl"],
    dataset_adapter=banking_to_evaluation_row,          # our adapter, §5
    completion_params=[{
        "model": "fireworks_ai/accounts/fireworks/models/qwen3-8b",
        "temperature": 0.7,
        "max_tokens": 4096,
    }],
    rollout_processor=MCPGymRolloutProcessor(),
    server_script_path="FIREWORKS_TRAINING/banking_mcp/server.py",  # our MCP server, §3
    mode="pointwise",          # one independent score per row
    passed_threshold=0.4,      # τ²-bench tutorial uses 0.40; tune for banking
    num_runs=2,                # repeat rollouts to average out user-sim noise
    max_concurrent_rollouts=32,
)
def test_banking(row: EvaluationRow) -> EvaluationRow:
    # `row` holds the FULL multi-turn trajectory after the rollout has run.
    score, reason = banking_partial_credit_reward(row)   # 0..1, see §4
    row.evaluation_result = EvaluateResult(score=score, reason=reason)
    return row

Key facts about this API:

  • mode="pointwise" → each row scored independently (what we want). Other modes (groupwise) exist for pairwise/ranking rewards — not for us.
  • num_runs repeats the whole rollout N times and the harness aggregates — use it to damp simulated-user variance (the LLM user is a reward-noise source; see playbook 012 §7).
  • The same decorated function is both the offline eval (ep local-test) and the live RFT reward — write it once.

3. The MCP server (wrapping our banking env)

The τ²-bench tutorial ships an airline server at examples/tau2_mcp/server.py. We mirror it but mount our banking domain. Pattern from the tutorial:

from typing import Optional
from eval_protocol.mcp import McpGym, EnvironmentAdapter   # verify import path
from tau2.domains.banking_knowledge.environment import BankingKnowledgeEnvironment  # our env

class BankingDomainMcp(McpGym):
    def __init__(self, seed: Optional[int] = None):
        default_config = {"domain": "banking_knowledge", "max_turns": 20}  # cap turns (see §6)
        self.adapter = EnvironmentAdapter(
            env_class=BankingKnowledgeEnvironment,
            default_config=default_config,
        )
        super().__init__("banking_knowledge", self.adapter, seed)

if __name__ == "__main__":
    BankingDomainMcp().run()

This is Path B from playbook 012 (copy the tutorial, swap domain). The lighter Path A (wrap our env directly via a remote rollout / RemoteRolloutProcessor over HTTP) is also valid and reuses more of our existing tau2 runner — but Path B stays closest to the shipped, working test_tau2_bench example, so start there.

The EnvironmentAdapter must drive the agent ↔ tools ↔ simulated user loop and expose the final reward_info (db_check, action_checks, reward_breakdown) that our reward reads. Our BankingKnowledgeEnvironment already computes that per simulation — the adapter's job is just to surface it on the finished trajectory. Confirm the exact env entrypoint/class name in src/tau2/domains/banking_knowledge/environment.py before wiring.

4. The reward function — the most important part

Do not use the binary τ² DB check as the RFT reward. A small model passing <10% means most GRPO groups are all-zero → advantage 0 → no gradient → training stalls. We densify.

What the τ²-bench tutorial does (the default, multiplicative-binary)

reward = 1.0
reward *= env_reward_info.reward          # DB-state integrity (hash of final state)
reward *= action_reward_info.reward       # required tool actions called w/ right args
reward *= nl_reward_info.reward           # LLM-as-judge NL assertions
reward *= communicate_reward_info.reward  # agent communicated required content

Each factor is binary, so this is effectively all-or-nothing — fine for evaluation, bad for training a weak model. Keep it as the eval/benchmark reward (it matches the headline metric), but train against the densified reward below.

What we train against: our gated partial-credit reward

Use reward_functions/recompute_partial_credit.py. It's already banking-specific and gated/monotonic so a failure can never out-score a real solve:

full solve (τ² reward == 1.0)  →  exactly 1.0
otherwise                      →  CAP(0.90) * partial,  partial ∈ [0,1]

partial composite (weights live at the top of the file):
    W_SUBGOAL * subgoal             # AgenticQwen: Σ action_match / N   (headline term)
  + W_PARAM   * param_partial       # ToolRL-lite: tool name + arg-name + arg-value match
  + W_DB      * db_match            # full-correctness bonus
  + W_VERIFY  * verified_before_act # banking guard: identity verification precedes a state change
  - P_REP     * repetition_penalty  # penalize degenerate loops (we saw GLM do 300+ msgs)

Why each piece matters (full rationale in playbook 012 §§1,8):

  • subgoal is the densifier: a rollout that did 3 of 5 right actions scores ~0.9*3/5 instead of 0. Groups stop being all-zero. This single term is the main lever.
  • verified_before_act + param_partial are the anti-reward-hacking guards — in banking, an unverified transfer that "looks right" must score low, not partially. Partial credit invites "right tool, wrong args"; these catch it.
  • Gating (CAP=0.90) keeps it monotonic: partial credit never reaches a true solve, so the model can't farm partial credit instead of solving.
  • Prefer the deterministic terms over the LLM-judged NL assertions — judge noise is indistinguishable from signal for a small model.

Wiring it into EP

recompute_partial_credit.py today reads finished results.json files offline. For RFT you need the same scoring math to read EP's live EvaluationRow trajectory instead. Plan:

  1. Factor the scoring core (the subgoal / param_partial / db_match / verified_before_act / repetition computation) out of the file-walking __main__ into a pure function: score_from_reward_info(reward_info, tool_calls) -> (float, reason).
  2. In test_banking (§2), pull reward_info + the agent's tool calls off row, call that function, return EvaluateResult(score=..., reason=...).
  3. Keep the offline __main__ path working — it's how you validate the reward on existing baselines without spending rollout budget.

Masking note: our tool_calling_evaluator.py has the right instinct (FIX 3 — mask server 5xx/429/timeout episodes so the model can't farm reward by triggering env errors). EP scores every returned row, so emulate masking by detecting env errors in the evaluator and returning a neutral/dropped score per EP's convention — check whether your EP version supports skipping a row vs. forcing a fixed score; do not silently score env errors as 0.

4b. Two RL problems specific to this task (read before launching)

These are the two things that break banking RFT in practice. Both come back to the same fact: EP/Fireworks managed RFT is trajectory-scalar — one reward per full rollout, and the trainer updates weights per batch of rollouts, never mid-rollout. So "give feedback earlier" and "update more often" are achieved by making rollouts shorter and denser, not by injecting a gradient step in the middle of a conversation.

Problem 1 — 300-message degenerate loops (mostly repeated turns)

Can I do multiple weight updates inside one rollout? In managed Fireworks RFT: no. A single agent conversation produces one trajectory → one scalar → it joins a batch → one gradient step. Literal mid-episode updates need a self-hosted token-level multi-turn trainer (verl / Megatron / rLLM — that's the stack in [[002_Multi-Turn-RL-Iterative-Reward-Calibration]], not Fireworks). Don't go there yet.

What you can do — and what actually fixes the loops — are three deterministic guards that make the rollout end early and score low, so the policy learns not to loop:

  1. Hard turn cap in the MCP server (max_turns, §3). Deterministically ends the episode at e.g. 20 turns instead of 300. Seed the cap per task from avg_messages in the difficulty CSV (the build script writes a max_turns_hint per row).
  2. Early-stop on repetition — should_early_stop() in workspace/score_from_reward_info.py. The rollout processor calls it every turn; the moment the agent repeats the same (tool, args) call REP_FULL times, terminate the rollout now and score the short trajectory. This is the "feedback much earlier" you asked for — the episode dies at message ~12, not 300.
  3. Repetition penalty in the reward — already in the partial-credit reward (P_REP, REP_FREE=5, REP_FULL=25). A looping rollout is penalized, so even short loops score low.

Net effect: loops are cut short (guards 1–2, saving rollout budget) and penalized (guard 3, creating the gradient that teaches "don't loop"). That's the realistic "update more often" — more, shorter (rollout, reward) pairs per unit compute.

Problem 2 — reward only at the end; wanting "smaller steps"

EP gives one scalar at episode end. The reward-timing spectrum (full table in workspace/training_plan.md §3) offers three ways to make that scalar feel step-wise, in increasing engineering cost:

  • B — graded final reward (do this now): the partial-credit reward is dense in value even though it fires once. A rollout that got the tool name + most args right scores ~0.6, not 0 — the gradient reflects partial progress without any new plumbing. This is the main lever and it's already built (§4).
  • C — truncated short rollouts (do this now): low max_turns on Easy tasks → the episode is short, so "one reward at the end" is one reward after a few steps. Pairs with B.
  • D — sub-episode decomposition (later, only if needed): split a task into one mini-episode per expected action, each with its own reward. This is the literal "train step by step." But it needs the env to checkpoint/reset mid-task, which tau2 may not support cleanly — gate it on Phase-1 results.

Banking-specific reality check (from our data): most banking tasks have one expected action (task_001 has exactly 1; the build script counts and warns). For a single-action task there is nothing to decompose — method D buys you nothing. The long rollout there isn't a long chain; it's the agent failing to converge on one correct call. So the lever is intra-action density — param_partial already grades tool-name + arg-name + arg-value separately, turning a binary "did the one call match" into a smooth 0→1 — plus the loop guards above. Reserve sub-episode decomposition for the genuinely multi-action tasks (disputes, multi-step transfers) where there's a real chain to split.

Bottom line for your two asks: you get "earlier feedback" from the early-stop loop guard + turn cap (Problem 1), and "smaller/denser steps" from graded partial credit + truncated rollouts (Problem 2) — both available now in EP. True per-turn weight updates are a self-hosted-trainer feature; don't reach for it until B+C+loop-guards demonstrably plateau.

5. The dataset

One JSONL row per task. The τ²-bench tutorial row carries:

  • id — unique scenario id
  • user_prompt_template — simulated-user context/persona
  • environment_context — domain spec (here: banking_knowledge)
  • user_simulation — simulated-user behavior definition
  • evaluation_criteria — required actions / communications / NL assertions

Build FIREWORKS_TRAINING/data/banking_easy.jsonl by filtering sort-tasks/task_difficulty.csv to bucket == Easy and emitting one row per task from data/tau2/domains/banking_knowledge/tasks/. Per the global rule, write the build script into workspace/ and run it yourself — don't have me execute it.

The dataset_adapter (banking_to_evaluation_row) converts each JSONL row into an EvaluationRow: load the banking system prompt, attach env context + user-sim + criteria as input metadata, set the row id. Mirror tau_bench_airline_to_evaluation_row from the SDK's tests/pytest/test_tau_bench_airline.py.

Phase 1 = Easy only. Don't add Medium/Hard until Easy demonstrably learns (§7).

6. Hyperparameters that matter for the <10%-pass problem

Flag Start value Why
--n (rollouts/prompt, aka group size) 16 With pass p≈0.1, P(≥1 success in group) = 1−(1−p)^n: n=8→57%, n=16→81%. A group needs at least one success to create advantage. Costs more rollouts but makes the gradient exist.
--rl-loss-method grpo (or dapo) GRPO is the default group-relative method. DAPO adds dynamic filtering (drops all-zero and all-one groups) — ideal here; use it if available.
--temperature 0.7–0.8 Exploration: too low and every rollout in a group is identical → no variance → no gradient.
--max-tokens tight (e.g. 4096) Small models ramble into degenerate loops. Cap hard; keep the repetition penalty in the reward. mcpmark uses 6 steps / 192 tokens for reference.
max_turns (MCP server) low (≈20, lower for Easy) Easy tasks are short; capping turns directly cuts the long-rollout cost and tightens credit assignment. Seed from avg_messages/max_messages in the difficulty CSV.
--epochs 1–2 Start at 1 for the smoke run.
--learning-rate 1e-4 (default) Leave at default until the loop is proven.
--lora-rank 8–16 LoRA is the cheap default.

Cold start first. RFT is the follow-on after SFT plateaus, not the cold-start tool. If pass rate is genuinely <10%, SFT the base on rejection-sampled (reward==1.0) trajectories from our existing strong baselines (gpt-oss-120b / gpt-4.1 banking runs in data/simulations/.../*__banking_knowledge/results.json) to lift pass rate into ~25–40% first, then RFT. See playbook 012 §2.

7. Step-by-step run procedure

Each step is a gate — if reward is flat or rollouts blow up in length/cost, stop and diagnose before spending more compute.

  1. Build the Easy dataset → FIREWORKS_TRAINING/data/banking_easy.jsonl (script in workspace/, you run it).

  2. Write the MCP server (§3) → FIREWORKS_TRAINING/banking_mcp/server.py. Smoke-test it serves the banking tools and runs one sim end-to-end.

  3. Refactor the reward (§4) into score_from_reward_info(...) and write the evaluator test_banking (§2).

  4. Validate offline: ep local-test (a.k.a. eval-protocol local-test) on a handful of Easy tasks. Confirm scores look sane and that high-reward trajectories actually solve the task (reward-hacking check — want ≥~80% agreement before trusting the signal).

  5. Baseline eval (no training): run both Qwen models through the evaluator on Easy; record pass rate + mean graded score. This is the bar to beat. Cross-check against our existing leaderboard numbers for sanity.

  6. Cold-start SFT (if pass <10%): rejection-sample reward==1.0 episodes from the strong baselines, SFT the small base, re-baseline.

  7. Smoke RFT run — smallest base, Easy only, --n 16, low max_turns, --epochs 1, small batch. Goal: confirm the loop runs and mean reward moves, not SOTA.

    eval-protocol create rft \
  --base-model accounts/fireworks/models/qwen3-8b \
  --output-model banking-easy-rft-v1 \
  --n 16 \
  --rl-loss-method dapo \
  --temperature 0.8 \
  --max-tokens 4096 \
  --epochs 1 \
  --lora-rank 8 \
  --wandb-project tau2-banking-rft
# the evaluator + server + dataset are picked up from the EP test in the working dir;
# confirm exact wiring with: eval-protocol create rft --help

  8. Analyze exploration — does mean reward rise across steps? Is the model finding the right tool-call pathway, or collapsing/looping? Inspect trajectories in result-visualizer/ (difficulty tags + reward-checklist tabs already help).

  9. Scale carefully, one knob at a time — only after Easy learns: raise max_turns, add the Medium bucket, then try the larger Qwen. Widen the bucket only once the previous one is learning (curriculum, playbook 012 §3).

  10. Hold out a real test split — report on unseen banking tasks (split_*), not the training tasks. That's the whole point of τ2-vs-τ1.

Deploy the tuned model

# one-click deploy is built into Fireworks; via CLI:
firectl create deployment accounts/fireworks/models/banking-easy-rft-v1
# then point the evaluator's completion_params.model at the new deployment and re-run the
# Step-5 eval to measure lift on the held-out split.

8. Known risks / open items (resolve as you go)

  • Reward refactor is the critical path. recompute_partial_credit.py reads files, not live EvaluationRows. Factoring out a pure scorer (§4) is the one piece of real new code.
  • Env-error masking in EP — confirm how your EP version lets you drop a row vs. force a score; don't let server 5xx/429 episodes score 0 (teaches the model to break the env).
  • Adapter must expose reward_info off the finished trajectory — verify the EnvironmentAdapter surfaces db_check + action_checks, or the dense terms have nothing to read.
  • Simulated-user stability — pin a fixed, low-temperature, capable user simulator (we have gpt-oss-120b runs) so the same agent behavior earns the same reward run-to-run. Use num_runs ≥ 2 to average residual noise.
  • Base-model eligibility — confirm the exact Fireworks slugs for our Qwen variants and that they're in the tunable size tier. See the model.request scope / API-key item in TODO.md — RFT rollouts hit the same request path; fix it first.
  • SDK drift — re-verify every flag/import against eval-protocol create rft --help and pip show eval-protocol before the first paid job.

9. Sources