Модуль 11. RL для LLM-агентов: инструменты, среда и длинные траектории

from __future__ import annotations

import os
import random
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from IPython.display import display


def find_repo_root(start: Path) -> Path:
    """Найти корень курса независимо от каталога запуска ноутбука."""
    start = start.resolve()
    for candidate in (start, *start.parents):
        if (candidate / "modules/module-11/module_11_reference.py").is_file():
            return candidate
    raise FileNotFoundError("Не найден корень репозитория RL for LLM")


ROOT = find_repo_root(Path.cwd())
MODULE_DIR = ROOT / "modules/module-11"
if str(MODULE_DIR) not in sys.path:
    sys.path.insert(0, str(MODULE_DIR))

from module_11_reference import (
    collector_target,
    discounted_return,
    discounted_td_sum,
    informative_group_probability,
    macro_action_logprob,
    potential_shaping,
    semantic_action_probability,
    semi_markov_return,
    shaped_return_identity,
    td_residuals,
    trajectory_mask_report,
    turn_probability_ratio,
    validate_rollout_record,
)
from module_11_agentic_training import (
    ACTIONS,
    KeyDoorAgentEnv,
    RecurrentActorCritic,
    collect_episode,
    evaluate_policy,
    rollout_contract_from_transitions,
    train_actor_critic,
)

PROFILE = os.getenv("RL_LLM_PROFILE", "smoke").strip().lower()
RUN_REAL_MODEL = False
if PROFILE not in {"smoke", "full", "real-model"}:
    raise ValueError("RL_LLM_PROFILE должен быть smoke, full или real-model")

if PROFILE == "full":
    SEEDS = (1, 2, 3, 4, 5)
    UPDATES = 30
    EPISODES_PER_UPDATE = 48
    EVALUATION_EPISODES = 100
else:
    SEEDS = (1, 2, 3)
    UPDATES = 8
    EPISODES_PER_UPDATE = 16
    EVALUATION_EPISODES = 20

np.random.seed(0)
random.seed(0)
torch.manual_seed(0)
torch.set_num_threads(1)

print(f"Корень курса: {ROOT}")
print(f"Профиль: {PROFILE}; семена={SEEDS}; обновления={UPDATES}; эпизодов на обновление={EPISODES_PER_UPDATE}")

env = KeyDoorAgentEnv(seed=0)
observation, reset_info = env.reset(target="red")
rows = []
for action in ("inspect", "pickup_red", "open_red", "finish"):
    next_observation, base_reward, terminated, truncated, info = env.step(action, gamma=0.99)
    rows.append({
        "наблюдение": observation,
        "действие": action,
        "следующее наблюдение": next_observation,
        "базовая награда": base_reward,
        "потенциал до": info["potential_before"],
        "потенциал после": info["potential_after"],
        "формированная награда": info["shaped_reward"],
        "истинное завершение": terminated,
        "внешнее усечение": truncated,
    })
    observation = next_observation

trace_df = pd.DataFrame(rows)
display(trace_df)
print("Сумма базовых наград:", trace_df["базовая награда"].sum())
print("Дисконтированный возврат сформированных наград:", discounted_return(trace_df["формированная награда"], 0.99))

segments = [
    ("system", 44),
    ("user", 63),
    ("assistant", 18),
    ("tool", 162),
    ("assistant", 21),
]
report = trajectory_mask_report(segments)
mask_table = pd.DataFrame(segments, columns=["роль", "число токенов"])
mask_table["входит в функцию потерь политики"] = mask_table["роль"].eq("assistant")
mask_table["роль"] = mask_table["роль"].map({
    "system": "системное сообщение",
    "user": "пользователь",
    "assistant": "политика",
    "tool": "инструмент",
})
display(mask_table)
print({
    "токенов политики": report.policy_tokens,
    "всего токенов": report.total_tokens,
    "доля исключённых токенов": report.masked_fraction,
    "коэффициент разбавления": report.dilution_factor,
})
print(
    f"Если ошибочно усреднить функцию потерь по всем {report.total_tokens} токенам, "
    f"сигнал {report.policy_tokens} токенов политики уменьшится в {report.dilution_factor:.2f} раза."
)

gamma = 0.9
boundary_table = pd.DataFrame([
    {
        "случай": "истинный терминал",
        "награда": 0.0,
        "V следующего состояния": 0.7,
        "целевая величина": collector_target(0.0, 0.7, gamma=gamma, terminated=True),
    },
    {
        "случай": "внешнее усечение сборщиком",
        "награда": 0.0,
        "V следующего состояния": 0.7,
        "целевая величина": collector_target(0.0, 0.7, gamma=gamma, terminated=False),
    },
])
display(boundary_table)

collector_env = KeyDoorAgentEnv(task_horizon=6, collector_limit=2, seed=0)
collector_env.reset(target="blue")
collector_env.step("inspect")
final_observation, reward, terminated, truncated, info = collector_env.step("pickup_blue")
print({
    "наблюдение": final_observation,
    "награда": reward,
    "истинное завершение": terminated,
    "внешнее усечение": truncated,
    "причина остановки": info["stop_reason"],
})

behavior_logprobs = [-0.42, -0.31, -0.27]
current_logprobs = [-0.35, -0.26, -0.24]

string_probabilities = {
    '{"name":"inspect"}': 0.25,
    '{"name": "inspect"}': 0.15,
    '{"name":"pickup_red"}': 0.60,
}
parser = {
    '{"name":"inspect"}': "inspect",
    '{"name": "inspect"}': "inspect",
    '{"name":"pickup_red"}': "pickup_red",
}

macro_table = pd.DataFrame([
    {"величина": "логарифм вероятности сгенерированной строки", "значение": macro_action_logprob(behavior_logprobs)},
    {"величина": "отношение политик за ход", "значение": turn_probability_ratio(current_logprobs, behavior_logprobs)},
    {"величина": "вероятность семантического действия inspect", "значение": semantic_action_probability(string_probabilities, parser, "inspect")},
    {"величина": "полумарковский возврат (2 и 3 шага)", "значение": semi_markov_return([0.0, 1.0], [2, 3], 0.9)},
])
display(macro_table)

gamma = 0.9
rewards = [0.0, 0.0, 1.0]
values = [0.4, 0.5, 0.6, 0.0]
deltas = td_residuals(rewards, values, gamma, terminated=[False, False, True])
left = discounted_td_sum(deltas, gamma)
right = discounted_return(rewards, gamma) - values[0] + gamma**len(rewards) * values[-1]

potentials = [0.0, 0.25, 0.55, 0.0]
shaped = potential_shaping(rewards, potentials, gamma)
original_return, shaped_return, predicted_return = shaped_return_identity(rewards, potentials, gamma)

display(pd.DataFrame({
    "t": range(len(rewards)),
    "r_t": rewards,
    "V_t": values[:-1],
    "V_next": values[1:],
    "delta_t": deltas,
    "Phi_t": potentials[:-1],
    "Phi_next": potentials[1:],
    "r_shaped": shaped,
}))
print(f"Телескопическая проверка: {left:.6f} = {right:.6f}")
print(f"Исходный возврат={original_return:.6f}; сформированный возврат={shaped_return:.6f}; формула={predicted_return:.6f}")

probabilities = np.linspace(0.001, 0.999, 500)
group_sizes = (4, 8, 16)
fig, ax = plt.subplots(figsize=(8.5, 4.8))
for group_size in group_sizes:
    values = [informative_group_probability(p, group_size) for p in probabilities]
    ax.plot(probabilities, values, label=f"G={group_size}")
ax.set_xlabel("Вероятность успеха одного роллаута p")
ax.set_ylabel("Вероятность смешанной группы")
ax.set_title("Группа информативна только при наличии успехов и провалов")
ax.legend()
ax.grid(alpha=0.25)
plt.show()

check_ps = (0.005, 0.05, 0.5, 0.95, 0.995)
group_df = pd.DataFrame([
    {"p": p, **{f"G={g}": informative_group_probability(p, g) for g in group_sizes}}
    for p in check_ps
])
display(group_df)

records = []
models = {}
for reward_mode in ("sparse", "potential"):
    for seed in SEEDS:
        model, curve = train_actor_critic(
            seed=seed,
            reward_mode=reward_mode,
            updates=UPDATES,
            episodes_per_update=EPISODES_PER_UPDATE,
            evaluation_interval=max(1, UPDATES // 5),
            evaluation_episodes=EVALUATION_EPISODES,
        )
        final_metrics = evaluate_policy(model, episodes=EVALUATION_EPISODES, seed=50_000 + seed)
        models[(reward_mode, seed)] = (model, final_metrics)
        for row in curve:
            records.append({"reward_mode": reward_mode, "seed": seed, **row})

curve_df = pd.DataFrame(records)
summary_df = (
    curve_df.groupby(["reward_mode", "update"], as_index=False)["success_rate"]
    .agg(median="median", low="min", high="max")
)
display(
    summary_df.rename(columns={
        "reward_mode": "режим награды",
        "update": "обновление",
        "median": "медиана",
        "low": "минимум",
        "high": "максимум",
    }).replace({"режим награды": {
        "sparse": "редкая итоговая",
        "potential": "потенциальное формирование",
    }})
)

fig, ax = plt.subplots(figsize=(8.5, 4.8))
for reward_mode, label in (("sparse", "Редкая итоговая награда"), ("potential", "Потенциальное формирование награды")):
    part = summary_df[summary_df["reward_mode"] == reward_mode].sort_values("update")
    x = part["update"].to_numpy(dtype=float)
    median = part["median"].to_numpy(dtype=float)
    low = part["low"].to_numpy(dtype=float)
    high = part["high"].to_numpy(dtype=float)
    ax.plot(x, median, marker="o", label=label)
    ax.fill_between(x, low, high, alpha=0.18)
ax.set_xlabel("Обновление политики (-1 означает оценку до обучения)")
ax.set_ylabel("Доля успешных жадных эпизодов")
ax.set_ylim(-0.03, 1.03)
ax.set_title("Синтетический CPU-эксперимент KeyDoorAgentEnv")
ax.grid(alpha=0.25)
ax.legend()
plt.show()

best_key, (best_model, best_metrics) = max(models.items(), key=lambda item: item[1][1]["success_rate"])
mode_ru = {
    "sparse": "редкая итоговая награда",
    "potential": "потенциальное формирование награды",
}[best_key[0]]
metrics_ru = {
    "доля успехов": best_metrics["success_rate"],
    "среднее число ходов": best_metrics["mean_turns"],
    "доля недопустимых действий": best_metrics["invalid_action_rate"],
}
print("Лучшая конфигурация для следующей диагностики:", (mode_ru, best_key[1]), metrics_ru)

rollout_env = KeyDoorAgentEnv(seed=123)
transitions, _, rollout_report = collect_episode(
    best_model,
    rollout_env,
    reward_mode="potential",
    greedy=True,
    target="red",
    policy_version="m11-notebook-policy",
)
trace = pd.DataFrame([
    {
        "ход": index,
        "наблюдение": transition.observation,
        "действие": transition.action,
        "следующее наблюдение": transition.next_observation,
        "логарифм вероятности поведения": transition.behavior_logprob,
        "базовая награда": transition.base_reward,
        "сформированная награда": transition.shaped_reward,
        "истинное завершение": transition.terminated,
        "внешнее усечение": transition.truncated,
    }
    for index, transition in enumerate(transitions)
])
display(trace)
print("Отчёт роллаута:", {
    "цель": rollout_report["target"],
    "успех": rollout_report["success"],
    "истинное завершение": rollout_report["terminated"],
    "внешнее усечение": rollout_report["truncated"],
    "причина остановки": rollout_report["stop_reason"],
    "версия политики": rollout_report["policy_version"],
    "выбранный возврат": rollout_report["selected_return"],
})

contract = rollout_contract_from_transitions(
    transitions,
    task_id="m11-keydoor-red",
    environment_version="1.0",
    policy_version="m11-notebook-policy",
)
print("Нарушения корректного контракта:", validate_rollout_record(contract))

broken_contract = dict(contract)
broken_contract["policy_versions"] = {}
broken_violations = validate_rollout_record(broken_contract)
translated_violations = [
    "отсутствует версия политики для одного или нескольких ходов"
    if message.startswith("missing policy versions") else message
    for message in broken_violations
]
print("Нарушения после удаления версии политики:", translated_violations)

if not RUN_REAL_MODEL:
    print(
        "Факультативная ветвь отключена. Для запуска создайте отдельную среду, "
        "зафиксируйте версии trl/transformers/datasets и установите RUN_REAL_MODEL = True."
    )
else:
    from datasets import Dataset
    from trl import GRPOConfig, GRPOTrainer

    class TextKeyDoorEnv:
        """Минимальная stateful-среда для официального environment_factory-контракта TRL."""

        def __init__(self):
            self.target = "red"
            self.inspected = False
            self.carrying = None
            self.opened = None
            self.reward = 0.0

        def reset(self, target: str, **kwargs) -> str:
            self.target = target
            self.inspected = False
            self.carrying = None
            self.opened = None
            self.reward = 0.0
            return "You are in a room with two locked doors. Inspect before choosing a key."

        def inspect(self) -> str:
            """Reveal the target color.

            Returns:
                A textual clue naming the required color.
            """
            self.inspected = True
            return f"The target door is {self.target}."

        def pickup(self, color: str) -> str:
            """Pick up one key.

            Args:
                color: Either ``red`` or ``blue``.

            Returns:
                The resulting environment observation.
            """
            if not self.inspected or color not in {"red", "blue"}:
                return "The pickup failed."
            self.carrying = color
            return f"You are carrying the {color} key."

        def open_door(self, color: str) -> str:
            """Try to open one door with the carried key.

            Args:
                color: Either ``red`` or ``blue``.

            Returns:
                The resulting environment observation.
            """
            if self.carrying != color:
                return "The door did not open."
            self.opened = color
            return f"The {color} door is open."

        def finish(self) -> str:
            """Finish the episode and store the task reward.

            Returns:
                A textual terminal observation.
            """
            self.reward = float(self.opened == self.target)
            return "success" if self.reward == 1.0 else "failure"

        def get_reward(self) -> float:
            return self.reward

    dataset = Dataset.from_dict({
        "prompt": [
            [{"role": "user", "content": "Inspect the room, choose the correct key, open the target door, and finish."}],
            [{"role": "user", "content": "Use the tools to open the required door and finish."}],
        ],
        "target": ["red", "blue"],
    })

    trainer = GRPOTrainer(
        model="Qwen/Qwen3-0.6B",
        args=GRPOConfig(
            num_generations=4,
            max_steps=20,
            max_completion_length=1024,
            chat_template_kwargs={"enable_thinking": False},
            output_dir="m11-agentic-grpo-output",
        ),
        train_dataset=dataset,
        environment_factory=TextKeyDoorEnv,
    )
    print("Trainer создан. trainer.train() намеренно не вызывается автоматически.")