Boas práticas de aprendizado por reforço multi-turno no Amazon SageMaker AI

O desafio do RL multi-turno

Treinar um agente para resolver tickets de suporte ou moderar conteúdo é bem diferente de treinar um modelo para responder a uma pergunta isolada. Em cenários multi-turno, o agente precisa ler instruções, chamar ferramentas, interpretar os resultados, decidir a próxima ação e, quando necessário, se recuperar de um erro antes de chegar a uma resposta final. É exatamente essa flexibilidade que torna o Aprendizado por Reforço (RL) agnóstico tão desafiador: quanto mais caminhos de ação disponíveis, maior a chance do modelo encontrar formas de satisfazer a recompensa sem realmente resolver a tarefa.

A AWS publicou um guia técnico detalhado sobre boas práticas para treinar agentes com Amazon SageMaker AI multi-turn RL (SageMaker AI MTRL) de forma confiável. Os exemplos ao longo do artigo são baseados no SOP-Bench, um benchmark da Amazon Science que avalia a capacidade de agentes resolverem tarefas baseadas em Procedimentos Operacionais Padrão (SOP) em 12 domínios de negócios.

O que o SageMaker AI MTRL oferece

O SageMaker AI MTRL fornece o loop de treinamento para tarefas agênticas. O agente pode rodar no Amazon Bedrock AgentCore, no Amazon Elastic Kubernetes Service (Amazon EKS), no Amazon Elastic Compute Cloud (Amazon EC2), no AWS Fargate ou em infraestrutura própria. A conexão é feita por um adaptador que expõe as ferramentas disponíveis ao servidor de rollout.

Entre os principais recursos do serviço estão:

• Interface modular agente-ambiente que mantém a integração com pouco código, mas oferece controle algorítmico completo.
• Execução serverless com precificação por token, sem necessidade de provisionar clusters de GPU.
• Rollout e coleta de trajetórias assíncronos com staleness off-policy limitado — geração e atualizações de gradiente correm em paralelo sem se afastar demais da política atual.
• Biblioteca nativa de algoritmos com Proximal Policy Optimization (PPO), Clipped Importance Sampling Policy Optimization (CISPO) e perdas de importance-sampling (IS), além de múltiplos estimadores de vantagem baseados em grupo (GRPO, GRPO pass@k, RLOO, entre outros).
• Treinamento com extensão de sequência para manter o tempo de execução razoável em trajetórias longas.
• Observabilidade de trajetórias e recompensas via MLflow gerenciado pelo Amazon SageMaker AI.

Construa um ambiente de treinamento confiável

RL de turno único precisa apenas de um prompt e uma função de recompensa. O RL multi-turno exige também um ambiente com o qual o agente interage ao longo dos turnos — as ferramentas que ele chama e os sistemas por trás delas. Esse ambiente faz parte da configuração de treinamento e influencia diretamente o que o modelo aprende e se as métricas são confiáveis.

A recomendação da AWS é construir um ambiente sandboxed ou simulado que replique a produção, mas permaneça isolado do tráfego real. Um treinamento típico gera milhares de rollouts, cada um com várias chamadas de ferramentas — por exemplo, um batch de 128 com group size 8 resulta em 1.024 rollouts por passo. Apontar esse volume para sistemas reais pode causar impacto em clientes e gerar efeitos colaterais indesejados, como reembolsos disparados, registros deletados ou fluxos acionados por engano.

Três padrões cobrem a maioria dos casos de uso:

• Ferramentas somente leitura: replique respostas gravadas indexadas pelos argumentos de entrada. No SOP-Bench, por exemplo, dez ferramentas mockadas retornam respostas determinísticas a partir de fixtures.
• Ferramentas com estado: sandboxes com estado inicializado por seed, válido pelo tempo de um episódio. Recursos são alocados no início do rollout e destruídos em um bloco try/finally ao final, sem vazamento de estado entre episódios.
• Resultados verificáveis: execução real em ambiente isolado. Para código, SQL ou matemática, é possível usar Docker exec, SQLite em memória por rollout ou eval Python puro. O AgentCore Code Interpreter oferece ambientes gerenciados e isolados para execução de código.

Independentemente do padrão escolhido, duas propriedades devem ser mantidas: reprodutibilidade (a mesma ferramenta com os mesmos argumentos retorna sempre o mesmo resultado) e representatividade (o ambiente deve ser construído a partir dos schemas e distribuições de dados reais de produção).

Configure uma avaliação externa antes de treinar

Antes de escrever qualquer função de recompensa, é essencial ter uma forma independente de medir o sucesso. O RL otimiza literalmente o sinal de recompensa — se a recompensa for o único número monitorado, não é possível separar progresso real na tarefa de progresso em satisfazer os critérios de recompensa.

A recomendação é criar uma avaliação externa em um conjunto de teste fixo, calculada de forma independente da recompensa. No SOP-Bench, essa avaliação é um match exato no objeto JSON final dentro das tags <final_output>: todos os campos precisam corresponder ao ground truth, ou o rollout recebe zero.

Antes de qualquer treinamento, é importante estabelecer uma baseline: rodar o modelo base e um modelo de referência (como um modelo frontier hospedado no Amazon Bedrock) pela mesma avaliação. Isso mostra o quanto o modelo base precisa evoluir e o que “bom” significa para aquela tarefa específica.

Um anti-padrão comum é usar a recompensa de treinamento como medida de sucesso. Agentes multi-turno são especialmente suscetíveis a isso: uma recompensa que paga por chamadas de ferramentas ensina o agente a chamar o máximo de ferramentas possível; uma recompensa que penaliza o número de turnos ensina o agente a responder antes de ter as informações necessárias. Em ambos os casos, a recompensa de treinamento sobe, mas o sucesso real na tarefa cai.

Desenhando uma boa função de recompensa multi-turno

O design de recompensa é um dos problemas mais desafiadores em RL. A mesma flexibilidade que permite ao agente resolver uma tarefa real também permite que ele encontre formas de satisfazer a recompensa sem resolver a tarefa. O modelo otimiza exatamente o que foi escrito, não o que foi intencionado.

A orientação padrão é usar a mesma regra de pontuação para treinamento e avaliação, desviando apenas quando há uma razão concreta. Existem dois motivos legítimos para usar uma recompensa mais densa:

• Razão algorítmica: uma pontuação binária pode colapsar a variância entre rollouts de um grupo. Quando todos os rollouts de um grupo pontuam igual, o sinal relativo é zero e o grupo não contribui com gradiente.
• Velocidade de convergência: uma recompensa densa dá ao modelo gradiente em direção ao progresso parcial em cada rollout, não apenas nos que tiveram sucesso completo.

O artigo apresenta um exemplo de recompensa densa por campo para o SOP-Bench:

class SOPBenchReward:
    """Dense per-field reward for the SOP-Bench aircraft-inspection task.
    Returns a scalar in [0, 1] plus a metrics dict surfaced in MLflow."""
    ground_truth: dict[str, str]
    format_coef: float = 0.1  # format is a small shaping term, not the objective

    async def __call__(self, history: list[Message]) -> tuple[float, dict[str, float]]:
        fields = parse_final_output(last_assistant(history))  # JSON inside <final_output>
        emitted = float(fields is not None)
        if fields is None:  # no parseable answer
            return self.format_coef * (emitted - 1), {"completion": 0.0, "field_acc": 0.0}
        matched = sum(1 for k, v in self.ground_truth.items()
                      if str(fields.get(k)).strip().lower() == str(v).strip().lower())
        field_acc = matched / len(self.ground_truth)  # partial credit: 5/6 > 0
        reward = field_acc + self.format_coef * (emitted - 1)  # correctness dominates
        return reward, {"completion": emitted, "field_acc": field_acc}

Um ponto crítico destacado no artigo: um parser de recompensa mais permissivo do que a avaliação é em si uma forma de reward hacking. Em um dos testes com o SOP-Bench, o parser de recompensa aceitava um wrapper <final_response> enquanto o benchmark só lia <final_output>. O modelo aprendeu exatamente o que foi pedido: a recompensa de treinamento subiu, mas a avaliação externa caiu.

Para mais detalhes sobre design de recompensa, recompensas esparsas vs. densas e modelos juízes, a AWS disponibiliza as melhores práticas de design de recompensa do SageMaker AI.

Gerenciando contexto e orçamento de turnos

Agentes multi-turno precisam lidar com questões que o RL de turno único não enfrenta. Cada chamada de ferramenta estende a conversa — o call, os argumentos, o resultado e o raciocínio do modelo entre eles. Trajetórias longas acumulam contexto rapidamente.

Dois orçamentos limitam isso: max_turns (controlado pelo loop do agente) e o orçamento de tokens por turno (definido via sampling_max_tokens para rollout e val_sampling_params.sampling_max_tokens para avaliação). A recomendação é calibrar ambos para o que a tarefa exige. No SOP-Bench, oito turnos e 2.048 tokens por turno cobrem o procedimento canônico com margem.

Uma regra prática: se um walkthrough humano da tarefa leva N turnos, defina max_turns = ceil(N * 1.5) no loop do agente. Se mais de 5% dos rollouts atingem o limite de sampling_max_tokens, aumente esse valor — caso contrário, o modelo aprende de trajetórias truncadas sem ver a recompensa que teria recebido ao terminar.

Separando conclusão de correção

Uma trajetória que termina com a resposta errada e uma que nunca termina são falhas diferentes. Confundi-las esconde onde o modelo está quebrando. As métricas das famílias rollout e val no MLflow fornecem ambos os sinais separadamente:

• rollout/reward/mean: recompensa média das trajetórias
• rollout/reward/zero_frac: fração de trajetórias que pontuaram exatamente zero
• rollout/turns/mean: média de turnos por trajetória
• analysis/zero_adv_groups: grupos onde todos os rollouts pontuaram igual, desperdiçando rollouts
• val/reward/mean: recompensa média de validação
• val/reward/pass_k_1, pass_k_8: pass@1 e pass@k no conjunto de validação

Um val/reward/pass_k_1 alto com baixa taxa de conclusão indica que o modelo acerta os casos fáceis, mas trava nos difíceis — sugerindo ajuste no orçamento de turnos. Uma alta taxa de conclusão com val/reward/pass_k_1 baixo indica que o agente responde com fluência, mas de forma errada — sugerindo redesenho da recompensa.

Monitorando o treinamento e iterando

O SageMaker AI MTRL disponibiliza os construtores MultiTurnRLTrainer e MultiTurnRLEvaluator para treinar e avaliar o agente. Para avaliação de modelos base e fine-tuned, a documentação de avaliação de modelos traz orientações detalhadas.

Durante o treinamento, o sinal mais importante é a divergência: quando rollout/reward/mean sobe mas val/reward/mean permanece estável, a recompensa está sendo hackeada. A recomendação é abrir as trajetórias no MLflow e comparar o que a recompensa creditou com o que a avaliação pontuou.

O artigo relata dois experimentos reais com o SOP-Bench. No primeiro, o treinamento foi feito com todas as tarefas ao mesmo tempo, o que levou tarefas a competirem entre si e a recompensa a oscilar sem convergir:

Ao limitar o treinamento à tarefa aircraft_inspection, a recompensa de rollout saturou em torno de 3,7 (com máximo de 5,0), enquanto a recompensa de validação caiu. A investigação das trajetórias revelou dois problemas: o exemplo one-shot do SOP omitia um campo obrigatório (cross_check_response) e usava uma tag de saída diferente da esperada pelo benchmark:

Após alinhar o exemplo com os dados e remover o campo sem resposta possível, tanto a recompensa de treinamento quanto a de validação convergiram de forma saudável. O resultado foi um aumento de 13% na Taxa de Sucesso de Tarefas (TSR) e aproximadamente 16% de crescimento na acurácia por campo no benchmark externo:

O loop de iteração recomendado

O artigo propõe uma sequência clara de etapas para o treinamento:

1. Coletar dados representativos da tarefa e dividir em conjuntos de treino, validação e teste.
2. Construir o ambiente de treinamento a partir dos schemas de produção: hermético, com seed, reproduzível.
3. Configurar a avaliação externa sobre o conjunto de teste, calculada de forma independente da recompensa.
4. Estabelecer baseline rodando o modelo base e um modelo de referência frontier pela avaliação.
5. Desenhar a recompensa e validá-la em saídas reais do modelo antes de qualquer treinamento.
6. Treinar monitorando rollout/reward, taxa de conclusão e uma amostra de trajetórias.
7. Avaliar o modelo treinado com a avaliação externa e ler as trajetórias onde recompensa e avaliação divergem.
8. Ajustar recompensa, ambiente ou dados e repetir.

Para situações em que a curva trava ou colapsa, o artigo oferece um guia de diagnóstico por sintoma — por exemplo, se a recompensa de treinamento sobe mas val/reward/mean fica estável, o sinal é de reward hacking e o parser de recompensa precisa ser ajustado. A recomendação é fazer uma mudança por vez e observar as métricas por 25 a 50 passos de treinamento antes de decidir o próximo ajuste.

Conclusão

A qualidade da recompensa e da avaliação é o que determina se o treinamento produz um agente útil — muito mais do que o algoritmo ou os hiperparâmetros. O SageMaker AI MTRL cuida da maior parte do trabalho operacional de um treinamento de RL agêntico distribuído, abstraindo hardware, orquestração e o engine de treinamento. O que fica a cargo do time é o que realmente importa: ambiente preciso, design de recompensa e avaliação alinhada com o objetivo final.

Para começar, a AWS disponibiliza um notebook de exemplo para configuração do MTRL, a documentação do SageMaker AI multi-turn RL, as melhores práticas de design de recompensa e um post sobre GRPO com recompensas verificáveis. O paper e dataset do SOP-Bench são a fonte dos exemplos práticos usados ao longo do artigo.

Fonte

Best practices for multi-turn reinforcement learning in Amazon SageMaker AI (https://aws.amazon.com/blogs/machine-learning/best-practices-for-multi-turn-reinforcement-learning-in-amazon-sagemaker-ai/)