Как оптимизировать траектории агента (trajectory optimization)?

Краткий тезис

Траектория агента — это последовательность шагов (observations, thoughts, actions), которые агент выполняет для решения задачи. Оптимизация траектории направлена на сокращение длины, повышение успешности и уменьшение стоимости (токенов, времени) при сохранении качества ответа. Основные методы: сбор и логирование успешных/неуспешных траекторий, абляция (удаление избыточных шагов), слияние похожих шагов (например, объединение нескольких поисков в батч), дистилляция (обучение маленького агента имитировать короткие траектории большого) и обучение с подкреплением (RL) с штрафом за длину траектории.

1. Что такое траектория агента и зачем её оптимизировать

Термин траектория агента (agent trajectory) — это полный лог взаимодействия агента с окружением: наблюдения (observation), мысли (thought), действия (action) и результаты действий (observation). В контексте Agentic RAG агент может искать документы, вызывать инструменты, уточнять запросы, проверять факты.

Зачем оптимизировать:

• Стоимость: каждый шаг потребляет токены LLM и время выполнения.
• Надёжность: длинные траектории склонны к накоплению ошибок (ошибка на раннем шаге каскадируется).
• Пользовательский опыт: пользователь ждёт быстрого ответа, длительные цепочки шагов увеличивают latency.

Термин оптимизация траектории (trajectory optimization) — набор техник, позволяющих получить более короткую, эффективную и надёжную последовательность шагов агента без потери качества конечного ответа.

2. Сбор данных: логирование успешных и неуспешных траекторий

Первый шаг — логгирование всех траекторий в production или при тестировании. Нужно сохранять:

• Исходный запрос пользователя.
• Полную историю шагов (действия, мысли, результаты).
• Финальный ответ и метку успешности (успех / неуспех), например, по оценке пользователя или автоматической проверке.

Термин успешная траектория (successful trajectory) — та, которая привела к правильному ответу. Неуспешная (failed trajectory) — либо привела к неверному ответу, либо превысила лимит шагов, либо агент зациклился.

Сбор данных позволяет:

• Выявить частые паттерны ошибок (например, агент дважды ищет один и тот же документ).
• Обучить модели ранжирования шагов.
• Сформировать датасет для методов, описанных ниже.

Пример структуры лога:

{
  "query": "Какая зарплата у ML-инженера в Германии?",
  "trajectory": [
    {"thought": "Нужно найти свежие данные по зарплатам", "action": "search(query='ML engineer salary Germany 2025')", "observation": "Результаты поиска: ..."},
    {"thought": "Возможно, нужен фильтр по городам", "action": "search(query='ML engineer salary Berlin 2025')", "observation": "..."},
    ...
  ],
  "final_answer": "Средняя зарплата ...",
  "success": true,
  "total_steps": 3,
  "cost_tokens": 4500
}

3. Абляция шагов (Ablation)

Термин абляция (ablation) — удаление одного или нескольких шагов из траектории с последующей проверкой, изменился ли ответ. Если ответ остался корректным, шаг считается избыточным.

Методика:

1. Взять успешные траектории из логов.
2. Для каждой траектории поочерёдно удалять каждый шаг (начиная с последнего).
3. Прогнать оставшуюся траекторию через агента (или симулировать, используя сохранённые observation предыдущих шагов).
4. Сравнить ответ с эталонным (golden answer). Если ответ совпадает — шаг можно исключить.
5. Статистически определить, какие шаги чаще всего оказываются лишними.

Термин эталонный ответ (golden answer) — правильный ответ, полученный от эксперта или из доверенного источника.

Пример: агент делает два поиска подряд, но второй не добавляет новой информации — его можно удалить.

def ablate_step(trajectory, step_index):
    # Удаляем шаг с индексом step_index
    new_traj = trajectory[:step_index] + trajectory[step_index+1:]
    # Запускаем агента с новой траекторией (первые шаги уже выполнены)
    # и смотрим, изменился ли ответ
    answer = run_agent_with_partial_trajectory(new_traj)
    return is_correct(answer, golden_answer)

for traj in successful_trajectories:
    for i in range(len(traj)):
        if ablate_step(traj, i):
            print(f"Шаг {i} можно удалить")

Абляция даёт понимание минимально необходимого набора шагов для конкретной задачи.

4. Слияние шагов (Merging)

Термин слияние (merging) — объединение нескольких похожих или последовательных шагов в один, более мощный. Примеры:

• Два search → один batch search: если агент выполняет поиск по двум связанным запросам, можно объединить их в один запрос с уточнением. Например, search('salary Munich') и search('salary Berlin') → search('salary in Munich and Berlin') или выполнить параллельный поиск.
• Два tool call → один: если агент вызывает get_weather и get_time для одного города, можно объединить в get_weather_and_time.
• Thought + Action → один шаг: агент сначала думает, потом действует; иногда можно объединить мысль и действие в одну LLM-итерацию, если мысль тривиальна.

Метод требует анализа частых паттернов. Можно автоматически искать последовательности шагов, которые всегда выполняются вместе, и предлагать их объединение.

from collections import Counter

pairs = Counter()
for traj in trajectories:
    for i in range(len(traj)-1):
        pair = (traj[i]['action_type'], traj[i+1]['action_type'])
        pairs[pair] += 1

# Найти пары, которые встречаются в >80% траекторий
frequent_pairs = [(p, c) for p, c in pairs.items() if c / len(trajectories) > 0.8]

Слияние сокращает количество шагов и уменьшает latency, но может усложнить реализацию инструмента.

5. Дистилляция (Distillation) траекторий

Термин дистилляция траекторий (trajectory distillation) — обучение маленькой модели (студента) генерировать короткие траектории, имитируя успешные траектории большой модели (учителя). Большая модель (например, GPT-4) генерирует успешные траектории, а маленькая (например, LLaMA-7B) учится повторять их, но с меньшим числом шагов.

Шаги:

1. Собрать датасет успешных траекторий от учителя (большой LLM).
2. Применить абляцию и слияние, чтобы получить оптимизированные траектории (ещё более короткие).
3. Обучить студента на парах (запрос → последовательность действий и мыслей). При обучении можно добавить penalty за длину (штрафовать за большое количество шагов).
4. После обучения студент способен сам генерировать короткие траектории для новых запросов.

Термин имитационное обучение (behavioral cloning) — частный случай дистилляции, когда студент учится по записям действий учителя без дополнительного взаимодействия со средой.

from transformers import Trainer, TrainingArguments

dataset = load_dataset("trajectories")  # содержит (query, optimized_trajectory)
model = AutoModelForCausalLM.from_pretrained("student-base")

training_args = TrainingArguments(
    output_dir="./student_distilled",
    per_device_train_batch_size=8,
    num_train_epochs=3,
    learning_rate=5e-5,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

Дистилляция даёт компактного агента с низкой стоимостью вызова, но требует качественных данных учителя.

6. Обучение с подкреплением (RL) с наградой за длину

Термин обучение с подкреплением (Reinforcement Learning, RL) — агент обучается взаимодействовать со средой, максимизируя накопленную награду. Для траекторной оптимизации награда включает:

• +1 за успешный ответ (или больший positive, если ответ правильный).
• -0.1 за каждый шаг (штраф за длину).
• -1 за неуспех (или за превышение максимального числа шагов).

Таким образом, агент стремится решить задачу за минимальное количество шагов. Можно использовать алгоритмы PPO или REINFORCE.

Этапы:

1. Определить среду (набор инструментов, эмулятор RAG).
2. Определить пространство действий (какие инструменты доступны).
3. Задать функцию награды.
4. Инициализировать политику (LLM-агент) и запустить обучение.

Термин функция награды (reward function) — правило, которое после каждого эпизода вычисляет численную оценку поведения агента.

Пример функции награды:

def reward(trajectory, final_answer_correct):
    base = 10 if final_answer_correct else -10
    for step in trajectory:
        base -= 0.5  # штраф за каждый шаг
    # дополнительный штраф за повторяющиеся действия
    if has_duplicate_actions(trajectory):
        base -= 5
    return base

RL-оптимизация может открыть новые неочевидные короткие пути, но обучение нестабильно и требует много эпизодов.

7. Практические инструменты и метрики

Для измерения качества оптимизации траекторий нужны метрики:

Инструменты:

• Weights & Biases — логирование траекторий и метрик.
• LangSmith / Langfuse — специализированные платформы для трейсинга агентов.
• NeMo Guardrails — для добавления ограничений на количество шагов.

8. Сравнение методов оптимизации

Обычно комбинируют: сначала абляция и слияние как быстрые improvements, затем дистилляция для постоянной модели, и RL для тонкой настройки.

9. Пет-проект для закрепления

Задача: Создать агента для ответов на вопросы по документации Python, который использует два инструмента: поиск по docstring (search_doc) и парсинг примера кода (parse_example). Оптимизировать его траектории.

Инструменты: Python, transformers, langchain, gym (для RL), wandb.

Шаги:

1. Реализовать агента с фиксированной политикой (вызывать search_doc, потом parse_example, потом генерировать ответ).
2. Собрать 1000 траекторий в симуляции (задать вопросы, для которых известен правильный ответ).
3. Провести абляцию: удалить вызов parse_example и проверить, как часто ответ остаётся верным (возможно, для простых вопросов он не нужен).
4. Применить слияние: объединить два вызова search_doc (если они идут подряд) в один с более широким запросом.
5. Обучить маленькую модель (DistilBERT) дистилляцией на оптимизированных траекториях.
6. Обучить политику с помощью PPO, добавив штраф за каждый шаг.

Ожидаемый результат:

• Средняя длина траектории сократится с 4 до 2 шагов.
• Успешность останется на уровне >90%.
• Стоимость токенов уменьшится в 2 раза.

10. Связь с другими вопросами

Навигация

• Предыдущий: 894
• Следующий: 896
• Индекс: 00. Индекс разборов