Afina un agente de búsqueda multi-turno con MTRL en Amazon SageMaker AI

Los agentes de búsqueda que usan LLMs pueden decidir qué buscar y cómo hacerlo a lo largo de varias rondas de interacción. Amazon SageMaker AI ofrece MTRL para entrenar esos agentes enfocándose en la trayectoria completa y optimizando la calidad de recuperación sin depender solo de modelos frontera.

Por Redaccion TD
Afina un agente de búsqueda multi-turno con MTRL en Amazon SageMaker AI

Introducción

Los agentes de búsqueda potenciados por grandes modelos de lenguaje (LLMs) están cambiando la forma en que las empresas recuperan información: en lugar de exigir consultas perfectas al usuario, el agente decide de forma autónoma qué buscar, qué estrategia de recuperación usar y cuándo dar por concluida la búsqueda. Esto ocurre a lo largo de varias rondas de interacción (multi-turn), donde cada paso depende del contexto acumulado.

Sin embargo, lograr un comportamiento consistente y eficiente en múltiples turnos no es trivial. Los modelos base no conocen sus herramientas ni el entorno específico de una organización. Hacer prompting a un modelo pequeño rara vez produce comportamiento multi-turno confiable; usar un modelo frontera puede resolverlo, pero a costa de mayor latencia y costo. La alternativa es el fine-tuning: enseñar al modelo pequeño cómo operar en su entorno para obtener velocidad y economía con una conducta confiable.

En este artículo explicamos cómo Amazon SageMaker AI soporta el entrenamiento por aprendizaje por refuerzo multi-turno (MTRL) para afinar agentes de búsqueda, qué ventajas ofrece y cuáles son las consideraciones prácticas para implementarlo en entornos empresariales.

¿Qué es Amazon SageMaker AI MTRL?

MTRL (Multi-Turn Reinforcement Learning) en Amazon SageMaker AI es una capacidad para afinar LLMs mediante aprendizaje por refuerzo en escenarios de interacción multi-turno. La idea central es modelar la tarea del agente como una secuencia de decisiones y optimizar la política en función de trayectorias completas, no de respuestas aisladas.

Principales características destacadas:

  • Interfaz modular agente-entorno: integración de bajo código para definir recompensas personalizadas, bucles de herramientas y la forma de la conversación multi-turno.
  • Ejecución serverless: entrenamiento y despliegue sin administrar clústeres GPU, con facturación por token y escala automática para producción.
  • Rollouts asíncronos y recolección de trayectorias: generación de muestras y actualizaciones de gradiente en paralelo, manteniendo el entrenamiento rápido y controlando la obsolescencia off-policy.
  • Biblioteca nativa de algoritmos: soporte para PPO, variantes con importance sampling y estimadores de ventaja agrupados (por ejemplo, GRPO y variantes pass@k).
  • Entrenamiento resumible: posibilidad de dividir corridas largas en múltiples jobs para superar límites de tiempo por job.
  • Observabilidad de trayectorias y recompensas: inspección detallada de cada turno y evolución durante el entrenamiento, integrada con MLflow.
  • Jobs de evaluación: métricas de recompensa, pass@k y otras antes de desplegar a endpoints de SageMaker AI o Amazon Bedrock.

Estas capacidades hacen que MTRL sea particularmente adecuado para agentes de búsqueda: hay una señal de recompensa clara (calidad de recuperación), un bucle de múltiples turnos (consultas y resultados) y un entorno bien definido (herramientas de búsqueda).

Por qué optimizar por trayectorias completas importa

En enfoques tradicionales:

  • El fine-tuning supervisado (SFT) requiere demostraciones humanas de trayectorias ideales, costosas y a menudo inexistentes para escenarios corporativos específicos.
  • El RL por respuesta única (p. ej. RL con recompensas verificables) evalúa y optimiza respuestas aisladas, perdiendo las dependencias entre turnos.

Un agente de búsqueda toma decisiones interdependientes: la elección del tipo de búsqueda en un turno afecta las posibilidades en los siguientes. MTRL optimiza la política sobre la secuencia completa y usa una señal de recompensa que puede reflejar únicamente si el resultado final fue satisfactorio. Esto enseña al agente a ejecutar estrategias eficaces durante toda la interacción.

Entorno, herramientas y modelo usados en el ejemplo

En la implementación descrita por AWS se usó lo siguiente:

  • Modelo base: Qwen3.6-27B (disponible en la región US West (Oregon), us-west-2).
  • Herramientas de búsqueda expuestas al agente endpoint: una búsqueda léxica (BM25) y una búsqueda vectorial por embeddings.
    • BM25 es adecuado para queries con términos específicos o identificadores exactos.
    • La búsqueda vectorial es recomendable para consultas semánticas o conceptuales, usando similitud entre embeddings.
  • Límite en el número de turnos: para evitar respuestas excesivamente largas y fomentar búsquedas eficientes.
  • Datos y artefactos: datasets de entrenamiento y validación almacenados en Amazon S3 en el formato requerido por MTRL, y un endpoint que permite llamadas a las herramientas durante los rollouts.

Configuración clave de entrenamiento

Tres componentes fundamentales en el setup son los datasets, la función de recompensa y la configuración del job MTRL.

Datasets

En el ejemplo de AWS se utilizaron conjuntos públicos y sintéticos para cubrir distintos retos de recuperación:

  • FRAMES (Train) — un benchmark de QA multihop que requiere síntesis a partir de múltiples artículos de Wikipedia (disponible en HuggingFace).
  • BRIGHT (Train) — un conjunto enfocado en recuperación que exige razonamiento sobre 12 dominios, donde la relevancia va más allá de la coincidencia de palabras (disponible en HuggingFace).
  • Enterprise RAG (Train) — un benchmark con documentos sintéticos empresariales y preguntas para entrenar y evaluar sistemas RAG en contextos internos (referencia en GitHub; incluye un amplio corpus y un conjunto de 500 preguntas para evaluación).

Función de recompensa

La clave en MTRL es que la señal de recompensa puede ser global, evaluando el resultado final de la trayectoria. En lugar de asignar recompensa por cada salida, se puede definir una métrica de éxito que refleje si la respuesta final cumple criterios de relevancia y calidad.

Esto simplifica la ingeniería de recompensas: basta con una señal que capture el objetivo del agente (por ejemplo, si la respuesta final satisface la consulta del usuario) y, opcionalmente, componentes auxiliares para eficiencia (p. ej., penalizaciones por turnos innecesarios).

Configuración del job MTRL

SageMaker AI permite elegir el algoritmo (PPO, CISPO, IS, etc.), parámetros de rollout asíncrono, límites de off-policy y mecanismos de evaluación continuos. El entrenamiento puede dividirse en múltiples jobs y las trayectorias quedan registradas para inspección y auditoría.

Buenas prácticas y consideraciones para empresas en América Latina

  • Definan recompensas alineadas con el objetivo del negocio: en entornos corporativos la señal puede provenir de métricas de utilidad humana, de verificaciones automáticas o sistemas de clasificación internos.
  • Seleccionen herramientas complementarias: combinar BM25 para identificadores y vector search para consultas semánticas mejora la cobertura.
  • Controlen la cantidad de turnos: un tope razonable promueve eficiencia y reduce costos de inferencia.
  • Privacidad y residencia de datos: almacenen datos sensibles en regiones y cuentas que cumplan con las políticas de la organización; en muchos proyectos latinoamericanos es clave revisar requisitos regulatorios.
  • Evaluación y observabilidad: aprovechen las capacidades de SageMaker AI para inspeccionar trayectorias y entender decisiones del agente, lo que facilita auditoría y mejora iterativa.
  • Balance costo-rendimiento: afinar un modelo pequeño especializado puede ofrecer la latencia y el costo de modelos compactos con la fiabilidad de modelos más grandes.

Conclusión

Multi-turn RL en Amazon SageMaker AI ofrece una ruta práctica para convertir modelos pequeños en agentes de búsqueda robustos dentro de entornos empresariales. Al optimizar políticas sobre trayectorias completas y aprovechar una plataforma que facilita la integración de herramientas, la ejecución serverless y la observabilidad, las organizaciones pueden mejorar la calidad de recuperación sin depender exclusivamente de modelos frontera. Para equipos en América Latina, esto representa una alternativa atractiva para implementar agentes de búsqueda eficientes, controlables y alineados con requisitos locales de datos y operación.

Fuente original: AWS ML Blog