Reward setup
The reward turns what the policy did into the number the trainer optimizes. Most rewards combine some of these:
- A verifier checks a result against ground truth: an exact answer, a passing test suite, or a state the environment can confirm. It is deterministic and cheap.
- A rule-based score rewards properties of the output, such as the requested format or a length limit.
- A judge is a model that grades what a verifier can’t check, such as helpfulness or efficiency. It is slower, costs money per call, and can fail.
In Flyte, scoring runs in its own CPU environment. It never holds a GPU, and its retries and timeouts are separate from generation and training.
Write a verifiable reward
A verifiable reward is a plain Python function. This one gives 1.0 for the correct answer and 0.2 for using the required answer format:
def _extract_answer(text: str) -> str | None:
"""Pull the integer following the last '####' marker; fall back to the last integer in the text."""
import re
if "####" in text:
tail = text.rsplit("####", 1)[1]
m = re.search(r"-?\d+", tail)
if m:
return m.group(0)
nums = re.findall(r"-?\d+", text)
return nums[-1] if nums else None
def _reward(rollout: Rollout) -> float:
"""Verifiable reward: 1.0 for the correct answer, +0.2 format bonus for emitting the '####' marker."""
reward = 0.0
if "####" in rollout.completion:
reward += 0.2
predicted = _extract_answer(rollout.completion)
if predicted is not None and predicted == rollout.answer:
reward += 1.0
return reward
Keep the reward function outside of any task, so you can unit-test it against hand-labeled completions without running the loop.
Score a group per task
A rule-based reward takes microseconds, so a task per completion would spend nearly all of its time starting containers. Score a whole group of completions in one task:
reward_env = flyte.TaskEnvironment(
name="rl-grpo-reward",
image=image,
resources=flyte.Resources(cpu=1, memory="2Gi"),
)
@reward_env.task
async def score_group(rollouts: list[Rollout]) -> list[float]:
"""Score a whole prompt group in one task — one reward task per group, not per rollout.
The rule-based reward is microseconds of pure-Python work, so a task *per rollout* would pay pod
startup over and over for trivial compute. Scoring at the group granularity (the unit `generate`
already returns) keeps reward an observable, pipelined task while cutting the pod count ~GROUP_SIZE×.
"""
return [_reward(r) for r in rollouts]
The group is the unit the rollout task already returns, and the unit GRPO normalizes rewards over. Scoring at that granularity keeps reward observable as its own step in the UI, at a fraction of the task count. The driver scores each group as soon as its rollouts arrive. See Training loop.
Call a judge without paying for it twice
A judge call is an API request that costs money and can fail after it returns, for example while uploading the result. Wrap the call in @flyte.trace, so a retried task replays the recorded result instead of calling the judge again:
@flyte.trace
async def call_judge(traj: Trajectory, rubric: Rubric) -> float:
"""Stand-in for the vendor LLM judge. Traced: replayed, not re-called, on a task retry."""
print(f"JUDGE CALLED trial={traj.trial_id} rubric={rubric.name}", flush=True) # count these in logs
await asyncio.sleep(0.5) # network
rng = random.Random(_seed("judge", traj.trial_id, rubric.name))
correct = 1.0 if traj.answer == _world_secret_digest(traj.world_id) else 0.0
efficiency = max(0.0, 1.0 - (len(traj.turns) - 1) / N_DOCS)
tidy = 1.0 if traj.files_written == ["answer.txt"] else 0.0
score = (
rubric.correctness_weight * correct + rubric.efficiency_weight * efficiency + rubric.tidiness_weight * tidy
) + rng.uniform(-0.05, 0.05)
return round(score, 3)
@judge_env.task(retries=2)
async def verify_and_judge(traj: Trajectory, rubric: Rubric, judge_flake_rate: float = 0.0) -> Reward:
verified = traj.answer == _world_secret_digest(traj.world_id) # the deterministic verifier
score = await call_judge(traj, rubric) # [J] memoized across retries
attempt = int(os.environ.get("FLYTE_ATTEMPT_NUMBER", "0"))
if attempt == 0 and random.Random(_seed("flake", traj.trial_id)).random() < judge_flake_rate:
raise RuntimeError("simulated post-judge failure (e.g. upload) — retry must NOT re-call the judge")
found_doc = next((t["doc"] for t in traj.turns if t["found"]), None)
return Reward(
trial_id=traj.trial_id,
world_id=traj.world_id,
verified_correct=verified,
n_turns=len(traj.turns),
judge_score=score,
reward=round((1.0 if verified else 0.0) + score, 3),
found_doc=found_doc,
)
verify_and_judge combines both kinds of signal: the deterministic verifier decides whether the answer is correct, and the judge grades how the policy got there. The task returns a typed Reward, so each component is visible in the UI and the trainer can weight them. See Traces.
Change the reward without regenerating rollouts
Reward design is iterative. When you change the rubric, the rollouts from earlier steps are still valid inputs, and regenerating them costs GPU time. Pass the rubric to the loop as an input, and choose how to reuse the rollouts:
Recover the finished run with a new rubric. Every action whose inputs didn’t change is reused, and judging and training run again:
def fork_with_new_rubric(run_name: str) -> None:
"""Fork a finished run with a new rubric.
Step 0's `generate` actions are reused, because their inputs are unchanged. `verify_and_judge`
(new rubric input) and `train_step` re-run, and so do later steps' rollouts, which now see
different weights.
"""
new = Rubric(correctness_weight=1.0, efficiency_weight=1.0, tidiness_weight=0.5, name="v2-efficiency")
r = flyte.rerun(run_name, recover=True, rubric=new)
print("forked run:", r.url)
See Recover a failed run.
Start a new run and let the cache find the reuse. Set cache="auto" on the world-building and rollout tasks. A new run with a new rubric gets every rollout whose inputs are unchanged from the cache:
@driver_env.task(report=True)
async def rescore(
n_steps: int = 3,
n_worlds: int = 3,
prompts_per_step: int = 3,
n_samples: int = 2,
) -> list[StepResult]:
"""Re-run training with a new rubric. Cached `build_world` and `generate` actions are reused
wherever their inputs are unchanged; judging and training re-run everywhere.
Step 0 reuses every rollout (same worlds, same v0 weights). The new rewards change the
weights after step 0, so later steps' `generate` inputs differ and regenerate — cache-key
identity finds that invalidation frontier by itself; nothing here knows about "forks".
"""
new = Rubric(correctness_weight=1.0, efficiency_weight=1.0, tidiness_weight=0.5, name="v2-efficiency")
return await train(
n_steps=n_steps, n_worlds=n_worlds, prompts_per_step=prompts_per_step, n_samples=n_samples, rubric=new
)
Either way, the first step’s rollouts are reused, because the worlds and starting weights are the same. After the new reward changes the weights, later rollouts have different inputs and run again. Flyte finds that boundary from the inputs, without any extra code.
Related
- RL for LLMs with GRPO and LoRA: the verifiable reward and group scoring on this page.
- Durable agentic RL with SkyRL and Harbor: the judge on this page, and re-scoring old rollouts with a new rubric.
- Training environment: what produces the trajectories you score.
- Training loop: how rewards become advantages for the policy update.
- Run model-written code safely: sandbox the code a reward executes.
- Evaluation gates: score a model before it is promoted.