Aprendizado por Reforço Multi-Turno no Amazon SageMaker AI: Guia Prático

Por que o ajuste fino tradicional não resolve tarefas agênticas

A maioria das técnicas de ajuste fino de modelos de linguagem foi projetada para melhorar respostas a prompts individuais. Mas tarefas agênticas funcionam de forma diferente: um agente que processa um relatório de despesas ou triagem de um chamado de suporte executa uma sequência de ações, onde cada passo depende do que aconteceu antes.

O problema é que o ajuste fino supervisionado não tem rótulos para esse tipo de trajetória, e o aprendizado por reforço de turno único não consegue creditar uma decisão tomada três passos antes do resultado final. É exatamente para esse cenário que a AWS desenvolveu o Aprendizado por Reforço Multi-Turno (MTRL, do inglês Multi-Turn Reinforcement Learning) no Amazon SageMaker AI.

Como o MTRL funciona no SageMaker AI

No MTRL, o SageMaker AI envia prompts para um agente hospedado no Amazon Bedrock AgentCore — ou em infraestrutura própria, como Amazon EKS, Amazon EC2 ou AWS Fargate, via encaminhador AWS Lambda. O agente chama o modelo de política sendo treinado, age no ambiente por quantos turnos a tarefa exigir, e reporta uma única recompensa ao final do episódio. Não é necessário especificar a trajetória correta — apenas como avaliar o resultado.

O algoritmo Otimização de Política por Gradiente Relativo de Grupo (GRPO) usa a variação de recompensas entre diferentes tentativas da mesma tarefa para atualizar um adaptador Adaptação de Baixo Posto (LoRA) no modelo de política. O MTRL é serverless: basta escolher o modelo base e o serviço provisiona toda a infraestrutura de treinamento e amostragem.

O exemplo prático: treinando com o jogo Wordle

Para demonstrar o pipeline, a AWS usou o Wordle como ambiente de treinamento. No jogo, o jogador tem seis tentativas para adivinhar uma palavra de cinco letras. Após cada tentativa, cada letra recebe uma cor: verde se estiver na posição correta, amarela se estiver na palavra mas em posição errada, e cinza se não estiver na palavra.

O Wordle é um exemplo didático ideal para o MTRL porque as regras são simples, cada tentativa depende do feedback da anterior e o resultado é avaliado de forma objetiva. Na prática, o mesmo pipeline se aplica a fluxos de trabalho internos que modelos públicos nunca foram treinados para executar.

No experimento descrito, o MTRL levou o modelo GPT-OSS-20B de uma taxa de resolução de 61% para 84% em palavras não vistas durante o treinamento, com custo de aproximadamente US$ 339 em tokens de treinamento.

Arquitetura da solução

O sistema tem duas metades. O SageMaker AI gerencia o job de treinamento e o modelo de política. O AgentCore hospeda o agente, o ambiente do Wordle e a função de recompensa. A cada rollout, o serviço MTRL invoca o runtime do AgentCore com uma linha de treinamento contendo a palavra secreta. Um agente Strands no runtime chama o modelo de política através do endpoint compatível com OpenAI do Job Runtime do SageMaker. O modelo faz até seis tentativas por meio de uma única ferramenta — guess_word — e ao final do jogo o agente retorna uma recompensa.

Modelos e regiões suportados

No momento da publicação, o MTRL suporta os seguintes modelos base, cada um treinado com um adaptador LoRA:

  • Amazon Nova 2 Lite — US East (N. Virginia) e US West (Oregon)
  • GPT-OSS-20B — US East (N. Virginia) e US West (Oregon)
  • Gemma 4 31B (instruction-tuned) — US West (Oregon)
  • Qwen 3.6 27B — US West (Oregon)

A AWS escolheu o GPT-OSS-20B para o exemplo por ter o menor preço por token entre os quatro e por permitir controle do esforço de raciocínio por requisição. Consulte a documentação do MTRL para a lista atualizada.

Custos do experimento

O MTRL cobra por milhão de tokens em três dimensões: Prefill (tokens de entrada processados pelo modelo), Sample (tokens gerados, incluindo tokens de raciocínio) e Train (tokens processados no passo de retropropagação que atualiza o adaptador LoRA). A execução de 100 passos consumiu aproximadamente 1,29 bilhão de tokens ao longo de 12 horas e 40 minutos, com o seguinte custo a preços de tabela na região us-east-1:

  • Prefill: 491,6 milhões de tokens × US$ 0,12/M = US$ 58,99
  • Sample: 159,7 milhões de tokens × US$ 0,30/M = US$ 47,92
  • Train: 643,3 milhões de tokens × US$ 0,36/M = US$ 231,59
  • Total: US$ 338,50

O runtime do AgentCore é cobrado separadamente por vCPU e memória utilizados. Nesse experimento, esse custo foi de aproximadamente US$ 79. Consulte a página de preços do Amazon SageMaker AI para valores atualizados.

Passo a passo do pipeline

Pré-requisitos

Para seguir o tutorial, são necessários Python 3.10 ou superior, uv, Node 20 ou superior e a CLI do AgentCore (npm install -g @aws/agentcore). Veja como começar com o Amazon Bedrock AgentCore para instalação e permissões de Gerenciamento de Identidade e Acesso (IAM). Também são necessários botocore 1.43.37 ou superior, dois papéis IAM, e um bucket do Amazon S3 na mesma região do agente. Veja os pré-requisitos do MTRL para detalhes completos.

Etapa 1: Criar os papéis IAM

O MTRL usa dois papéis. O papel do runtime do agente é criado automaticamente pelo comando agentcore deploy e precisa da política gerenciada AmazonSageMakerJobRuntimeAccess. O papel do job de treinamento deve ser criado manualmente e precisa confiar em job.sagemaker.amazonaws.com com as ações sts:AssumeRole e sts:TagSession, além da política AmazonSageMakerJobFullAccess.

{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Principal": {"Service": "job.sagemaker.amazonaws.com"},
    "Action": ["sts:AssumeRole", "sts:TagSession"]
  }]
}

Atenção: o principal correto é job.sagemaker.amazonaws.com, não sagemaker.amazonaws.com. Com o principal genérico, o CreateJob falha com “Could not assume role”.

Etapa 2: Implantar o agente no Amazon Bedrock AgentCore

Clone o repositório de exemplos, configure conta e região em Wordle/agentcore/aws-targets.json e execute:

git clone https://github.com/aws-samples/sample-multi-turn-rl-wordle
cd sample-multi-turn-rl-wordle/Wordle
agentcore validate
agentcore deploy

A CLI compila as dependências do agente para a plataforma Linux ARM64 do runtime, cria o papel de execução com a política AmazonSageMakerJobRuntimeAccess e implanta via AWS CloudFormation. Anote o ARN do runtime na saída para usar no job de treinamento.

O agente é um único arquivo, main.py. Simplificado, seu ponto de entrada tem a seguinte estrutura — para uma tarefa própria, basta alterar o parsing da tarefa, as ferramentas e a recompensa:

app = BedrockAgentCoreApp()

@app.entrypoint # outermost: starts the HTTP server AgentCore invokes
@sagemaker_rft_handler # reports the returned reward to SageMaker AI
def handle_rollout(payload):
    prompt, answer = parse_task(payload["prompt"])
    game = WordleGame(answer)

    @tool
    def guess_word(guess: str) -> str:
        """Submit a 5-letter Wordle guess and get feedback."""
        return game.guess(guess)

    agent = Agent(model=build_model(), system_prompt=SYSTEM_PROMPT, tools=[guess_word])
    agent(prompt)
    return {"reward": game.final_reward()}

Importante: o decorator sagemaker_rft_handler não inicia um servidor HTTP por si só. É necessário envolvê-lo com BedrockAgentCoreApp e @app.entrypoint, com o decorator entrypoint mais externo. Sem isso, cada invocação expira após 30 segundos.

Etapa 3: Verificar o runtime antes de gastar GPU

Antes de iniciar o treinamento, vale confirmar que toda a cadeia funciona com um payload de teste. O repositório inclui um em Wordle/payload.json. Substitua o ID de conta fictício pelo seu e execute:

ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
sed "s/111122223333/$ACCOUNT/" payload.json > /tmp/payload.json
aws bedrock-agentcore invoke-agent-runtime \
  --agent-runtime-arn  \
  --runtime-session-id smoke-test-$(date +%s)000000000000000 \
  --payload fileb:///tmp/payload.json /dev/stdout

Nos logs do Amazon CloudWatch do runtime, o agente deve tentar autenticar no endpoint do Job Runtime e receber um erro 400 com “Could not find job” — o que confirma que parsing, autenticação e reporte de erros funcionam corretamente.

Etapa 4: Preparar o dataset

O dataset do MTRL não tem respostas-alvo, transcrições nem rótulos. Cada linha é uma tarefa para o agente tentar. As três primeiras linhas do arquivo de treinamento têm esta forma:

{"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"riser\", \"id\": \"wordle_train_0000\"}"}
{"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"swoon\", \"id\": \"wordle_train_0001\"}"}
{"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"dirty\", \"id\": \"wordle_train_0002\"}"}

O campo answer é usado apenas pelo ambiente e pela função de recompensa — o modelo de política nunca o vê. O script make_dataset.py do repositório gera 640 palavras de treinamento e 128 de validação sem sobreposição, a partir das listas de palavras de código aberto SCOWL. Após gerar, faça upload para o Amazon S3:

uv run python make_dataset.py
aws s3 cp training-data.jsonl s3:///wordle-mtrl/train/
aws s3 cp validation-data.jsonl s3:///wordle-mtrl/validation/

Etapa 5: Lançar o job de treinamento

O launcher usa a classe MultiTurnRLTrainer do SDK sagemaker-train. Um app MLflow é obrigatório — se não tiver um, crie com aws sagemaker create-mlflow-app:

uv run python run_mtrl_training.py \
  --model openai-reasoning-gpt-oss-20b \
  --agent-runtime-arn  \
  --role-arn arn:aws:iam:::role/ \
  --train-dataset s3:///wordle-mtrl/train/training-data.jsonl \
  --val-dataset s3:///wordle-mtrl/validation/validation-data.jsonl \
  --s3-output-path s3:///wordle-mtrl/output/ \
  --mlflow-app-arn arn:aws:sagemaker:::mlflow-app/

O launcher sobrescreve três hiperparâmetros: sampling_max_tokens para 8192 (o padrão de 4096 fazia o modelo esgotar tokens antes da primeira chamada de ferramenta), max_epochs para 20 e max_steps para 100.

Atenção: não reimplante o agente enquanto um job estiver em execução. Reimplantar reinicia o runtime, os rollouts em andamento não recebem requisições de amostragem e o serviço falha o job inteiro.

Etapa 6: Monitorar o treinamento no MLflow

O serviço registra métricas de recompensa, tokens e turnos a cada passo, além de executar validação a cada val_every passos. A métrica principal a observar é val/reward/pass_at_1 — a fração de palavras de validação resolvidas. No experimento, a recompensa de treinamento subiu para cerca de +0,57 nos primeiros 15 passos e se estabilizou. A taxa de resolução na validação saltou de 2% para 31% nos primeiros 10 passos, depois oscilou entre 35% e 48% — variação esperada com apenas 128 palavras.

Etapa 7: Avaliar o modelo treinado nas configurações de serviço

Um job de avaliação (com JobCategory definido como AgentRFTEvaluation) executa os prompts de avaliação pelo mesmo runtime do AgentCore e pela mesma função de recompensa, tanto no modelo base quanto no pacote treinado, registrando eval/reward/pass_at_1 no MLflow. Diferentemente da validação durante o treinamento, aqui é possível definir a temperatura.

uv run python run_mtrl_eval.py --runs base:1.0 trained:1.0 \
  --model-package-arn  \
  --agent-runtime-arn  \
  --role-arn arn:aws:iam:::role/ \
  --dataset s3:///wordle-mtrl/validation/validation-data.jsonl \
  --s3-output-path s3:///wordle-mtrl/eval/ \
  --mlflow-app-arn arn:aws:sagemaker:::mlflow-app/

Os resultados mostram que o adaptador adiciona entre 14 e 54 pontos percentuais em todas as configurações testadas. A curva de validação do treinamento reportou 37% no passo 100 (temperatura 0), mas nas configurações de serviço (esforço médio, temperatura 1.0) o modelo treinado resolve 84% das palavras, contra 61% do modelo base. A referência de hiperparâmetros recomenda temperatura 0 para avaliação durante o treinamento por tornar as curvas determinísticas e comparáveis — mas o ideal é sempre avaliar nas mesmas configurações que serão usadas em produção.

Dica: avalie nas configurações que você vai usar em produção. A curva de treinamento diz que o adaptador termina em 37%. Nas configurações de serviço, ele resolve 84%.

Etapa 8: Limpeza dos recursos

# 1. Agent runtime and its execution role. remove all resets the project config;
# the next deploy sees the empty state and tears down the deployed resources.
cd sample-multi-turn-rl-wordle/Wordle
agentcore remove all
agentcore deploy

# 2. MLflow app
aws sagemaker delete-mlflow-app --arn arn:aws:sagemaker:::mlflow-app/

# 3. Datasets, job output, and evaluation output
aws s3 rm s3:///wordle-mtrl/ --recursive

Como projetar a função de recompensa: lições aprendidas

A função de recompensa usada no experimento passou por uma revisão importante. A primeira versão somava a pontuação por tentativa de um projeto de referência (wordle-lora-rl) ao longo de todo o jogo. Isso gerou duas inversões problemáticas: um jogo perdido pagava recompensa positiva (cerca de +0,49 em seis turnos de bônus), e resolver na primeira tentativa valia menos do que resolver na terceira.

A solução foi separar um termo de resultado dominante de um pequeno termo de modelagem limitado:

reward = outcome + w * clip(shaping, -1, 1), onde w = 0.04 se resolvido, caso contrário 0.15

O termo de resultado é definido exclusivamente pelo desfecho do jogo:

  • Resolvido na tentativa n: 1,05 + 0,5 × (6 – n) / 5 — de 1,55 na primeira até 1,05 na sexta
  • Seis tentativas sem resolver: -0,5
  • Nunca fez uma tentativa válida: -1,5

O termo de modelagem soma penalidades por contradizer letras conhecidas, por repetir tentativas e por palavras fora do dicionário, além de bônus por tentar letras novas e reduzir o espaço de possibilidades. Esse valor é limitado a [-1, 1] e escalado pelo peso w, garantindo que nunca mova um jogo para outra faixa de resultado.

Como a documentação do SageMaker AI orienta: comece com uma recompensa baseada em resultado, adicione modelagem com parcimônia e verifique a recompensa em trajetórias reais.

Problemas comuns durante o treinamento

A AWS documentou os principais erros encontrados durante o experimento, em ordem de ocorrência:

  • Tokens de autenticação sob concorrência: criar um token por chamada esgota o provedor de credenciais em 32 rollouts simultâneos. A solução é uma sessão auto-renovável com cache curto e re-autenticação em erros 403.
  • Rollout que nunca chega ao modelo falha o job inteiro: quando todas as chamadas ao modelo falham por autenticação, o serviço encerra o job com “No sampling requests were received”. Retornar um status de erro também causa falha — sempre retorne uma recompensa.
  • Chamadas de ferramenta derivando para texto puro: com taxa de aprendizado muito alta, o modelo passou a escrever chamadas de ferramenta como texto, todos os rollouts receberam a recompensa mínima e o gradiente desapareceu. A solução foi recuperar tentativas de texto bruto com uma pequena penalidade de formatação.
  • Platô de treinamento: a recompensa de treinamento parou de subir após cerca de 15 passos e a validação após 30 passos. Um currículo que inicializa jogos com tentativas anteriores, ou mais rollouts por palavra, seriam os próximos passos a testar.

Conclusão

O MTRL no SageMaker AI treina modelos com base no resultado de um episódio completo — o que o torna adequado para tarefas agênticas onde a trajetória correta é desconhecida e apenas o resultado pode ser avaliado. O exemplo usou uma recompensa única por episódio, mas quando passos intermediários também podem ser avaliados, o MTRL aceita uma recompensa por turno via API UpdateReward.

Um modelo de código aberto com 20 bilhões de parâmetros atingiu 84% de taxa de resolução nas configurações de serviço, partindo de 61%, com cerca de US$ 339 em tokens de treinamento e US$ 79 em uso do runtime do AgentCore. O mesmo pipeline se aplica a procedimentos internos, ferramentas e fluxos de trabalho que modelos públicos nunca foram treinados para executar.

Para começar, clone o repositório de exemplos, siga os passos descritos e execute o smoke test antes de comprometer tempo de GPU. Para as APIs subjacentes, consulte a documentação do MTRL no SageMaker AI.

Fonte

Multi-Turn Reinforcement Learning on Amazon SageMaker AI: A hands-on guide (https://aws.amazon.com/blogs/machine-learning/multi-turn-reinforcement-learning-on-amazon-sagemaker-ai-a-hands-on-guide/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *