Funções de recompensa personalizadas para aprendizado por reforço multi-turno com Amazon Nova Forge

O coração do treinamento por reforço: a função de recompensa

No treinamento de modelos de linguagem com aprendizado por reforço (RL), a função de recompensa é o mecanismo que define o que o modelo aprende. Uma recompensa sutilmente errada pode ensinar o comportamento errado enquanto todas as métricas de treinamento parecem saudáveis. Projetar uma recompensa robusta para tarefas agênticas multi-turno é um dos desafios mais complexos ao personalizar os modelos Amazon Nova.

A AWS publicou um guia técnico aprofundado sobre esse tema, focando especificamente no Amazon Nova Forge e na capacidade Bring Your Own Orchestration (BYOO) — que permite executar a lógica de recompensa no próprio ambiente do cliente. O post explora como desenhar uma recompensa composta, como executar código gerado pelo modelo com segurança e como instrumentar os componentes para que o time possa confiar no que o treinamento está aprendendo de fato.

RFT vs. SFT: por que o ajuste fino por reforço se destaca

O Amazon Nova oferece múltiplas abordagens de personalização. Entre elas, o ajuste fino por reforço (RFT) se destaca por ensinar comportamentos desejados por meio de feedback iterativo — sem exigir exemplos curados com raciocínios anotados, como ocorre no ajuste fino supervisionado (SFT).

O RFT multi-turno vai além: ele otimiza a recompensa acumulada ao longo de uma sequência de passos, como chamar ferramentas, executar código ou se recuperar de erros. O modelo aprende a partir de sinais de avaliação sobre suas próprias saídas.

O algoritmo utilizado pelo Nova Forge é o Otimização de Política Relativa de Grupo (GRPO). Para cada conversa, o GRPO usa a função de recompensa para ranquear K execuções do modelo e atualiza os pesos com base na vantagem normalizada do lote. Um ponto crítico: um sinal de recompensa só influencia o aprendizado se criar variação dentro do grupo. Se um componente retorna o mesmo valor para todas as execuções, ele não contribui nada para o gradiente.

Como o Nova Forge executa recompensas multi-turno

Para tarefas de turno único, o Nova Forge aceita a função de recompensa como uma função AWS Lambda. Já para tarefas multi-turno — que ultrapassam o limite de 15 minutos de invocação do Lambda — o caminho é o BYOO. O time configura rollout.delegate: true e executa a lógica de recompensa em um container próprio, por exemplo no Amazon ECS.

O container gerencia toda a interação multi-turno: executa o simulador de usuário, roda o código gerado e chama o verificador. Ao final, retorna uma recompensa agregada por amostra (aggregate_reward_score) e, opcionalmente, pontuações por componente (metrics_list). O Parte 1 desta série cobre a infraestrutura com Amazon SageMaker HyperPod e o deployment via AWS CDK.

Estruturando uma recompensa multi-turno

Recompensas escalares simples são fáceis de burlar, e uma única recompensa terminal costuma ser esparsa demais para o modelo aprender em tarefas multi-turno. A abordagem recomendada combina três tipos de sinal:

• Recompensas de resultado (nível de episódio): capturam se o artefato final atingiu o objetivo — por exemplo, se os testes unitários passaram. São esparsas no início do treinamento.
• Recompensas comportamentais (nível de turno): capturam se o modelo exibiu o comportamento intermediário desejado, como perguntar antes de agir ou chamar a ferramenta correta. São úteis para moldar comportamentos que a recompensa de resultado é esparsa demais para ensinar.
• Penalidades: desencorajam explicitamente modos de falha como adivinhar, repetir ou travar. Separam estratégias boas e ruins para que o otimizador enxergue um gradiente.

Exemplo prático: ensinando o Nova Lite 2.0 a perguntar antes de codificar

A AWS construiu uma tarefa de codificação colaborativa multi-turno com 500 problemas de programação únicos. O Amazon Nova Lite 2.0 foi treinado nessa tarefa com RFT multi-turno, usando GRPO com Adaptação de Baixo Rank (LoRA), no Amazon SageMaker HyperPod, com a recompensa implementada dentro de um container de ambiente gerenciado pelo cliente (o caminho BYOO do Nova Forge).

A mecânica é a seguinte: o modelo recebe uma solicitação de código vaga e subespecificada. Um simulador de usuário guarda a especificação completa em privado e revela um detalhe apenas quando o modelo pergunta. Em cada turno, o modelo pode fazer uma pergunta de esclarecimento ou submeter código. Se perguntar, o simulador responde e a conversa continua. Se submeter código, o rollout termina e o handler de recompensa executa esse código contra testes unitários ocultos.

A recompensa é uma soma ponderada de quatro componentes:

• correctness (peso 1.0): fração dos testes unitários ocultos que passam no código final.
• asked_before_coding (peso 0.6): 1.0 se perguntou no turno 1 e depois submeteu; 0.6 se perguntou depois e submeteu; 0 caso contrário.
• guessed_immediately (peso 0.4): penalidade de -1.0 se o primeiro turno foi código sem nenhuma pergunta.
• loop_penalty (peso 0.2): -0.5 se os dois últimos turnos forem mais de 80% similares.

Dois princípios guiam o design: primeiro, não condicionar o comportamento desejado a outro componente — asked_before_coding é creditado por conta própria, não condicionado à correção do código. Segundo, penalizar explicitamente o modo de falha — guessed_immediately torna adivinhar estritamente pior do que perguntar, o que preserva a variação entre estratégias dentro de um grupo GRPO.

def asked_before_coding(completion, answer, **kw) -> float:
    msgs = _messages(completion, kw)
    first_q = _first_question_turn(msgs, parser)
    final = _final_code(completion, parser)
    committed = bool(final) and not _is_question(final)
    if first_q == 1 and committed:
        return 1.0  # asked first, then committed (ideal)
    if first_q is not None and committed:
        return 0.6  # asked later, then committed
    return 0.0  # never asked, or asked but never committed

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

Executando código gerado pelo modelo com segurança

O componente de correção executa código gerado pelo modelo contra testes unitários. Sob RL, esse código é produto de exploração — trate-o como não validado. O container deve rodar em ambiente de execução isolado, sem expor credenciais ou rede ao código gerado, com limites de recursos e em diretório temporário. Um sentinel aleatório por execução impede que o modelo forje o resultado escrevendo o marcador esperado na saída de erro.

import resource, secrets, subprocess, sys, tempfile
from pathlib import Path

def run_tests(code: str, test: str, timeout_s: int = 30) -> float:
    nonce = secrets.token_hex(8)  # unforgeable per-run marker
    harness = (
        "import sys, unittest, json\n"
        f"{code}\n\n{test}\n\n"
        '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("
        "{'total': r.testsRun, 'passed': r.testsRun - len(r.failures) - len(r.errors)}) + '__"
        f"{nonce}__')\n"
    )

    def _limit():
        resource.setrlimit(resource.RLIMIT_CPU, (timeout_s, timeout_s))
        resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3))  # 2 GB
        resource.setrlimit(resource.RLIMIT_NPROC, (64, 64))

    with tempfile.TemporaryDirectory() as cwd:
        path = Path(cwd) / "h.py"
        path.write_text(harness)
        try:
            proc = subprocess.run(
                [sys.executable, str(path)],
                capture_output=True,
                text=True,
                timeout=timeout_s,
                cwd=cwd,
                env={"PATH": "/usr/bin"},  # no creds, no network env
                preexec_fn=_limit,
            )
        except Exception:
            return 0.0
        # parse the nonce-delimited summary and validate the test count before scoring
        ...

Também é importante validar o número de testes efetivamente executados contra o número esperado, para impedir que o modelo dilua a pontuação com testes trivialmente aprovados que ele mesmo escreveu.

Armadilhas que colapsam uma recompensa — e como detectá-las

O design de recompensas multi-turno tem um conjunto bem conhecido de modos de falha. O mais perigoso é o colapso silencioso: a recompensa agregada, a perda e as curvas de comprimento de resposta podem parecer saudáveis enquanto um componente que você conta não contribui nada para o aprendizado.

Quando a recompensa colapsa para uma única estratégia

Uma versão anterior da recompensa condicionava o bônus de “perguntar” à correção do código — você só ganhava o bônus se o código final também passasse nos testes. Também havia um termo de eficiência que recompensava conversas mais curtas. O treinamento colapsou: o modelo convergiu para adivinhar no primeiro turno. A vantagem GRPO foi a zero e o aprendizado parou.

Dois erros de design causaram isso. Primeiro, o gate estava atrás de uma condição inalcançável — a correção era próxima de zero nas tarefas difíceis, então o bônus de perguntar quase nunca disparava. Segundo, o termo de eficiência tinha um ótimo degenerado — menos turnos o maximizava, então a política colapsou para um único turno não-comprometido. A correção é o design descrito acima: descondicionar o comportamento desejado e penalizar o modo de falha explicitamente.

Componente silenciosamente morto

Quando um componente de recompensa retorna o mesmo valor para todas as execuções de um grupo GRPO, sua variância intra-grupo é zero. Ele não contribui nada para a vantagem nem para o gradiente — mesmo com peso máximo. Os componentes que ainda variam mantêm a recompensa agregada, a perda e o comprimento de resposta parecendo saudáveis, então as curvas nunca revelam o problema.

Na execução descrita no artigo original, isso aconteceu exatamente assim: a taxa de perguntas de esclarecimento do modelo subiu de cerca de 34% para 96%, mas a correção do código mal se moveu — porque o scorer de correção retornava o mesmo valor em cada rollout. A causa raiz era que o harness nunca executava de fato o código do modelo.

Como instrumentar para pegar esses problemas cedo

• Acompanhe o desvio padrão intra-grupo de cada componente, não apenas a recompensa agregada. Qualquer componente com variância intra-grupo próxima de zero não está treinando, independentemente do seu peso.
• Leia transcrições ordenadas pelo componente que está testando, não pela recompensa total. Ordenar por recompensa total esconde um componente morto atrás dos vivos.
• Faça ablação de cada componente que você afirma estar funcionando. Se remover um componente não muda nada, ele não estava contribuindo.
• Projete para variância intra-grupo. Gates inalcançáveis, ótimos degenerados e termos saturados colapsam essa variação e param o aprendizado mesmo quando a recompensa parece boa.
• Fique atento a um componente denso sufocando outro. Quando o sinal comportamental satura, o sinal de resultado pode nunca receber gradiente. Considere reduzir o peso do componente de molde quando ele saturar.

Pré-requisitos para seguir o exemplo

Para reproduzir o ambiente descrito, são necessários: uma assinatura do Amazon Nova Forge (que fornece o Nova Customization SDK e as APIs de RFT multi-turno); a infraestrutura de RFT multi-turno da Parte 1 (cluster Amazon SageMaker HyperPod, ambiente gerenciado no Amazon ECS, bucket Amazon S3 para dados de rollout e checkpoints); e o código de exemplo do repositório aws-samples/sample-nova-multi-turn-rl-infra.

O ambiente de recompensa personalizado é opt-in: no arquivo cdk.json, defina use_custom_env como "true" e custom_env_id com o ID do seu ambiente antes de fazer o deploy. Por padrão, o stack usa o ambiente built-in de wordle.

Conclusão

A função de recompensa é a parte do RFT que o time projeta — e é onde vivem as falhas sutis. O modelo pode aprender o comportamento treinado enquanto um componente que você considera importante não contribui nada para o aprendizado, sem que nenhuma métrica agregada revele isso. Melhor instrumentação, não um algoritmo melhor, é o que resolve esse problema.

Com uma função de recompensa personalizada no Amazon Nova Forge, o time tem controle total sobre a recompensa — o que significa que a responsabilidade de acertar é também do time. Para a infraestrutura e o deployment via AWS CDK que tornam essas execuções reproduzíveis, consulte a Parte 1 desta série.

Fonte

Custom reward functions for multi-turn reinforcement learning with Amazon Nova Forge (https://aws.amazon.com/blogs/machine-learning/custom-reward-functions-for-multi-turn-reinforcement-learning-with-amazon-nova-forge/)