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

In multi-turn reinforcement studying (RL), your {custom} reward operate decides what the mannequin truly learns. A subtly mistaken reward can quietly educate the mistaken factor whereas each coaching curve appears to be like wholesome. Designing a reward that holds up over multi-turn, agentic duties is without doubt one of the hardest components of customizing Amazon Nova fashions. For multi-turn coaching, Amazon Nova Forge runs your reward logic in your individual atmosphere via its Deliver Your Personal Orchestration (BYOO) functionality. You possibly can deal with defining what consequence appears to be like like whereas Nova Forge coordinates rollouts, message passing, and dialog state throughout turns. Nova Forge additionally gives a serverless multi-turn RL possibility, now usually out there, for groups that choose to not handle that atmosphere. This submit makes use of the BYOO path.

Amazon Nova gives a number of customization approaches, with reinforcement fine-tuning (RFT) standing out as a result of it may educate fashions the behaviors you need via iterative suggestions. RFT takes a unique method from supervised fine-tuning (SFT). Fairly 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, similar to calling instruments, executing code, or recovering from a mistake. It optimizes cumulative reward throughout the entire trajectory quite 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 process variants whereas SFT degrades. Tailored from Chu et al., 2025

This submit focuses on the reward operate itself: find out how to design a composite multi-turn reward that Group Relative Coverage Optimization (GRPO) can study from. This submit additionally exhibits find out how to execute model-generated code safely contained in the reward, and why to instrument every part so you may 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 find out how to catch them. The code all through is illustrative. Use it as a place to begin in your personal reward implementation.

Conditions

To comply with alongside, you want the next:

• An Amazon Nova Forge subscription, which offers 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 submit, 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 educated reward mannequin. It may be a rule-based verify that verifies the output (reinforcement studying with verifiable rewards), or it may name one other giant 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 doubtless. Nova Forge makes use of GRPO. For every dialog, GRPO makes use of the reward operate to rank Okay mannequin rollouts. GRPO makes use of the highest-ranked mannequin completions to replace the mannequin based on the normalized reward (the benefit) of the batch. RFT with GRPO is a basic approach reaching noticeable efficiency features over preliminary SFT.

A reward sign influences studying solely via 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 due to this fact contributes nothing to the gradient.

How your reward operate runs with Nova Forge relies on the duty. With single-turn RFT, you register the reward as an AWS Lambda operate and level your recipe at it via reward_lambda_arn. Multi-turn duties just like the one on this submit 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 person simulator, executes code, and calls a verifier. It then returns an combination reward per pattern (aggregate_reward_score), plus an non-compulsory record of per-component scores (metrics_list). Half 1 of this sequence covers this infrastructure and its AWS Cloud Growth Equipment (AWS CDK) deployment. This submit 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 process, 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 process, this may embrace a person simulator that responds to the mannequin flip by flip.
• Scores the finished trajectory throughout a number of reward elements (for instance, process correctness, an intermediate-behavior sign, and penalties), reporting every via metrics_list.
• Returns an combination 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 truly rewards, which, as we present, just isn’t at all times what you suppose you wrote.

Selecting the construction of a multi-turn reward

Single scalar rewards are simple to sport, and a single terminal reward is usually too sparse to study from in multi-turn duties. Most manufacturing multi-turn rewards due to this fact mix three sorts of sign: consequence rewards, behavioral rewards, and penalties.

Episode-level (consequence) rewards seize whether or not the ultimate artifact happy the purpose. For instance, did the unit assessments cross, or did the workflow full? They aim the factor you in the end 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 conduct you need, similar to asking earlier than appearing, calling the precise device, or avoiding loops. They’re finest for shaping conduct the end 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 similar to guessing, repeating, or stalling. They separate good and dangerous methods so the optimizer sees a gradient.

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

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

We constructed a multi-turn collaborative-coding process over 500 distinctive programming duties. We educated 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 short, under-specified coding request.
• A person simulator holds the complete 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 mistaken code, whereas asking surfaces the hidden element and results in appropriate code. “Ask first” must be pressured by the duty.

Designing the reward

Make the goal conduct immediately and independently rewardable, and penalize the failure mode explicitly. For this process, the reward is a weighted sum of 4 elements:

Two ideas drive the design. First, un-gate the conduct you need: asked_before_coding is credited by itself, not conditioned on correctness, however 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 via metrics_list:

def asked_before_coding(completion, reply, **kw) -> float:
    msgs = _messages(completion, kw)
    first_q = _first_question_turn(msgs, parser)
    remaining = _final_code(completion, parser)
    dedicated = bool(remaining) and never _is_question(remaining)
    if first_q == 1 and dedicated:
        return 1.0          # requested first, then dedicated (preferrred)
    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 via 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’t forge the end result by writing the anticipated marker to stderr. For execution that requires further isolation, name a devoted sandbox. This harness exhibits 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 rely earlier than scoring
    ...

Additionally validate the variety of assessments truly run towards the quantity anticipated, so the mannequin can’t dilute the rating with its personal trivially-passing assessments. For reward capabilities deployed in dwell environments, implement these safety measures quite than treating them as non-compulsory.

Pitfalls: What makes a reward collapse, and find out how to repair it

Multi-turn reward design has a well known set of failure modes. Reward hacking is the place the mannequin video games a proxy as an alternative of reaching the purpose. 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 often announce themselves in transcripts or in loss and KL curves. Collapse is the harmful one: combination 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 probably the most time on this process, and find out how 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 brought on it. First, the gate sat behind an unreachable situation. Correctness was close to zero on these exhausting duties, so the asking bonus virtually by no means fired. The conduct 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 regarded alike, within-group variation vanished, and studying stopped.

The repair is the design within the earlier part: un-gate the conduct 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 study.

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 elements that also differ preserve combination reward, coverage loss, benefit, and completion size trying 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 mixture reward curve. Combination curves conceal a lifeless channel behind the dwell 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 truly binds to and runs the mannequin’s output. Repair that and make sure 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 via 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 may dismiss a flat reward imply of 0.000 as “these duties are simply exhausting,” 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 guide 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 dwell 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 modifications 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 nice. Un-gate the goal conduct 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 end result time period you care about could by no means get a gradient. Think about down-weighting a shaping time period as soon as it saturates, or up-weighting the end 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 submit use SageMaker HyperPod and Amazon ECS assets that incur price whereas they run. If 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 costs. Take away the rollout information and checkpoints out of your Amazon S3 bucket should you now not want them.

Conclusion

The reward operate is the a part of RFT you design, and it’s the place the delicate failures dwell. In your runs, the mannequin could study the conduct you practice for whereas a time period you care about contributes nothing to studying, with no combination metric revealing it. Higher instrumentation, not a greater algorithm, mounted the problem. Measure every part’s contribution to the benefit, learn transcripts via 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 will have full management over the reward, which suggests 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 assessment and contributions to this submit.

In regards to the authors