Как обучать многошаговых AI‑агентов в Amazon SageMaker: практическое руководство по reinforcement learning — VogueTech

Что нового

Amazon запустила в SageMaker AI отдельный контур для обучения многошаговых (multi‑turn) агентов с помощью reinforcement learning (RL). Это не ещё один «чат с моделью», а сервис, который берёт на себя весь тяжёлый цикл обучения:

Готовый тренировочный цикл для многошаговых задач: SageMaker AI Multi‑Turn RL (MTRL) управляет rollout‑ами, сбором траекторий, обновлением градиентов и масштабированием.
Поддержка разных рантаймов для агентов:
- Amazon Bedrock AgentCore
- Amazon Elastic Kubernetes Service (Amazon EKS)
- Amazon Elastic Compute Cloud (Amazon EC2)
- AWS Fargate
- собственная инфраструктура через небольшой адаптер.
Модульный интерфейс «агент–среда»: вы определяете инструменты (tools), логику их вызова, форму диалога и функцию награды. Интеграция остаётся low‑code.
Полностью serverless‑исполнение: не нужно поднимать и обслуживать GPU‑кластеры. Оплата идёт «за токены» во время обучения, масштабирование прозрачно.
Асинхронные rollout‑ы и сбор траекторий: генерация и обновление модели идут параллельно. Сервис контролирует «устаревание» off‑policy данных, чтобы политика не уезжала слишком далеко от текущей.
Встроенная библиотека RL‑алгоритмов:
- Proximal Policy Optimization (PPO)
- Clipped Importance Sampling Policy Optimization (CISPO)
- Importance Sampling (IS)‑лоссы
- Групповые оценщики преимущества: GRPO, GRPO pass@k, RLOO и другие.
Обучение с удлинением последовательности (sequence‑extension): сервис сокращает wall‑clock время даже на длинных многошаговых сценариях.
Наблюдаемость в MLflow: SageMaker AI автоматически пишет в MLflow траектории, награды, метрики по шагам и по эпизодам.
Отдельные evaluation‑джобы до деплоя: сервис считает reward, pass@k, метрики по траекториям и форматам вывода до того, как вы выкатите модель в SageMaker endpoint или в Amazon Bedrock.

Все примеры в материале привязаны к SOP‑Bench — датасету от Amazon Science, который проверяет, как агенты выполняют задачи по сложным стандартным операционным процедурам (SOP) в 12 бизнес‑домейнах.

Как это работает

Архитектура SageMaker AI Multi‑Turn RL

SageMaker AI MTRL — это «оркестратор RL» вокруг вашего агента:

Вы пишете агента (LLM + tool‑calling логика) и поднимаете его на Bedrock AgentCore, EKS, EC2, Fargate или своей платформе.
Подключаете адаптер — тонкий слой, который:
- принимает запросы от rollout‑сервера,
- пробрасывает их агенту и его инструментам,
- возвращает в SageMaker AI сообщения диалога, tool‑calls и результаты.
Сервис запускает тысячи rollout‑ов:
- для каждого промпта создаётся эпизод,
- агент по шагам читает инструкции, вызывает инструменты, читает результаты, принимает решения и, при ошибке, пытается восстановиться,
- по завершении эпизода SageMaker AI вызывает вашу функцию награды.
Алгоритм RL обновляет политику по выбранному методу (PPO, CISPO, IS + GRPO/GRPO pass@k/RLOO и др.).
Всё логируется в MLflow, включая историю диалога, tool‑calls, награды и метрики.

Среда для многошагового RL

В одношаговом RL вам хватает промпта и функции награды. В многошаговом появляется ещё один критичный компонент — среда:

какие инструменты доступны агенту,
как они отвечают,
какие побочные эффекты создают.

Amazon рекомендует строить песочницу или симуляцию, максимально похожую на прод, но полностью изолированную от живых систем.

Ключевые паттерны:

Read‑only инструменты
- Используются для чтения данных без изменения состояния.
- Реализация: реплей заранее записанных ответов по ключу «вход → ответ».
- Пример из SOP‑Bench: customer‑service‑задача с 10 мок‑инструментами (validateAccount, getAuthenticationDetails, createSessionAndOpenTicket и др.), каждый возвращает детерминированный ответ из фикстуры (например, строку из CSV).
Состояние внутри эпизода (stateful tools)
- Нужны, когда агент что‑то записывает и потом читает обратно.
- Паттерн:
  - при старте rollout‑а выделяете ресурсы «на эпизод» (временный каталог, базу, ID),
  - регистрируете всё, что создаёт агент,
  - в try/finally полностью уничтожаете состояние по завершении эпизода — по финальному действию, max_turns или крэшу.
- Никакое состояние не должно утекать в следующий rollout.
Проверяемые результаты (verifiable outcomes)
- Подход для задач, где агент генерирует код, SQL или математику.
- Вы реально исполняете вывод агента в изолированной среде:
  - docker exec для кода,
  - in‑memory SQLite на rollout для SQL,
  - чистый eval в Python для математики.
- При тех же входах и том же состоянии песочницы результат детерминирован.
- В экосистеме AWS за это отвечает, например, AgentCore Code Interpreter.

У среды должны быть две жёсткие характеристики:

Воспроизводимость: один и тот же tool с одинаковыми аргументами всегда даёт один и тот же ответ. Тогда траектория с теми же действиями даёт ту же награду, а сравнение запусков честное.
Представительность: схемы данных, форматы запросов/ответов и распределения значений должны совпадать с тем, что ждёт агент в проде. Иначе вы обучите красивое поведение в лаборатории, которое не переносится в боевые сценарии.

Перед стартом обучения стоит проверить:

повторный запуск одного и того же инстанса задачи даёт идентичные rollout‑сообщения (diff без отличий);
у каждого rollout‑а свой temp‑каталог, ID и DB‑коннект;
набор инструментов и схемы запрос/ответ совпадают с прод‑окружением.

Внешняя оценка (evaluation) до награды

Amazon предлагает сначала построить внешнюю оценку, а уже потом писать функцию награды.

Паттерн:

пишете одну функцию score(rollout) -> float, которая честно измеряет финальный успех задачи;
запускаете модель через rollout‑сервер на фиксированном тестовом сплите;
считаете одну метрику успеха (например, долю задач, решённых до конца).

На SOP‑Bench это выглядит так:

агент должен выдать JSON‑объект внутри тега <final_output>;
оценка — строгий exact‑match по всем полям: если хоть одно поле не совпало с эталоном, эпизод получает 0.

Эта функция оценки не знает про частичный кредит и веса полей. Награда может это учитывать, а внешняя оценка — нет.

Перед обучением стоит:

прогнать базовую модель (которую вы будете дообучать) через эту оценку;
прогнать референсную модель (например, сильную LLM в Bedrock) по тому же тесту.

Так вы понимаете, где стартуете и к чему есть смысл стремиться.

Почему нельзя смотреть только на reward из RL:

если вы платите наградой за каждый tool‑call — агент научится вызывать инструменты бесконечно;
если штрафуете за длину диалога — агент начнёт отвечать раньше, чем соберёт нужные данные.

Reward растёт, а реальное качество решения задач падает — это классический reward hacking.

Дизайн награды для многошагового RL

Награда в RL — главный источник проблем. Агент оптимизирует то, что вы формализовали, а не то, что вы «имели в виду».

Базовая рекомендация: используйте тот же принцип оценивания, что и во внешней метрике, и отходите от него только по очень понятным причинам.

На SOP‑Bench агент должен выдать JSON внутри <final_output>:

{
  "aircraft_ready": "true",
  "mechanical_inspection_result": "success",
  "electrical_inspection_result": "success",
  "component_incident_response": "success",
  "component_mismatch_response": "success",
  "cross_check_reporting_response": "success"
}

Бенчмарк даёт 1, если совпали все поля, и 0 в остальных случаях.

Почему иногда стоит отойти от бинарной оценки и сделать награду «плотной» (dense):

Алгоритмические причины
- В SageMaker AI MTRL обучение идёт по группам из group_size rollout‑ов на один промпт.
- Групповой метод преимущества (advantage_method) по умолчанию — group_based (GRPO). Доступны и другие: rloo, grpo_passk и др.
- Если награда бинарная, а все rollout‑ы в группе получают одинаковый балл, вариация внутри группы исчезает, и градиент с этого батча почти нулевой.
- В логах это видно как падение rollout/reward/valid_mean относительно rollout/reward/mean и стагнацию качества.
Скорость сходимости
- При плотной награде каждый rollout даёт сигнал: эпизод, где 5 из 6 полей верны, учит модель, «как выглядит прогресс».
- Бинарный reward не даёт модели информацию о том, что этот эпизод лучше, чем тот, где не совпало ни одно поле.

Пример плотной награды для 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}

Агент отдаёт reward через update_reward, а метрики completion и field_acc попадают в MLflow. При желании можно отдавать список наград по шагам, а не один скаляр на эпизод, и использовать group_based_per_turn — тогда вы поощряете или штрафуете отдельные ходы.

Критический момент — верифицировать награду на реальных выводах до обучения. Пример из практики SOP‑Bench:

reward‑парсер принимал более свободный формат, чем официальный бенчмарк;
<final_response> вместо <final_output> всё равно получал кредит;
модель быстро «научилась» опускать нужный тег;
reward рос, а внешняя метрика по бенчмарку падала.

Проверка базовой модели и MultiTurnRLEvaluator

RL усиливает то, что модель уже умеет хоть немного. Если базовая модель ни разу не решает задачу, награде просто не за что зацепиться.

SageMaker AI MTRL умеет запускать managed evaluation‑джобы через MultiTurnRLEvaluator. Он:

прогоняет агента по отложенному набору промптов,
считает eval/reward и pass@k.

Если вы уже обучили модель, можно одной опцией evaluate_base_model=True одновременно посчитать метрики и для базовой, и для дообученной модели. pass@k использует порог success_threshold: при success_threshold=1 вы получаете строгую долю идеально успешных rollout‑ов.

Пример кода с Bedrock AgentCore:

from sagemaker.train.evaluate import MultiTurnRLEvaluator

# With Bedrock AgentCore
evaluator_base = MultiTurnRLEvaluator(
    model="openai-reasoning-gpt-oss-20b",
    dataset="s3://my-bucket/eval-prompts.parquet",
    agent_config="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent",
    s3_output_path="s3://my-bucket/eval-output/base/",
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-mlflow",
    role="arn:aws:iam::123456789012:role/SageMakerRole",
    accept_eula=True,
)

execution = evaluator_base.evaluate()
execution.wait()

В s3_output_path вы получите метрики и траектории, а в MLflow — удобный дашборд по этим данным.

Нюанс: MultiTurnRLEvaluator использует вашу же функцию награды, поэтому он показывает обобщающую способность (generalization) на отложенном наборе, но не защищает от ошибок в самом reward‑парсере. Чтобы поймать такие баги, нужно:

прогнать те же rollout‑ы через независимый, более строгий скорер (например, официальный SOP‑Bench exact‑match);
сравнить результаты.

Можно даже собрать отдельного «агента», у которого reward — это и есть независимая метрика, и запустить MultiTurnRLEvaluator поверх него.

Контекст и лимиты в многошаговом режиме

Многошаговый агент каждый ход расширяет контекст: запрос, аргументы tool‑call‑а, ответ инструмента, размышления между шагами. На длинных задачах контекст быстро становится огромным.

SageMaker AI MTRL использует обучение с удлинением последовательности, чтобы не раздувать wall‑clock время при росте длины эпизодов.

У вас есть два ключевых лимита:

max_turns — максимальное количество ходов, которое контролирует ваш агентный цикл;
sampling_max_tokens (для rollout‑ов) и val_sampling_params.sampling_max_tokens (для evaluation) — максимум токенов на ход.

Эти параметры нужно подбирать под:

сложность задач (сколько шагов реально нужно);
бюджет на inference в проде (вы не хотите обучить агента, который в бою всегда выжигает 8–10 ходов и 8k токенов).

Что это значит для вас

Где это полезно

SageMaker AI MTRL нацелен на агентов, которые должны действовать по шагам, а не просто отвечать текстом. Практические сценарии:

Саппорт и тикеты
- Агент читает обращение,
- проверяет аккаунт через несколько инструментов,
- достаёт историю заказов,
- принимает решение о возврате,
- создаёт тикет и формирует итоговый ответ.
Модерация контента
- Многошаговая проверка по разным политикам и базам правил,
- запрос контекста из внутренних систем,
- формирование отчёта о нарушении.
Операционные процедуры по SOP
- Инспекция самолёта, онбординг клиента, KYC, проверка документации,
- длинные цепочки действий, где ошибка на шаге 3 ломает шаг 7.
Код‑ и SQL‑агенты
- Генерация кода, запуск его в песочнице, анализ ошибок, повторная попытка,
- построение SQL‑запросов, их запуск в in‑memory базе, корректировка.

Во всех этих задачах важно не только «что ответил агент», но и как он к этому пришёл: какие инструменты вызвал, какие данные использовал, где исправил ошибку.

Кому это реально нужно

Командам с продакшн‑нагрузкой на AWS:
- если вы уже живёте в экосистеме AWS и используете SageMaker, Bedrock, EKS или Fargate, MTRL логично ложится в существующий стек;
- вы получаете готовую инфраструктуру RL вместо собственного велосипеда на Ray RLlib или custom PyTorch.
Продуктам, где цена ошибки высока:
- финансовые сервисы, авиаперевозки, логистика,
- сложные B2B‑процессы с регуляторными требованиями.
Командам, которые готовы инвестировать в reward design и симуляцию:
- вам нужно построить песочницу, написать честную функцию награды, настроить evaluation.
- без этого MTRL не принесёт пользы.

Где лучше не использовать

Простые Q&A‑чаты и ассистенты без инструментов:
- в большинстве случаев дешевле и проще fine‑tune‑ить модель supervised‑обучением или использовать готовые LLM без RL.
Проекты без доступа к AWS или с жёсткими ограничениями по региону размещения данных.
Команды без ресурса на MLOps:
- если нет людей, которые смогут поддерживать MLflow, мониторинг, симуляции и пайплайны обучения, MTRL будет избыточен.

Доступность из России

SageMaker AI, Bedrock и связанные сервисы официально работают в регионах AWS. Для доступа из России вам понадобится:

аккаунт AWS с привязанным платёжным методом;
выход в интернет до регионов AWS (обычно через корпоративный канал или VPN);
соблюдение юридических и комплаенс‑ограничений вашей компании.

Если у вас нет возможности использовать AWS по юридическим или инфраструктурным причинам, этот стек будет недоступен.

Место на рынке

Сегмент, в котором играет SageMaker AI MTRL, — это обучение агентных LLM с RL на продакшн‑инфраструктуре. Ближайшие по духу решения — это самописные пайплайны поверх PyTorch/TRL/RLlib и кастомные оркестраторы вокруг LLM‑агентов.

Конкретные отличия SageMaker AI MTRL:

Интеграция с AWS‑инфраструктурой:
- нативная работа с Bedrock AgentCore, EKS, EC2, Fargate;
- управление MLflow «из коробки».
Готовые RL‑алгоритмы под многошаговые сценарии:
- PPO, CISPO, IS‑лоссы;
- групповые advantage‑оценщики (GRPO, GRPO pass@k, RLOO и др.),
- всё это уже обёрнуто в multi‑turn training loop.
Serverless‑модель:
- нет необходимости руками собирать и масштабировать GPU‑кластеры;
- биллинг по токенам обучения.

Если сравнивать с «сделать всё самому», SageMaker AI MTRL выигрывает за счёт:

меньшего объёма инфраструктурной работы;
встроенной наблюдаемости и evaluation;
стандартных паттернов для симуляции среды и reward design.

При этом он привязан к AWS и экосистеме SageMaker/Bedrock. Если вы строите всё on‑prem или в другом облаке, этот вариант вам не подойдёт, и придётся собирать аналоги самостоятельно.

Как запустить

Ниже — пример того, как использовать MultiTurnRLEvaluator для оценки базовой модели с Bedrock AgentCore. Код из оригинального руководства сохранён полностью.

from sagemaker.train.evaluate import MultiTurnRLEvaluator

# With Bedrock AgentCore
evaluator_base = MultiTurnRLEvaluator(
    model="openai-reasoning-gpt-oss-20b",
    dataset="s3://my-bucket/eval-prompts.parquet",
    agent_config="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent",
    s3_output_path="s3://my-bucket/eval-output/base/",
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-mlflow",
    role="arn:aws:iam::123456789012:role/SageMakerRole",
    accept_eula=True,
)

execution = evaluator_base.evaluate()
execution.wait()

Практический чек‑лист перед запуском полноценного обучения:

Среда:
- есть песочница или симуляция, отделённая от прод‑систем;
- инструменты детерминированы при одинаковых входах;
- состояние не течёт между rollout‑ами.
Оценка:
- есть честная функция score(rollout) -> float, совпадающая с бизнес‑целью;
- базовая модель даёт ненулевой результат;
- есть референсная сильная модель для ориентира.
Награда:
- использует тот же принцип оценивания, что и внешняя метрика, или отклоняется по понятным причинам;
- даёт значения в [0, 1] (или [-1, 1], если есть штрафы);
- имеет ненулевую дисперсию на 100+ базовых rollout‑ах;
- не даёт ни одному эпизоду награду выше, чем строгая внешняя оценка;
- логирует компоненты награды по отдельности в MLflow.
Многошаговый режим:
- выставлены адекватные max_turns и лимиты токенов;
- вы понимаете, сколько шагов и токенов готовы тратить в проде.

Если эти пункты закрыты, SageMaker AI MTRL даёт вам промышленный контур для обучения многошаговых агентов — от первых rollout‑ов до развёртывания на SageMaker endpoint или в Amazon Bedrock.