Custom reward functions for multi-turn reinforcement learning with Amazon Nova Forge
Bad reward design in multi-turn reinforcement fine-tuning (RFT) can produce deceptively good training curves while teaching an agent the wrong behavior, here’s why and how to avoid it.
One experiment the Nova Forge team shared makes the point bluntly: a shaping bonus that only paid when final code passed tests effectively produced identical values across completions in many groups, so the optimizer ignored it, and the policy learned to guess on turn one. The curves looked healthy. The behavior was wrong.
The piece you control vs. the plumbing Nova Forge provides
Amazon Nova Forge supplies the multi-turn RFT plumbing: a Nova Customization SDK, multi-turn RFT APIs, and two deployment patterns for your evaluation logic. You can run evaluators serverless (now generally available per the Nova Forge post) or run them in a customer-managed container via BYOO (Bring Your Own Orchestration) when tasks exceed short execution windows.
Your environment must provide the executable pieces that actually judge behavior: a user simulator (holds hidden task specs and answers clarifying questions), a safe test harness for any generated code, and a grader that returns an aggregate_reward_score plus a metrics_list of per-component scores for every rollout. Example code and an infra starting point are available in the Nova Forge sample repo at github.com/aws-samples/sample-nova-multi-turn-rl-infra.
The core principle: variance is the signal
Nova Forge’s training loop (the post describes Group Relative Policy Optimization, or GRPO) ranks K candidate completions for the same conversation and computes normalized advantages within that group. In plain language: updates come from how candidates compare to each other, not from their absolute scores.
That means a reward component influences learning only to the extent it varies across the K completions in a group. If a term returns the same value for every completion in that group, it creates zero within-group variance and contributes nothing to the advantage, and therefore nothing to the gradient.
“A reward signal influences learning only through the variation it creates within a group. If a term takes the same value for every completion in a group, it contributes nothing to the advantage. It therefore contributes nothing to the gradient.”, from the Nova Forge post
Micro-example: imagine a group of four completions and a component X whose values are [1, 0, 0, 0]. Component X creates a strong learning signal because it separates the top completion from the rest. If X returns [1, 1, 1, 1], it creates no signal and GRPO won’t learn from it even if you weight it highly.
Worked example (from the Nova Forge post): teach Nova Lite 2.0 to ask before coding
Illustrative run: the team fine-tuned Amazon Nova Lite 2.0 across 500 unique programming tasks using LoRA. The environment simulated users and validated final code against hidden unit tests. Each rollout returned an aggregate_reward_score and a metrics_list so training could report per-component contributions.
- correctness (weight 1.0): fraction of hidden unit tests passing on the final code.
- asked_before_coding (weight 0.6): 1.0 if the model asked on turn 1 then committed; 0.6 if it asked later then committed; else 0.
- guessed_immediately (weight 0.4): penalty, -1.0 if the first turn is code with no question.
- loop_penalty (weight 0.2): -0.5 if the last two turns are more than 80% similar.
These weights are an illustrative configuration from the Nova Forge example. They show how outcome reward (correctness) can be blended with shaping terms (asking, not guessing) and an efficiency penalty. They are a starting point, not a universal default. Tune them per task with ablations and sensitivity sweeps.
Failure modes you should expect (and how they surface)
Three failure patterns appeared in the Nova Forge runs and are common in multi-turn RFT:
- Reward hacking. The policy finds shortcuts that maximize shaping terms without solving the true task. Dense shaping rewards (frequent signals) can swamp the sparse outcome signal.
- Training instability. Optimization can diverge, entropy can collapse, or KL terms can blow up if hyperparameters or update schedules are misaligned.
- Reward collapse (the “silently dead component”). A component that returns the same value across all completions in a group provides no learning signal. In one Nova Forge run a correctness scorer never executed the model output (due to a harness import/setup error), producing effectively constant scores. The clarifying-question rate rose from ~34% to ~96% while correctness hardly moved, a strong indicator the outcome scorer was silent.
Other practical pitfalls to watch for:
- Test overfitting. The model learns specifics of the hidden unit tests rather than the underlying problem. Rotate/expand tests and use invariant properties where possible.
- Flaky tests. Non‑deterministic unit tests inject noise into rewards. Ensure determinism or guard scores with repeated runs.
- Grader drift and bugs. Changes to grader code or its environment can silently change reward scales. Version and CI your grader like any production service.
Safe execution of model-generated code, practical, concrete protections
Treat model output as unvalidated. The Nova Forge team recommends running code in a disposable sandbox and applying defense‑in‑depth:
- No credentials and no network access.
- OS resource limits, example harness settings shown in the example: RLIMIT_AS set to (2 * 1024**3, 2 * 1024**3) (2 GB) and RLIMIT_NPROC set to (64, 64).
- Per-run timeouts (example: default timeout_s = 30 seconds on unit tests) and ephemeral temp directories.
- Unforgeable nonces (example: secrets.token_hex(8)) written to stderr alongside a JSON summary so harness output can be parsed and validated.
- Validation of test counts and expected artifacts to ensure the harness actually executed the model output.
“Model output under RL is optimized through exploration, so treat it as not validated.”, from the Nova Forge post
Caveats: these are example protections from the sample harness. Additional layers include container isolation, seccomp filters, signed artifacts, reproducible build environments, and separate audit logs, which further reduce risk. Always run grader and harness code through code review, CI, and security scanning.
Instrument ruthlessly, what to capture first
Since GRPO learns from within-group differences, the telemetry that matters is per-component statistics, not just an aggregate score. Prioritize this minimal, actionable telemetry first:
- Per-component within-group standard deviation (high priority). Flag components whose stddev is near-zero relative to peers.
- Per-component contribution to the advantage (high priority). Compute each component’s contribution to the normalized reward used by GRPO.
- Transcripts for top/bottom deciles (medium). Sort rollouts by a component (e.g., asked_before_coding) and inspect examples from the extremes.
- Ablation/revival pipelines (medium). Remove a shaping term to confirm its effect, then reintroduce it with adjusted weight.
- Longitudinal dashboards for drift (ongoing). Track metric drift and grader behavior over time.
Pragmatic heuristics to start with: flag components whose median within-group stddev is less than 1% of the median stddev across components, or whose explained variance across groups is below 0.01. These are conservative heuristics, tune thresholds to your task.
Also randomize and stratify how K completions are sampled into groups (shuffle seeds, vary temperatures, include diverse model checkpoints). Poor group composition can artificially reduce variance and hide learning signals.
Deployment trade-offs: BYOO vs serverless
Two deployment paths exist for the evaluation environment:
- Serverless multi-turn RL (generally available in Nova Forge): lower ops overhead and simpler scaling for short, stateless evaluators. Note that serverless MTRL outputs may be packaged as Restricted Model Packages (RMPs), which affects artifact access and governance.
- BYOO (Bring Your Own Orchestration): run evaluators in your container (ECS, SageMaker HyperPod, etc.). BYOO removes Lambda-style timeout limits and gives full control for stateful simulators and long-running evaluations, at the cost of more operational surface and higher complexity.
Choice depends on execution duration, statefulness, compliance needs, and cost. If you pick BYOO, ensure your orchestration includes deployment controls and CI for the grader and harness code. If you pick serverless, confirm how outputs are packaged and who can access the resulting RMPs.
Checklist before flip‑the‑switch (operational & reviewer checks)
- Return a metrics_list with every rollout and record per-component within-group stddev.
- Sanity‑check the harness: confirm it executes model output, emits the unforgeable nonce, and reports the expected number of tests.
- Sort transcripts by component and read samples from the top and bottom deciles.
- Ablate shaping terms to verify their contribution, then revive with adjusted weights if needed.
- Sandbox execution (no creds/network), enforce RLIMITs and timeouts, and validate test counts and sentinels on stderr.
- Reviewer questions: can the grader’s behavior be reproduced locally? Are grader and harness in source control with CI? Is there an audit trail and signed deployments for any grader change or RMP?
Operational and governance implications
Moving evaluation logic into executable grader code changes ML‑ops responsibilities. The grader, harness, and user simulator become production artifacts that require code review, CI/CD, access control, and audit logs. LoRA is a pragmatic parameter‑efficient adaptation used in the example run to reduce compute and iteration time, but the reward channels remain the primary lever for safety and behavior.
“With a custom reward function on Amazon Nova Forge you have full control over the reward, which means the responsibility for getting it right is yours.”, Maria Masood, Nick Biso, Laurent Mombaerts
Key takeaways: common questions and honest answers
-
How does a reward component actually influence learning?
Only through the variation it creates within a GRPO group; if a component returns the same value across all completions in a group, it contributes no advantage and produces no gradient. Action: monitor per-component within-group stddev and trigger ablation checks when variance is near-zero.
-
Should I use dense shaping rewards to speed up learning?
Dense shaping can accelerate early progress but can also create degenerate optima that ignore the sparse outcome. Action: start with small shaping weights, monitor variance, and down‑weight shaping terms if they saturate the signal.
-
When do I need BYOO instead of serverless?
Choose BYOO for long‑running, stateful, or tool‑dependent evaluations that exceed serverless runtime limits or need custom dependencies; choose serverless for simpler, shorter evaluators to reduce ops burden. Action: weigh execution time, stateful needs, governance, and cost before deciding.
-
How should I run model-generated code safely?
Run it in a sandbox with no credentials/network, OS-level resource limits (example: 2 GB address space, limited processes), short timeouts (example: 30 s), per-run nonces, and verification of expected test counts. Action: add container isolation, signed artifacts, and separate audit logs to strengthen defenses.
-
How do I detect a “silently dead” reward component?
Automate per-component within-group standard deviation panels; a near-zero stddev across groups indicates the component provides no learning signal. Action: inspect the harness for execution failures, mismatched outputs, or grader bugs; run the same input locally to reproduce the problem.
Reward design for multi-turn RFT sits at the intersection of ML engineering and product design. State your goal, then build incentives that produce the desired sequence of behaviors under exploration. Instrument each channel, isolate execution risk, and treat grader code with the same rigor you apply to your models.
Special thanks to Mahima Chaudhary for review and to Maria Masood, Nick Biso, and Laurent Mombaerts for the worked example and practical guidance.