AWS Machine LearningPublicación de la fuente

Aprendizaje por refuerzo multipaso en Amazon SageMaker AI: Una guía práctica

La mayoría del fine-tuning enseña a un modelo a producir una mejor respuesta a un único prompt. Las tareas agénticas son diferentes.

La mayoría del fine-tuning enseña a un modelo a producir una mejor respuesta a un único prompt. Las tareas agénticas son diferentes. Un agente que presenta un informe de gastos o clasifica un ticket de soporte realiza una secuencia de acciones, y cada paso depende de lo que ocurrió antes. Rara vez existe una única secuencia correcta que imitar, y la única retroalimentación fiable es si toda la secuencia tuvo éxito. El fine-tuning supervisado no tiene etiquetas para eso, y el aprendizaje por refuerzo de un solo turno no puede dar crédito a una decisión tomada tres pasos antes del resultado.

Multi-Turn Reinforcement Learning (MTRL) en Amazon SageMaker AI entrena modelos exactamente con esa señal. SageMaker AI envía prompts a un agente que usted aloja en Amazon Bedrock AgentCore, o en su propia infraestructura, como Amazon Elastic Kubernetes Service (Amazon EKS), Amazon Elastic Compute Cloud (Amazon EC2) o AWS Fargate, a través de un Reenviador de AWS Lambda. El ejemplo aquí usa AgentCore. El agente llama al modelo de política que se está entrenando, actúa en tu entorno durante tantos turnos como necesite la tarea, e informa una recompensa cuando el episodio termina. No especificas la trayectoria correcta, solo cómo puntuar el resultado.

Esta entrada utiliza Wordletranslate, ", un juego en el que el jugador tiene seis intentos para adivinar una palabra oculta de cinco letras, como el ejemplo. Después de cada intento, cada letra se vuelve verde si está en la posición correcta, amarilla si está en la palabra pero en una posición diferente, y gris si no está en la palabra. El juego es adecuado para MTRL porque las reglas son pocas, cada intento depende de la retroalimentación del anterior, y el resultado se evalúa objetivamente. Wordle es un ejemplo didáctico, y los modelos frontera ya lo juegan bien. En la práctica, se usaría MTRL para hacer que un modelo pequeño y de bajo costo sea bueno en una tarea de varios pasos específica de tu negocio, como un flujo de trabajo interno en el que los modelos públicos nunca fueron entrenados. El mismo pipeline se aplica a ese caso."

En nuestra ejecución, MTRL llevó a gpt-oss-20b de una tasa de resolución del 61 a un 84 por ciento en palabras reservadas con su nivel de esfuerzo de razonamiento predeterminado, y del 15 al 69 por ciento con el esfuerzo bajo con el que fue entrenado, por unos $339 en tokens de entrenamiento. En esta publicación, repasamos el pipeline de extremo a extremo: desplegar el agente y la recompensa en AgentCore, preparar el conjunto de datos, entrenar, monitorear en MLflow y evaluar en la configuración de servicio. También compartimos los errores de diseño de recompensas y los fallos de entrenamiento que encontramos. Sus conjuntos de datos, agente y modelo entrenado permanecen en su cuenta de AWS y no se usan para entrenar los modelos base. Todo el código está en el repositorio adjunto aws-samples.

Descripción general de la solución

El sistema tiene dos mitades. SageMaker AI posee el trabajo de entrenamiento y el modelo de política. AgentCore aloja el agente, el entorno de Wordle y la función de recompensa. El siguiente diagrama muestra cómo se conectan.

Figura 1: MTRL en SageMaker AI, donde el trabajo de entrenamiento invoca al agente alojado en AgentCore una vez por rollout para muestrear el modelo de política, jugar la partida y reportar una recompensa

Para cada rollout, el servicio MTRL invoca el runtime de AgentCore, una capacidad de Amazon Bedrock AgentCore, con una fila de entrenamiento, que contiene una palabra secreta. Un agente de Strands en el runtime llama al modelo de política a través del endpoint compatible con OpenAI del SageMaker Job Runtime. El modelo hace hasta seis intentos mediante una herramienta, guess_word, y después de cada intento el entorno devuelve retroalimentación por letra y un resumen de lo que se sabe. Cuando termina la partida, el agente devuelve una recompensa. Group Relative Policy Optimization (GRPO) usa la dispersión de recompensas entre varios rollouts de la misma palabra para actualizar un adaptador de Low-Rank Adaptation (LoRA) en el modelo de política.

Tres puntos de diseño importan para lo que sigue:

  • El agente y la recompensa viven en un solo lugar. Un único archivo Python desplegado en AgentCore contiene el entorno, la herramienta y la función de puntuación.
  • El conjunto de datos contiene solo la palabra secreta. Las listas de palabras incluidas con el agente se construyen a partir de las listas de palabras de código abierto SCOWL y se usan únicamente para la puntuación.
  • La recompensa está dominada por el resultado. Resolver paga entre 1.05 y 1.55 según la rapidez, perder paga -0.5, y no hacer nunca un intento válido paga -1.5. Un pequeño término de modelado recompensa una buena estrategia dentro de cada nivel.

Modelos compatibles y Regiones de AWS

MTRL es serverless: usted elige un modelo base y el servicio aprovisiona la infraestructura de entrenamiento y muestreo. En el momento de escribir este artículo, admite los siguientes modelos base, cada uno entrenado con un adaptador LoRA:

Modelo base Regiones de AWS
Amazon Nova 2 Lite US East (N. Virginia) (us-east-1), US West (Oregon) (us-west-2)
GPT-OSS-20B US East (N. Virginia) (us-east-1), US West (Oregon) (us-west-2)
Gemma 4 31B (ajustado para instrucciones) US West (Oregon) (us-west-2)
Qwen 3.6 27B US West (Oregon) (us-west-2)

Elegimos GPT-OSS-20B porque tiene el precio más bajo por token de los cuatro y su esfuerzo de razonamiento es controlable por solicitud. Consulte la documentación de MTRL para ver la lista actual.

Precios

No elige instancias ni dimensiona un clúster. Paga por millón de tokens en tres dimensiones, con precios según el modelo base:

  • Prefill: tokens de entrada que procesa el modelo de política. Estos se acumulan en una tarea de varios turnos, porque cada turno vuelve a leer el historial.
  • Sample: tokens que genera el modelo de política, incluidos los tokens de razonamiento.
  • Train: tokens procesados en la pasada hacia atrás que actualiza el adaptador LoRA.

La ejecución de 100 pasos consumió 491.6 millones de tokens de prefill, 159.7 millones de sample y 643.3 millones de train en 105,486 rollouts en 12 horas 40 minutos. Con los GPT-OSS-20B precios de lista en la Región EE. UU. Este (N. Virginia) (us-east-1) en el momento de escribir esto, eso equivale a unos $339:

Dimensión Tokens Precio por 1M Costo
Prefill 491,616,806 $0.12 $58.99
Sample 159,737,098 $0.30 $47.92
Train 643,308,554 $0.36 $231.59
Total 1,294,662,458 $338.50

La dimensión train domina y escala con la longitud de la trayectoria, por lo que los juegos más cortos y un razonamiento menos verboso reducen la factura. El runtime de AgentCore factura por separado la vCPU y la memoria que usa el agente. Para esta ejecución, fueron unos $79, según las métricas de uso de Amazon CloudWatch del runtime a precios de lista de us-east-1. DescribeJob informa los recuentos de tokens en BillableTokenUsage. Consulte precios de Amazon SageMaker AI para ver las tarifas actuales.

Requisitos previos

Para seguir esta guía, necesita lo siguiente. Consulte también los requisitos previos de MTRL.

  • Python 3.10 o posterior, uv, Node 20 o posterior y la CLI de AgentCore (npm install -g @aws/agentcore). Consulte Comience con Amazon Bedrock AgentCore para la instalación y los permisos de IAM que necesita la CLI.
  • botocore 1.43.37 o posterior, en la máquina que lanza el trabajo. Esta es la primera versión que incluye la AgentRFT CreateJob operación.
  • Dos roles de AWS Identity and Access Management (IAM), descritos en el Paso 1.
  • Un bucket de Amazon Simple Storage Service (Amazon S3) en la misma Región que el agente. El trabajo de entrenamiento también debe ejecutarse en esa Región, porque CreateJob rechaza un ARN de runtime de otra Región.

Paso 1: Crear los roles de IAM

MTRL usa dos roles, y uno de ellos tiene una política de confianza que la documentación general de SageMaker AI no menciona.

Rol de runtime del agente. agentcore deploy crea este rol por usted. Confía en bedrock-agentcore.amazonaws.com y necesita la AmazonSageMakerJobRuntimeAccess política administrada, que permite al agente llamar al modelo de política y reportar recompensas. El stack de CDK en el repositorio de ejemplo adjunta esta política automáticamente.

Rol del trabajo de entrenamiento. Usted crea este. Debe confiar en job.sagemaker.amazonaws.com con ambos sts:AssumeRole y sts:TagSession, y necesita la AmazonSageMakerJobFullAccess política administrada.

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

Nota: El principal de confianza es job.sagemaker.amazonaws.com, no sagemaker.amazonaws.com. Con el principal genérico, CreateJob falla con “Could not assume role,” y el mensaje nombra el principal correcto solo después de varios reintentos.

Paso 2: Desplegar el agente en Amazon Bedrock AgentCore

Clone el repositorio de ejemplo, configure su cuenta y Región en Wordle/agentcore/aws-targets.jsony despliegue:

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

La CLI compila las dependencias del agente para la plataforma Linux ARM64 del runtime, crea el rol de ejecución con AmazonSageMakerJobRuntimeAccess adjuntado, y despliega mediante AWS CloudFormation. La recomendamos en lugar de cargar un paquete de código en la consola, donde es común equivocarse con los aproximadamente 700 MB de dependencias compiladas o con la política adjuntada. Anote el ARN del runtime en la salida para el trabajo de entrenamiento.

Figura 2: El runtime del agente Wordle en la consola de Amazon Bedrock AgentCore con estado Ready

El runtime lee dos variables de entorno desde agentcore.json. REASONING_EFFORT=low establece reasoning_effort en cada solicitud al modelo de política. En nuestros trabajos de evaluación, el esfuerzo bajo utilizó entre un 63 y un 83 por ciento menos tokens de muestra por partida que el medio. REASONING_PROMPT_HINT=true también añade “Reasoning: low” al prompt del sistema, que es la forma en que gpt-oss lee esta configuración en su formato de chat Harmony. Ninguna de las dos variables da al modelo ninguna pista sobre la palabra secreta.

Figura 3: Variables de entorno del runtime que muestran REASONING_EFFORT configurado en low

El agente es un solo archivo, main.py. Simplificado, su punto de entrada se ve así. Para tu propia tarea, cambias el análisis de la tarea, las herramientas y la 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: El decorador sagemaker_rft_handler no inicia un servidor HTTP por sí mismo. Envuélvelo con BedrockAgentCoreApp y @app.entrypoint, con el decorador de entrada como el más externo. Sin esto, cada invocación agota el tiempo de espera tras 30 segundos.

Paso 3: Verificar el runtime

Antes de gastar tiempo de GPU, confirma que toda la cadena funciona con un payload con la forma del que envía el servicio de entrenamiento. El repositorio incluye uno en Wordle/payload.json:

{
  "prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"beach\", \"id\": \"smoke-test\"}",
  "metadata": {
    "jobArn": "arn:aws:sagemaker:us-east-1:111122223333:job/AgentRFT/smoke-test",
    "rolloutId": "r1"
  }
}

Se espera que el ARN de trabajo falso falle, y la forma en que falla es la prueba. Pon primero tu propio ID de cuenta: para un ARN de trabajo en otra cuenta, el servicio devuelve un 403 antes de verificar si el trabajo existe, lo que parece una configuración de IAM rota. Desde el directorio Wordle, ejecuta:

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_ARN> \
  --runtime-session-id smoke-test-$(date +%s)000000000000000 \
  --payload fileb:///tmp/payload.json /dev/stdout

En los registros de Amazon CloudWatch del runtime, ves al agente analizar el prompt, autenticarse en el endpoint de Job Runtime y recibir un error 400 con el mensaje “Could not find job.”. Ese fallo confirma que el análisis, la autenticación y el reporte de errores funcionan. El comando en sí imprime {“reward”: -1.5}, porque el agente nunca llegó al modelo de política. Un 403 con el mensaje “not authorized to access this resource” significa que el ARN del trabajo aún nombra otra cuenta. Un tiempo de espera de 30 segundos significa en cambio que el servidor HTTP no está iniciando (ver la nota Importante en el Paso 2).

Paso 4: Preparar el conjunto de datos

El conjunto de datos MTRL no tiene respuestas objetivo, transcripciones ni etiquetas. Cada fila es una tarea para que el agente la intente. Aquí están las primeras tres filas del archivo de entrenamiento:

{"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\"}"}

El servicio lee la columna prompt (o la primera columna) y pasa la cadena a tu agente textualmente, sin analizarla. Hicimos que la cadena fuera un pequeño objeto JSON con tres campos:

  • prompt: la instrucción que el agente muestra al modelo de política, idéntica en cada fila.
  • answer: la palabra secreta. Solo el entorno y la función de recompensa la usan. El modelo de política nunca la ve.
  • id: un ID de fila para rastrear un rollout en los registros de vuelta a su fila de origen.

El repositorio de ejemplo make_dataset.py extrae 640 palabras de entrenamiento (cinco lotes completos de 128) y 128 palabras de validación de la lista de respuestas sin superposición, por lo que la tasa de resolución de validación mide la generalización. El servicio también acepta Apache Parquet, arreglos JSON y CSV. Sube ambos archivos a Amazon S3:

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

Como el servicio nunca inspecciona la cadena prompt , una tarea de producción puede empaquetar en ella una lista completa de mensajes, una configuración de herramientas o una especificación de recompensa. Consulta la documentación de MTRL para ver ejemplos, incluido cómo proteger contenido sensible en la cadena.

Paso 5: Lanzar el trabajo de entrenamiento

El launcher usa la clase MultiTurnRLTrainer del SDK sagemaker-train . Apúntalo al ARN del runtime del Paso 2, al rol del Paso 1, a los conjuntos de datos del Paso 4 y a una app de MLflow para las métricas. La app de MLflow es obligatoria, así que si no tienes una, créala primero con aws sagemaker create-mlflow-app:

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

El lanzador sobrescribe tres hiperparámetros y deja todos los demás en el valor predeterminado de gpt-oss-20b. Establece sampling_max_tokens en 8192, el máximo del servicio, porque con el valor predeterminado de 4096 gpt-oss a menudo se quedaba sin tokens mientras razonaba, antes de su primera llamada a una herramienta. Establece max_epochs en 20 y max_steps en 100: 640 filas con el valor predeterminado global_batch_size de 128 son 5 pasos por época, y con el valor predeterminado group_size de 8, cada paso juega 1,024 partidas. El trabajo pasa por Starting, Downloading y Pending antes de llegar a Training. Una vez que está entrenando, tu runtime de AgentCore comienza a recibir invocaciones de rollout.

Importante: No redespliegues el agente mientras se está ejecutando un trabajo. Redesplegar reinicia el runtime, los rollouts en curso no ven solicitudes de muestreo y el servicio marca como fallido todo el trabajo.

Paso 6: Monitorear el entrenamiento en MLflow

El servicio registra estadísticas de recompensa, tokens y turnos en cada paso y ejecuta la validación cada val_every pasos. La métrica a observar es val/reward/pass_at_1, la fracción de palabras reservadas resueltas. Cada acierto puntúa al menos 1.01, por lo que supera el valor de success_threshold de 1.0 del servicio y pass@1 lo cuenta. El siguiente gráfico está construido a partir del historial de MLflow de esta ejecución:

Figura 4: Curvas de entrenamiento para la ejecución de 100 pasos, con la recompensa media por partida a la izquierda y la tasa de aciertos de validación a la derecha

La recompensa de entrenamiento sube a aproximadamente +0.57 en 15 pasos y se mantiene ahí. La recompensa de validación se mantiene más baja, principalmente debido a la temperatura de 1.0 frente a 0 en lugar de sobreajuste (ver Paso 7). La tasa de aciertos de validación salta de 2 a 31 por ciento en los primeros 10 pasos, y luego se mueve entre 35 y 48 por ciento. Con 128 palabras, esas oscilaciones son mayormente ruido. No probamos si más pasos ayudan en las configuraciones de servicio.

Los trazos de OpenTelemetry del runtime en CloudWatch registran cada turno del modelo y cada llamada a herramientas. Aquí hay una partida que el modelo entrenado jugó durante un trabajo de evaluación (esfuerzo medio, temperatura 1.0), con el razonamiento recortado. La palabra secreta era PASTE:

[model] Reasoning: [...] "CRANE" has 5 distinct letters: C,R,A,N,E. [...]
We'll guess "CRANE".
Tool call: guess_word(guess="CRANE")
[tool] Guess 1/6: CRANE -> C=gray, R=gray, A=yellow, N=gray, E=green
Current Knowledge:
- Correct Position (Green): _ _ _ _ E
- In Word, Wrong Position (Yellow): 'A' (at least 1)
- Not in Word (Gray): C, N, R
- Guesses Remaining: 5
[model] Reasoning: [...] "ALIVE": ends with E, includes A. Letters not C,N,R.
Tool call: guess_word(guess="ALIVE")
[tool] Guess 2/6: ALIVE -> A=yellow, L=gray, I=gray, V=gray, E=green
[model] Reasoning: [...] So positions not 1 or 3 for A. Positions left: 2,4
for A. [...] We might guess "HASTE"? [...]
Tool call: guess_word(guess="HASTE")
[tool] Guess 3/6: HASTE -> H=gray, A=green, S=green, T=green, E=green
[model] Reasoning: We have the word: _ASTE. [...] "PASTE" is a word. PASTE
fits. Also "WASTE" also fits [...] Choose "PASTE" maybe.
Tool call: guess_word(guess="PASTE")
[tool] Guess 4/6: PASTE --- all green. Correct, you solved it!

Un rollout del modelo entrenado a partir de los trazos del runtime de AgentCore. El modelo mantiene las letras verdes en su lugar, nunca reutiliza una letra gris y usa cada letra amarilla para descartar posiciones, aunque nada en los datos de entrenamiento le mostró cómo hacerlo.

Paso 7: Evaluar el modelo entrenado en las configuraciones de servicio

Cuando el trabajo se completa, su página en SageMaker AI Studio muestra el estado y el grupo de paquetes de modelos que contiene el modelo entrenado. DescribeJob, con JobCategory establecido en AgentRFT, devuelve el ARN del paquete de modelos de salida y los recuentos de tokens facturables.

Figura 5: El trabajo de entrenamiento completado en SageMaker AI Studio en la pestaña Jobs, Training, Multi-turn RL

Un trabajo de evaluación (JobCategory establecido en AgentRFTEvaluation) reproduce sus prompts de evaluación a través del mismo runtime de AgentCore y la misma recompensa, sobre el modelo base o el paquete entrenado, y registra eval/reward/pass_at_1 en MLflow. A diferencia de la validación de entrenamiento, puede establecer la temperatura. El lanzador del repositorio ejecuta las evaluaciones una tras otra, porque la cuota predeterminada permite un solo trabajo de evaluación a la vez (puede solicitar un aumento):

uv run python run_mtrl_eval.py --runs base:1.0 trained:1.0 \
  --model-package-arn <OUTPUT_MODEL_PACKAGE_ARN> \
  --agent-runtime-arn <RUNTIME_ARN> \
  --role-arn arn:aws:iam::<ACCOUNT>:role/<JOB_ROLE> \
  --dataset s3://<YOUR_BUCKET>/wordle-mtrl/validation/validation-data.jsonl \
  --s3-output-path s3://<YOUR_BUCKET>/wordle-mtrl/eval/ \
  --mlflow-app-arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID>

Figura 6: Esfuerzo de razonamiento y temperatura de muestreo en cada etapa, donde el trabajo de entrenamiento fija la validación en temperatura 0 y el esfuerzo desplegado, mientras que los trabajos de evaluación le permiten elegir

Cada evaluación de 128 palabras tomó de cuatro a siete minutos. El esfuerzo de razonamiento proviene de las variables de entorno del agente, así que para cambiarlo usted edita REASONING_EFFORT, vuelve a desplegar y evalúa. Evaluamos ambos modelos con esfuerzo bajo y medio y a temperatura 0 y 1.0:

Figura 7: Modelo base frente a modelo entrenado bajo cuatro configuraciones de decodificación, con la tasa de resolución a la izquierda y la recompensa media por juego a la derecha

El segmento gris de cada barra es el modelo base y el segmento azul es la ganancia del entrenamiento. El par de esfuerzo bajo y temperatura 0 es la validación del trabajo de entrenamiento en los pasos 0 y 100. El resultado entrenado de esfuerzo medio y temperatura 1.0 es la media de dos ejecuciones (80 y 88 por ciento).

A temperatura 0, ambos modelos a menudo terminan el juego sin resolver, con cualquiera de los dos esfuerzos. Con el muestreo activado, el esfuerzo medio añade 15 puntos para el modelo entrenado y 46 para el modelo base. El adaptador añade 14–54 puntos en cada configuración, así que lo que el entrenamiento enseñó se transfiere entre configuraciones.

La propia curva de validación del trabajo de entrenamiento en MLflow contaba una historia diferente: alcanzó un máximo de 48 por ciento en el paso 90 y terminó en 37 por ciento, puntuada a temperatura 0. La referencia de hiperparámetros recomienda la temperatura 0 para la evaluación porque la decodificación voraz hace que las curvas sean deterministas y comparables entre ejecuciones. Informamos a temperatura 1.0 porque así es como servimos el modelo. Con esfuerzo medio y temperatura 1.0, el modelo entrenado resuelve el 84 por ciento de las palabras reservadas, frente al 61 por ciento del modelo base.

Consejo: Evalúe con las configuraciones que va a servir. La curva de entrenamiento dice que este adaptador termina en 37 por ciento. Con las configuraciones de servicio resuelve el 84 por ciento. Los trabajos de evaluación que muestran esto costaron alrededor de $1.20 en tokens y tomaron aproximadamente una hora.

Paso 8: Limpieza

Los trabajos de entrenamiento y evaluación dejan de facturar cuando se completan, pero el runtime de AgentCore, su rol de IAM, la aplicación de MLflow y los objetos de S3 persisten. Desde el directorio del proyecto:

# 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:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID>

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

Debido a que agentcore remove all reescribe los archivos bajo agentcore/, ejecute git checkout -- agentcore/ si desea volver a desplegar más tarde. Eliminar el stack de AgentCore-Wordle-default CloudFormation elimina los mismos recursos y deja los archivos intactos. El rol del trabajo de entrenamiento y los paquetes de modelos no se facturan.

Análisis profundo: Diseño de la recompensa

Nuestra recompensa necesitó una revisión, y el error es uno común.

El primer diseño sumaba la puntuación por intento del wordle-lora-rl project durante todo el juego. Esa puntuación funciona allí porque cada intento es una muestra de entrenamiento independiente. Sumada en una sola recompensa por partida, produjo dos inversiones que encontramos solo leyendo los registros:

  • Una partida perdida pagaba una recompensa positiva. Seis turnos de bonos sumaban aproximadamente +0.49.
  • Resolver en el turno 1 puntuaba menos que resolver en el turno 3. Resolver al primer intento ganaba +0.99 y resolver al tercer intento ganaba +1.61, porque los turnos extra acumulaban más bonos.

La corrección separa un término de resultado dominante de un término pequeño de modelado (shaping) acotado:

reward = outcome + w * clip(shaping, -1, 1), w = 0.04 if solved, else 0.15

El término de resultado se establece completamente según cómo terminó la partida:

Resultado Valor
Resuelto en el intento n 1.05 + 0.5 * (6 – n) / 5, así que 1.55 en el intento 1 hasta 1.05 en el intento 6
Se jugaron los seis intentos, sin resolver -0.5
Nunca se hizo un intento válido -1.5

El término de modelado es esa puntuación por intento, sumada a lo largo de la partida:

  • Penalizaciones por contradecir una letra verde, amarilla o gris conocida, por repetir un intento y por palabras fuera del diccionario.
  • Un bono de entropía por una palabra de apertura fuerte.
  • Bonos por probar letras nuevas y por reducir el conjunto de posibilidades restantes.

La suma se recorta a [-1, 1] y se escala por 0.04 en caso de resolver y 0.15 en caso de perder, de modo que distingue partidas con el mismo resultado pero nunca mueve una partida a otro nivel. Con valores típicos de modelado:

Escenario Recompensa
Resuelto en el intento 1 +1.550
Resuelto en el intento 3 +1.375
Resuelto en el intento 6 +1.069
Perdido tras seis intentos -0.427
Nunca hizo un intento válido -1.500

Una resolución más rápida ahora siempre supera a una más lenta, toda resolución supera a cualquier derrota, y no jugar nunca es el peor resultado. Nuestra primera corrección usó un único peso de 0.15, lo que permitió que una resolución en el intento 3 con formado completo (+1.45) superara a una resolución en el intento 2 con el mínimo (+1.25). Los conteos de intentos adyacentes están separados por 0.1, por lo que un peso de resolución menor de 0.04 elimina el solapamiento, y la peor resolución (+1.01) todavía supera el umbral de éxito de 1.0. El formado sigue importando porque GRPO aprende de la varianza de recompensa dentro de un grupo, así que dos partidas perdidas que difieren en la consistencia de las pistas todavía le enseñan algo al modelo. Como aconseja la documentación de SageMaker AI, comience con una recompensa basada en el resultado, añada formado con moderación y verifique la recompensa en trayectorias reales.

Análisis detallado: Qué salió mal, en el orden en que lo encontrará

Cada uno de estos costó al menos una ejecución de entrenamiento.

  • Tokens Bearer bajo concurrencia. Emitir un token con una sesión de botocore nueva por llamada agotó el proveedor de credenciales con 32 rollouts concurrentes, y almacenar tokens en caché durante una hora devolvía errores 403 a medida que las credenciales rotaban. La solución es una sola sesión con auto-renovación, una caché corta y una re-firma ante un 403.
  • Un rollout que nunca llega al modelo de política hace fallar todo el trabajo. Cuando todas las llamadas al modelo en un rollout fallaban la autenticación, el servicio terminaba el trabajo con «No sampling requests were received.». Devolver un estado de error también lo hace fallar, así que devuelva siempre una recompensa. Una partida en la que se muestrea al modelo pero nunca hace un intento válido está bien: puntúa -1.5, y la ejecución de 100 pasos tuvo 51 de ellas en unos 105,000 rollouts.
  • Las llamadas a herramientas se convierten en texto plano. Con una tasa de aprendizaje demasiado alta, el modelo escribía las llamadas a herramientas como texto, cada rollout recibía la misma recompensa mínima y el gradiente desaparecía. El agente ahora recupera los intentos del texto plano con una pequeña penalización de formato.
  • La meseta. La recompensa de entrenamiento dejó de subir aproximadamente después del paso 15 y la validación dejó de mejorar aproximadamente después del paso 30, aunque a temperatura 0 la validación subestima al modelo (Paso 7). Un currículo que inicialice partidas con intentos previos, o más rollouts por palabra que el valor predeterminado group_size de 8, es lo que intentaríamos a continuación.

Conclusión

MTRL en SageMaker AI entrena un modelo con el resultado de un episodio completo. Eso encaja con tareas agénticas, donde la trayectoria correcta es desconocida y solo el resultado puede puntuarse. Nosotros usamos una única recompensa a nivel de episodio. Cuando también se pueden puntuar pasos intermedios, MTRL acepta además una recompensa por turno mediante la API UpdateReward. En esta publicación, un modelo de pesos abiertos de 20B alcanzó una tasa de resolución del 84 por ciento en configuraciones de servicio, frente al 61 por ciento anterior, por unos $339 en tokens de entrenamiento y unos $79 en uso del runtime de AgentCore. El 37 por ciento que reporta la curva de entrenamiento es un artefacto de cómo el trabajo de entrenamiento puntúa la validación, así que evalúe con las configuraciones con las que servirá.

El mismo pipeline se aplica a procedimientos internos, herramientas y flujos de trabajo sobre los que los modelos públicos nunca fueron entrenados.

Para empezar, clone el repositorio de ejemplo, siga los pasos anteriores y ejecute la prueba de humo en el Paso 3 antes de comprometer tiempo de GPU. Para las APIs subyacentes, consulte la documentación de MTRL de SageMaker AI.

 


Sobre los autores

Fuente original

AWS Machine Learning

Notas sobre el contenido

La publicación original y los derechos pertenecen a la fuente.

Traducción automática · Consulte el original