Agent Lightning: обучите любого AI-агента с помощью RL на основе реальных трассировок

Что такое Agent Lightning и зачем это нужно

Microsoft выпустила Agent Lightning, open source фреймворк, который превращает реальные запуски агентов в переходы, пригодные для обучения с подкреплением, без переписывания существующего стека агентов. Он разделяет обучение и исполнение, стандартизирует формат трассировок и вводит LightningRL, иерархический метод, который конвертирует сложные многозвенные выполнения в переходы для оптимизации одновременных RL тренеров.

Моделирование агентов как процессов принятия решений

Agent Lightning формализует выполнение агента как частично наблюдаемую Марковскую процесс принятия решений. Наблюдение соответствует текущему вводу в policy LLM, действие соответствует вызову модели, а вознаграждение может быть терминальным или промежуточным. Из каждого прогона система извлекает только вызовы policy модели вместе с входами, выходами и вознаграждениями. Это отрезает шум фреймворка и дает чистые переходы для обучения.

LightningRL и распределение кредитов

LightningRL решает задачу назначения кредитов в многозвенных эпизодах. Он преобразует длинные траектории агента в переходы на уровне отдельных вызовов, сохраняя информацию, нужную для корректного распределения заслуг за ранние решения. Полученные переходы совместимы с целями однократного RL, поэтому команды могут повторно использовать стандартные RL тренеры и алгоритмы вроде PPO или GRPO.

Training Agent Disaggregation и архитектура системы

Дизайн разделяет рантайм и обучающую часть. Lightning Client выполняется рядом с существующими агентами в продакшене, захватывает трассы промптов, вызовов инструментов и сигналы вознаграждения. Lightning Server принимает потоковые трассы, выполняет обучение на GPU и обслуживает обновленные модели через API, совместимый с OpenAI. Это позволяет держать инструменты и зависимости близко к продакшену, а обучение централизовать и масштабировать.

Трассировка, телеметрия и единый интерфейс данных

Рантайм поддерживает два пути трассировки. По умолчанию используются OpenTelemetry spans для интеграции с стандартными коллекторами. Также предложен компактный встроенный трейсер для команд, которые не хотят разворачивать OpenTelemetry. Оба пути сходятся в одном хранилище для обучения.

Каждый вызов модели и вызов инструмента фиксируется как span с входами, выходами и метаданными. Алгоритмический слой преобразует spans в упорядоченные триплеты промпт, ответ и вознаграждение. Такая селективная экстракция позволяет оптимизировать один агент внутри многоагентного рабочего процесса или сразу несколько агентов, не меняя оркестрационный код. Те же трассы можно использовать для оптимизации промптов или обучения с учителем.

Эксперименты и наборы данных

В работе оценили три задачи с использованием Llama 3.2 3B Instruct в роли policy модели. Для text to SQL использовали бенчмарк Spider и агент на LangChain, включающий writer, rewriter и checker. Writer и rewriter оптимизировались, checker оставлялся фиксированным, а вознаграждения стабильно росли в процессе обучения и на тесте.

Для retrieval augmented generation применяли бенчмарк MuSiQue и индекс масштаба Wikipedia с примерно 21 миллионом документов. Ретривер использовал BGE эмбеддинги и косинусную схожесть, агент был построен на OpenAI Agents SDK. Вознаграждение было взвешенной суммой формата и F1 корректности, и кривые вознаграждения показали устойчивый рост.

Для математических вопросов с использованием инструментов агент на AutoGen вызывал калькулятор на датасете Calc X. Обучение улучшило способность корректно вызывать инструменты и интегрировать их результаты в итоговые ответы.

Основные выводы для команд

Agent Lightning обеспечивает интеграцию с минимальными изменениями для существующих фреймворков агентов, таких как LangChain, OpenAI Agents SDK, AutoGen или CrewAI. LightningRL преобразует многозвенные траектории в переходы и применяет назначение кредитов, чтобы стандартные однократные RL методы могли быть использованы. Automatic Intermediate Rewarding превращает системные сигналы в плотные промежуточные вознаграждения, снижая проблему разреженных вознаграждений в длинных рабочих процессах. Трассы, записанные через OpenTelemetry или встроенный трейсер, стримятся на сервер обучения, который выдает обновленные модели через OpenAI совместимый эндпоинт, что позволяет масштабно выкатывать модели без переноса инструментов.