Customized reward features for multi-turn reinforcement studying with Amazon Nova Forge

In multi-turn reinforcement studying (RL), your {custom} reward operate decides what the mannequin really learns. A subtly fallacious reward can quietly train the fallacious factor whereas each coaching curve appears to be like wholesome. Designing a reward that holds up over multi-turn, agentic duties is likely one of the hardest elements of customizing Amazon Nova fashions. For multi-turn coaching, Amazon Nova Forge runs your reward logic in your personal atmosphere by way of its Deliver Your Personal Orchestration (BYOO) functionality. You possibly can concentrate on defining what consequence appears to be like like whereas Nova Forge coordinates rollouts, message passing, and dialog state throughout turns. Nova Forge additionally provides a serverless multi-turn RL possibility, now typically accessible, for groups that choose to not handle that atmosphere. This publish makes use of the BYOO path.

Amazon Nova provides a number of customization approaches, with reinforcement fine-tuning (RFT) standing out as a result of it could actually train fashions the behaviors you need by way of iterative suggestions. RFT takes a distinct method from supervised fine-tuning (SFT). Moderately than requiring curated examples with annotated reasoning paths, it learns from analysis alerts on the mannequin’s personal outputs. Multi-turn RFT extends this to brokers that act over a sequence of steps, equivalent to calling instruments, executing code, or recovering from a mistake. It optimizes cumulative reward throughout the entire trajectory reasonably than grading a single response. On the coronary heart of RFT lies the reward operate: the scoring mechanism that guides the mannequin, and the half you design.

Determine 1 — Out-of-distribution (OOD) efficiency after equal-compute post-training from a shared checkpoint. RL improves OOD generalization throughout all job variants whereas SFT degrades. Tailored from Chu et al., 2025

This publish focuses on the reward operate itself: methods to design a composite multi-turn reward that Group Relative Coverage Optimization (GRPO) can be taught from. This publish additionally reveals methods to execute model-generated code safely contained in the reward, and why to instrument every part so you’ll be able to belief what coaching is studying. Half 1 of this sequence covers the Amazon SageMaker HyperPod and Nova Forge infrastructure. It additionally covers the coaching configuration that runs these rewards. We shut with the pitfalls that may quietly collapse a reward, drawn from an actual run the place the highest-weighted part silently contributed no studying sign in any respect. We present methods to catch them. The code all through is illustrative. Use it as a place to begin to your personal reward implementation.

Stipulations

To comply with alongside, you want the next:

• An Amazon Nova Forge subscription, which supplies the Nova Customization SDK and the multi-turn RFT APIs.
• The multi-turn RFT infrastructure from Half 1 of this sequence:
  • An Amazon SageMaker HyperPod cluster, a customer-managed atmosphere on Amazon Elastic Container Service (Amazon ECS).
  • An Amazon Easy Storage Service (Amazon S3) bucket for rollout information and checkpoints.
• The instance code for this publish, together with the reward atmosphere and a walkthrough, from the aws-samples/sample-nova-multi-turn-rl-infra repository.
• The {custom} reward atmosphere is opt-in: in cdk.json, set use_custom_env to “true” and custom_env_id to your atmosphere ID (for instance, “my-custom-env”) earlier than you deploy. By default the stack makes use of the built-in wordle atmosphere.
• Familiarity with reinforcement fine-tuning and GRPO.

Constructing {custom} rewards with Amazon Nova Forge

RFT works by sampling completions from the present mannequin and scoring them with a reward operate. In Nova Forge, the reward operate is a grader you write in code, and never a individually skilled reward mannequin. It may be a rule-based test that verifies the output (reinforcement studying with verifiable rewards), or it could actually name one other massive language mannequin (LLM) to evaluate the response, an method generally known as LLM-as-Decide.

RFT then adjusts the mannequin weights to make higher-reward completions extra possible. Nova Forge makes use of GRPO. For every dialog, GRPO makes use of the reward operate to rank Ok mannequin rollouts. GRPO makes use of the highest-ranked mannequin completions to replace the mannequin in line with the normalized reward (the benefit) of the batch. RFT with GRPO is a basic method attaining noticeable efficiency positive aspects over preliminary SFT.

A reward sign influences studying solely by way of the variation it creates inside a bunch. If a time period takes the identical worth for each completion in a bunch, it contributes nothing to the benefit. It subsequently contributes nothing to the gradient.

How your reward operate runs with Nova Forge is determined by the duty. With single-turn RFT, you register the reward as an AWS Lambda operate and level your recipe at it by way of reward_lambda_arn. Multi-turn duties just like the one on this publish exceed what a single Lambda invocation helps. Multi-turn conversations and long-running scoring run previous the 15-minute Lambda invocation restrict. For these, Nova Forge makes use of BYOO. You set rollout.delegate: true and run your atmosphere and reward logic in an atmosphere container, for instance on Amazon ECS. Nova Forge delegates every rollout to your atmosphere. It then collects the finished episodes again for coaching. Your container manages the multi-turn interplay and dialog state: it runs the consumer simulator, executes code, and calls a verifier. It then returns an mixture reward per pattern (aggregate_reward_score), plus an elective checklist of per-component scores (metrics_list). Half 1 of this sequence covers this infrastructure and its AWS Cloud Improvement Equipment (AWS CDK) deployment. This publish focuses on the reward.

How reward analysis works

The coaching job generates candidate rollouts from the Nova mannequin for every immediate. In a multi-turn job, a rollout is a full episode with a sequence of turns (a trajectory), not a single response. Your reward operate receives every rollout and performs three steps:

• Runs the duty logic. For a conversational job, this may embody a consumer simulator that responds to the mannequin flip by flip.
• Scores the finished trajectory throughout a number of reward parts (for instance, job correctness, an intermediate-behavior sign, and penalties), reporting every by way of metrics_list.
• Returns an mixture reward per rollout (aggregate_reward_score), which coaching turns into within-group benefits.

Determine 2 — A single multi-turn rollout: Nova Forge delegates to your atmosphere container, which asks the simulator or runs the dedicated code, then returns a reward rating for GRPO

This cycle repeats over many coaching steps, progressively shaping the mannequin to maximise cumulative reward throughout the entire sequence. The mannequin optimizes towards no matter your reward really rewards, which, as we present, just isn’t all the time what you assume you wrote.

Selecting the construction of a multi-turn reward

Single scalar rewards are easy to sport, and a single terminal reward is usually too sparse to be taught from in multi-turn duties. Most manufacturing multi-turn rewards subsequently mix three sorts of sign: consequence rewards, behavioral rewards, and penalties.

Episode-level (consequence) rewards seize whether or not the ultimate artifact happy the aim. For instance, did the unit assessments move, or did the workflow full? They aim the factor you finally care about, however they are typically sparse and near-zero early in coaching.

Flip-level (behavioral) rewards seize whether or not the mannequin exhibited the intermediate habits you need, equivalent to asking earlier than appearing, calling the correct device, or avoiding loops. They’re greatest for shaping habits the result reward is simply too sparse to show, although they are often earned with out actual progress if not designed fastidiously. Penalties explicitly discourage a failure mode equivalent to guessing, repeating, or stalling. They separate good and unhealthy methods so the optimizer sees a gradient.

Mix these so the mannequin learns each the habits and the result, with out one part masking or ravenous the opposite. The remainder of this publish makes that concrete. We design a four-component reward for an actual job and execute model-generated code safely inside it. Then we stroll by way of the pitfalls that may collapse such a reward and methods to repair them.

Labored instance: Instructing Amazon Nova Lite 2.0 to ask earlier than coding

We constructed a multi-turn collaborative-coding job over 500 distinctive programming duties. We skilled Amazon Nova Lite 2.0 on it with multi-turn RFT, utilizing GRPO with Low-Rank Adaptation (LoRA), on Amazon SageMaker HyperPod, implementing the reward inside a customer-managed atmosphere container (the Nova Forge BYOO path).

The mechanics are as follows:

• The mannequin sees a quick, under-specified coding request.
• A consumer simulator holds the total specification privately and divulges a element solely when the mannequin asks.
• Every flip, the mannequin both asks a clarifying query or commits code. If it asks, the simulator solutions and the dialog continues. If it commits code, the rollout ends and your reward handler executes that code towards hidden unit assessments to attain correctness. (Working model-generated code safely is a priority we return to later.)

The design intent is that guessing produces fallacious code, whereas asking surfaces the hidden element and results in right code. “Ask first” ought to be pressured by the duty.

Designing the reward

Make the goal habits immediately and independently rewardable, and penalize the failure mode explicitly. For this job, the reward is a weighted sum of 4 parts:

Two ideas drive the design. First, un-gate the habits you need: asked_before_coding is credited by itself, not conditioned on correctness, nevertheless it does require the mannequin to finally commit code, which closes the “ask perpetually, by no means reply” loophole. Second, penalize the failure mode: guessed_immediately makes guessing strictly worse than asking, which restores variation between methods inside a GRPO group, the variation the algorithm wants to supply a gradient.

Name these part scorers contained in the reward handler within the atmosphere container, and report every worth by way of metrics_list:

def asked_before_coding(completion, reply, **kw) -> float:
    msgs = _messages(completion, kw)
    first_q = _first_question_turn(msgs, parser)
    ultimate = _final_code(completion, parser)
    dedicated = bool(ultimate) and never _is_question(ultimate)
    if first_q == 1 and dedicated:
        return 1.0          # requested first, then dedicated (supreme)
    if first_q just isn't None and dedicated:
        return 0.6          # requested later, then dedicated
    return 0.0              # by no means requested, or requested however by no means dedicated

def guessed_immediately(completion, reply, **kw) -> float:
    for m in _assistant_turns(completion, kw):
        code = _code_of(parser.parse(m["content"]))
        return -1.0 if (code and never _is_question(code)) else 0.0
    return 0.0

Executing model-generated code safely

The correctness part runs model-generated code towards unit assessments. Mannequin output beneath RL is optimized by way of exploration, so deal with it as not validated. The container runs in its personal remoted execution atmosphere, however it’s best to nonetheless take precautions. Don’t expose credentials or community to the generated code. Apply useful resource limits and run in a brief listing. Use a per-run random sentinel so the mannequin can not forge the consequence by writing the anticipated marker to stderr. For execution that requires further isolation, name a devoted sandbox. This harness reveals the sample:

import useful resource, secrets and techniques, subprocess, sys, tempfile
from pathlib import Path

def run_tests(code: str, take a look at: str, timeout_s: int = 30) -> float:
    nonce = secrets and techniques.token_hex(8)                       # unforgeable per-run marker
    harness = (
        "import sys, unittest, jsonn"
        f"{code}nn{take a look at}nn"
        'if __name__ == "__main__":n'
        "    r = unittest.TextTestRunner(stream=sys.stderr, verbosity=0).run(n"
        "        unittest.TestLoader().loadTestsFromModule(sys.modules[__name__]))n"
        f"    sys.stderr.write('__{nonce}__' + json.dumps("
        "{'whole': r.testsRun, 'handed': r.testsRun - len(r.failures) - len(r.errors)}) + '__"
        f"{nonce}__')n"
    )
    def _limit():
        useful resource.setrlimit(useful resource.RLIMIT_CPU, (timeout_s, timeout_s))
        useful resource.setrlimit(useful resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3))   # 2 GB
        useful resource.setrlimit(useful resource.RLIMIT_NPROC, (64, 64))
    with tempfile.TemporaryDirectory() as cwd:
        path = Path(cwd) / "h.py"
        path.write_text(harness)
        strive:
            proc = subprocess.run(
                [sys.executable, str(path)], capture_output=True, textual content=True,
                timeout=timeout_s, cwd=cwd, env={"PATH": "/usr/bin"},   # no creds, no community env
                preexec_fn=_limit,
            )
        besides Exception:
            return 0.0
    # parse the nonce-delimited abstract and validate the take a look at depend earlier than scoring
    ...

Additionally validate the variety of assessments really run towards the quantity anticipated, so the mannequin can not dilute the rating with its personal trivially-passing assessments. For reward features deployed in stay environments, implement these safety measures reasonably than treating them as elective.

Pitfalls: What makes a reward collapse, and methods to repair it

Multi-turn reward design has a widely known set of failure modes. Reward hacking is the place the mannequin video games a proxy as a substitute of attaining the aim. Coaching instability is the place updates diverge and entropy collapses or the Kullback-Leibler (KL) time period blows up. Reward collapse is the place the sign degenerates till within-group variation disappears and studying quietly stops. The primary two normally announce themselves in transcripts or in loss and KL curves. Collapse is the damaging one: mixture reward, loss, and completion-length curves can all look wholesome whereas a part you’re relying on contributes nothing. This part covers the 2 collapse failures that price us essentially the most time on this job, and methods to catch them.

When a reward collapses to a single technique

An earlier model of this reward gated the asking bonus behind correctness. You earned the asking reward provided that the ultimate code additionally handed. It additionally added an effectivity time period that rewarded shorter conversations. Coaching collapsed. The mannequin converged to guessing on flip one. The imply reward froze, and the GRPO benefit went to zero.

Two design errors triggered it. First, the gate sat behind an unreachable situation. Correctness was close to zero on these arduous duties, so the asking bonus virtually by no means fired. The habits we needed to reward was invisible to the optimizer. Second, the effectivity time period had a degenerate optimum. Fewer turns maximized it, so the coverage collapsed onto a single, non-committal flip. Each completion seemed alike, within-group variation vanished, and studying stopped.

The repair is the design within the earlier part: un-gate the habits you need, and penalize the failure mode explicitly. With each in place, distinct methods preserve producing distinct rewards inside a bunch, which preserves the variance GRPO must be taught.

Silently lifeless part

When a reward part returns the identical worth for each completion in a GRPO group, its within-group variance is zero. In consequence, it contributes nothing to the benefit or the gradient, even on the highest weight. The parts that also range preserve mixture reward, coverage loss, benefit, and completion size wanting wholesome, so the curves by no means reveal it. One frequent trigger in code rewards is a correctness scorer that returns 0 on each rollout as a result of the harness by no means executes the mannequin’s output. This may occur due to mismatched entry-point names, failed imports, or a setup error that makes each take a look at fail earlier than its assertions run. In our run, that is precisely what occurred: the mannequin’s clarifying-question fee rose from roughly 34–96 p.c. Code correctness barely moved, as a result of the correctness scorer was returning the identical worth on each rollout.

To catch a lifeless part, observe every part’s within-group commonplace deviation, not the combination reward curve. Combination curves disguise a lifeless channel behind the stay ones. If that unfold sits at or close to zero, the part isn’t coaching, no matter its weight. The standard root trigger in code rewards is a correctness scorer caught at 0 as a result of the harness by no means really binds to and runs the mannequin’s output. Repair that and ensure the unfold turns into non-zero.

Instrument so that you catch these early

A number of habits catch these failures, and would have caught ours on day one:

1. Instrument per-component contribution to the benefit, not simply per-component reward. Report every part by way of metrics_list, and observe its imply and its within-group commonplace deviation. Any part with near-zero within-group variance contributes nothing to studying, no matter its weight. You would possibly dismiss a flat reward imply of 0.000 as “these duties are simply arduous,” however a flat within-group variance is unambiguous. Automate this as a per-component advantage-variance panel so lifeless channels are flagged mechanically, with out handbook inspection.
2. Learn transcripts sorted by the part you’re testing, not by whole reward. Sorting by whole reward hides a lifeless part behind the stay ones. Sorting by the suspect part surfaces the issue instantly.
3. Ablate or revive each part you declare is doing work. If eradicating a part adjustments nothing, it was not doing work. If reviving a part recovers a metric you assumed was already optimized, it was not within the goal.
4. Design for within-group variance. GRPO learns from variations between completions of the identical immediate. Unreachable gates, degenerate shaping optima, and saturating phrases all collapse that variation and cease studying even when the reward appears to be like superb. Un-gate the goal habits and penalize the failure mode so methods separate.
5. Look ahead to one dense reward ravenous one other. As soon as our dense asking reward saturated, the sparse correctness reward couldn’t transfer the coverage. If a behavioral shaping time period dominates, the result time period you care about might by no means get a gradient. Contemplate down-weighting a shaping time period as soon as it saturates, or up-weighting the result time period.
6. Deal with mannequin output as not validated. Sandbox any execution of generated code (no credentials, no community, useful resource limits) and make verifiers unforgeable (random sentinels, test-count validation).

Clear up

The coaching run and atmosphere on this publish use SageMaker HyperPod and Amazon ECS assets that incur price whereas they run. Whenever you end experimenting, comply with the teardown steps in Half 1 of this sequence to delete the SageMaker HyperPod cluster and the Amazon ECS atmosphere, which stops the biggest fees. Take away the rollout information and checkpoints out of your Amazon S3 bucket if you happen to not want them.

Conclusion

The reward operate is the a part of RFT you design, and it’s the place the delicate failures stay. In your runs, the mannequin might be taught the habits you prepare for whereas a time period you care about contributes nothing to studying, with no mixture metric revealing it. Higher instrumentation, not a greater algorithm, mounted the problem. Measure every part’s contribution to the benefit, learn transcripts by way of the lens of the part you’re testing, and ablate what you declare is working. With a {custom} reward operate on Amazon Nova Forge you’ve full management over the reward, which implies the duty for getting it proper is yours. For the infrastructure and AWS CDK deployment that make these runs reproducible, see Half 1 of this sequence.

Acknowledgements

Particular because of Mahima Chaudhary for his or her overview and contributions to this publish.

In regards to the authors