Entrenamiento multi-turn en SageMaker AI: guía práctica con MTRL

Introducción

La mayoría del ajuste fino enseña a un modelo a responder mejor a un solo prompt. Pero muchas tareas empresariales son agentivas: implican una secuencia de acciones donde cada paso depende de lo anterior y la única retroalimentación confiable es si la secuencia completa tuvo éxito. Multi‑Turn Reinforcement Learning (MTRL) en Amazon SageMaker AI aborda este tipo de problemas, permitiendo entrenar una política con recompensa a nivel de episodio en lugar de etiquetas por paso.

En esta guía práctica revisamos cómo funciona MTRL, la arquitectura que conecta SageMaker AI con un agente hospedado en Amazon Bedrock AgentCore (o infraestructura propia), el ejemplo didáctico con Wordle, resultados reales de entrenamiento y consideraciones de costo y diseño relevantes para equipos en Latinoamérica.

¿Por qué usar MTRL?

Tareas como completar un informe de gastos, derivar y resolver un ticket de soporte, o ejecutar un flujo de trabajo interno pueden requerir múltiples decisiones encadenadas. El aprendizaje supervisado estándar no ofrece señales sobre qué decisión temprana llevó al resultado final, y el RL de un solo turno no puede atribuir crédito a acciones previas dentro de la misma interacción. MTRL entrena usando la recompensa global del episodio: no necesita que ustedes especifiquen la trayectoria correcta, solo cómo puntuar el resultado.

Arquitectura y flujo de datos

El sistema tiene dos componentes principales:

• SageMaker AI: gestiona el trabajo de entrenamiento y la política (el modelo que se está ajustando).
• AgentCore (Amazon Bedrock AgentCore) o infraestructura propia: aloja el agente, el entorno y la función de recompensa.

Flujo simplificado por rollout:

1. SageMaker AI inicia un rollout con una fila del dataset (en el ejemplo, la palabra secreta).
2. MTRL invoca el runtime de AgentCore, que ejecuta un agente Strands.
3. El agente llama al modelo de política a través de un endpoint OpenAI‑compatible proporcionado por SageMaker Job Runtime.
4. El agente realiza tantas acciones (turnos) como necesite el episodio; al final, devuelve una única recompensa.
5. Group Relative Policy Optimization (GRPO) usa la distribución de recompensas entre rollouts para actualizar una adaptación LoRA en el modelo de política.

Puntos clave de diseño observados en la implementación:

• El agente y la función de recompensa viven juntos (un archivo Python desplegado en AgentCore contiene entorno, herramienta y scoring).
• El dataset contiene solo la obligación mínima (en el ejemplo, la palabra secreta); las listas de palabras para validación provienen de SCOWL y se usan solo para el scoring.
• La recompensa se diseña en función del resultado del episodio, con términos de shaping pequeños para guiar la estrategia.

Ejemplo práctico: Wordle como entorno de enseñanza

Wordle es un buen ejemplo pedagógico: reglas sencillas, dependencia de retroalimentación previa y una métrica objetiva de éxito. En cada intento, las letras se marcan como verde (posición correcta), amarillo (letra en palabra pero en otra posición) o gris (no está en la palabra). El agente tiene hasta seis intentos y recibe un solo puntaje al final del episodio.

Este ejemplo permite probar la canalización MTRL sin depender de datos etiquetados paso a paso, y es representativo de flujos empresariales donde la lógica interna no fue entrenada en modelos públicos.

Diseño de la recompensa y errores comunes

En la implementación publicada, la recompensa fue principalmente outcome‑dominant:

• Resolver la palabra paga entre 1.05 y 1.55 según la rapidez.
• Perder (no hallar la palabra) paga -0.5.
• No realizar una conjetura válida paga -1.5.
• Se añadió un término de shaping pequeño para premiar buenas estrategias dentro de cada categoría.

Los autores comparten que el diseño de recompensa es crítico: recompensas mal calibradas o shaping excesivo pueden dirigir al modelo hacia atajos no deseados. Registrar y analizar rollouts ayuda a detectar comportamientos indeseados tempranamente.

Modelos soportados y despliegue serverless

MTRL en SageMaker AI funciona en modo serverless: ustedes eligen un modelo base y el servicio aprovisiona la infraestructura de entrenamiento y muestreo. Al momento del experimento, los modelos soportados (con LoRA) incluían:

• Amazon Nova 2 Lite (regiones us‑east‑1 y us‑west‑2)
• GPT‑OSS‑20B (us‑east‑1, us‑west‑2)
• Gemma 4 31B (instruction‑tuned) (us‑west‑2)
• Qwen 3.6 27B (us‑west‑2)

En el experimento se usó GPT‑OSS‑20B por su menor precio por token y por poder controlar el esfuerzo de razonamiento por petición.

Costos observados en el experimento

Los cargos en MTRL se miden por millón de tokens en tres dimensiones:

• Prefill: tokens de entrada que el modelo procesa (acumulados a lo largo de la historia multi‑turn).
• Sample: tokens que el modelo genera (incluidos los tokens de razonamiento).
• Train: tokens usados en el backward pass para actualizar la LoRA.

En la corrida de 100 pasos y 105,486 rollouts (12 h 40 m) se consumieron:

• Prefill: 491,616,806 tokens
• Sample: 159,737,098 tokens
• Train: 643,308,554 tokens
• Total: 1,294,662,458 tokens

Con los precios listados para GPT‑OSS‑20B en us‑east‑1 en ese momento, el costo aproximado por tokens fue de $338.50, con la dimensión Train dominando el costo. Además, el runtime de AgentCore facturó separado por vCPU y memoria; en ese experimento fue aproximadamente $79 según métricas de CloudWatch.

Conclusión práctica: la longitud de la trayectoria y la verbosidad del razonamiento incrementan notablemente la factura. Para entornos empresariales en Latinoamérica, esto implica balancear costo y precisión y, si es necesario, elegir modelos y parámetros de razonamiento más económicos.

Resultados concretos

En la ejecución reportada, la adaptación LoRA sobre GPT‑OSS‑20B mejoró la tasa de resolución (“solve rate”) en palabras no vistas de:

• 61% a 84% al nivel de razonamiento por defecto.
• 15% a 69% en el modo de bajo esfuerzo con que se entrenó (importante si buscan modelos rápidos y de bajo costo).

El experimento costó en torno a $339 en tokens, aparte del cargo por runtime de AgentCore.

Monitoreo, seguridad y propiedad de datos

SageMaker AI registra métricas que se pueden consultar con DescribeJob (incluyendo conteos de tokens bajo BillableTokens) y el experimento documenta uso de MLflow para monitoreo. Importante para equipos corporativos: los datasets, el agente y el modelo entrenado permanecen en la cuenta AWS del cliente y no se usan para entrenar los modelos base.

Aplicaciones prácticas para empresas en Latinoamérica

• Automatización de flujos internos: procesos contables, control de inventario o revisión de cumplimiento que requieren decisiones encadenadas.
• Soporte y triage: bots que manejan múltiples intercambios antes de resolver o escalar un ticket.
• Herramientas internas especializadas: tareas con vocabulario o reglas propias que los grandes modelos públicos no cubren.

MTRL permite convertir un modelo pequeño y económico en una política competente para tareas específicas de su empresa, sin exponer datos sensibles fuera de su cuenta AWS.

Conclusión

MTRL en SageMaker AI ofrece un enfoque práctico para entrenar modelos en tareas secuenciales donde solo existe una señal de éxito al final del episodio. La combinación de AgentCore para ejecutar el agente y SageMaker AI para el entrenamiento con LoRA y GRPO permite iterar rápidamente, pero requiere cuidado en el diseño de la recompensa y en la gestión de costos por tokens. Para equipos en Latinoamérica que buscan automatizar flujos internos o crear modelos especializados, MTRL es una alternativa interesante para obtener rendimiento con modelos más pequeños y control sobre sus datos.

Para reproducir el experimento y ver el código de referencia, los autores publicaron el repositorio aws-samples con los scripts y configuraciones usadas en la guía.