Модуль 5. Обучение по фиксированным данным: имитация, глубокие Q-методы и офлайн-RL

from __future__ import annotations

from pathlib import Path
import os
import sys
import time
import warnings

import matplotlib.pyplot as plt
import numpy as np


def find_course_root() -> Path:
    candidates = [Path.cwd(), *Path.cwd().parents]
    for candidate in candidates:
        if (candidate / "modules/module-05/module_5_training.py").is_file():
            return candidate
    raise FileNotFoundError("Could not locate the RL for LLM course root")

ROOT = find_course_root()
MODULE = ROOT / "modules/module-05"
DATA = MODULE / "data"
sys.path.insert(0, str(MODULE))

from module_5_reference import (
    behavior_support_mask,
    cql_discrete_penalty,
    double_dqn_targets,
    dqn_targets,
    expectile_loss,
    fqe_targets,
    iql_advantage_weights,
    return_to_go,
)
from module_5_training import (
    ACTION_NAMES,
    KeyDoorGrid,
    collect_offline_dataset,
    dagger_experiment,
    evaluate_decision_transformer,
    offline_experiment,
    train_decision_transformer,
    train_online_dqn,
)
from module_5_language import language_offline_experiment, token_text, PROMPTS

PROFILE = os.environ.get("RL_LLM_M5_PROFILE", "smoke").strip().lower()
if PROFILE not in {"smoke", "full"}:
    raise ValueError("RL_LLM_M5_PROFILE must be 'smoke' or 'full'")

# Full mode regenerates the numerical experiments. Smoke mode executes smaller
# end-to-end loops and loads pinned full-profile artifacts for stable figures.
print({"course_root": str(ROOT), "profile": PROFILE, "numpy": np.__version__})


LABELS = {
    "grid_title": "KeyDoorGrid: ключ, дверь, цель, ловушки и разные типы окончания",
    "grid_x": "столбец",
    "grid_y": "строка",
    "dagger_x": "раунд агрегации данных",
    "success": "доля успешных эпизодов",
    "dagger_title": "Зафиксированный прогон: BC и последующие раунды DAgger",
    "episode": "эпизод обучения",
    "rolling_success": "скользящая доля успеха",
    "dqn_title": "DQN и Double DQN: один зафиксированный seed, окно 15",
    "count": "число переходов",
    "log_title": "Состав действий в фиксированном журнале",
    "mean_return": "средняя награда",
    "offline_title": "Один журнал, разные защитные механизмы",
    "fqe_warning": "FQE — отдельная обучаемая модель; её калибровку нельзя предполагать по названию метода.",
    "identity": "достигнуто = запрошено",
    "requested_return": "запрошенный return-to-go",
    "achieved_return": "наблюдаемая средняя награда",
    "dt_title": "Условие Decision Transformer — запрос, а не гарантия",
    "exact_match": "доля точных ответов",
    "language_title": "Синтетический токеновый офлайн-контур",
    "summary_heading": "Ключевые выводы практикума",
    "summary_bc": "BC устойчиво решает задачу на распределении демонстраций, но DAgger показывает, как собирать метки на состояниях ученика.",
    "summary_dqn": "Replay и target network стабилизируют online Q-learning, а Double DQN разделяет выбор и оценивание.",
    "summary_offline": "Фиксированный журнал превращает покрытие в ограничение идентифицируемости; FQE и пессимизм сами требуют проверки.",
    "summary_dt": "Decision Transformer обучается как условная модель последовательности, но return-to-go не является обещанием результата.",
    "summary_llm": "Токеновые Q/V-функции возможны, однако уменьшенная ILQL-подобная реализация не подменяет полноразмерный LLM-эксперимент.",
}

env = KeyDoorGrid(seed=5)
obs, info = env.reset(start=(6, 0), seed=5)
print("observation_dim=", env.observation_dim, "action_dim=", env.action_dim)
print("initial_state=", info["state"])

# A compact vector drawing keeps the environment readable without requiring
# image or environment packages.
fig, ax = plt.subplots(figsize=(10, 5.5))
ax.set_xlim(0, env.width)
ax.set_ylim(env.height, 0)
ax.set_aspect("equal")
ax.set_xticks(np.arange(env.width + 1))
ax.set_yticks(np.arange(env.height + 1))
ax.grid(True, linewidth=0.8)
for r, c in env.wall_positions:
    ax.add_patch(plt.Rectangle((c, r), 1, 1, alpha=0.65))
for r, c in env.trap_positions:
    ax.text(c + 0.5, r + 0.5, "×", ha="center", va="center", fontsize=22)
kr, kc = env.key_position
ax.text(kc + 0.5, kr + 0.5, "K", ha="center", va="center", fontsize=18, fontweight="bold")
dr, dc = env.door_position
ax.text(dc + 0.5, dr + 0.5, "D", ha="center", va="center", fontsize=18, fontweight="bold")
gr, gc = env.goal_position
ax.text(gc + 0.5, gr + 0.5, "G", ha="center", va="center", fontsize=18, fontweight="bold")
for r, c in env.start_positions:
    ax.plot(c + 0.5, r + 0.5, ".", markersize=6)
ax.set_title(LABELS["grid_title"])
ax.set_xlabel(LABELS["grid_x"])
ax.set_ylabel(LABELS["grid_y"])
plt.show()

# Exact local mechanisms used later inside the neural loops.
r = np.array([0.2, -1.0])
terminal = np.array([0.0, 1.0])
next_q = np.array([[1.1, 1.6], [10.0, 20.0]])
print("DQN targets:", dqn_targets(r, terminal, next_q, gamma=0.95))
print("CQL term:", cql_discrete_penalty(np.array([2.0, 0.0, -1.0]), 0))
print("return-to-go:", return_to_go([1.0, 0.0, -0.5, 2.0]))

start = time.perf_counter()
if PROFILE == "full":
    dagger_run = dagger_experiment(
        seed=17, expert_episodes=4, rounds=3,
        learner_episodes_per_round=8, epochs=12,
        evaluation_episodes=24,
    )
else:
    dagger_run = dagger_experiment(
        seed=17, expert_episodes=4, rounds=2,
        learner_episodes_per_round=6, epochs=6,
        evaluation_episodes=12,
    )
print(f"DAgger wall time: {time.perf_counter() - start:.2f}s")
for row in dagger_run["records"]:
    print(row)

pinned = np.load(DATA / "M5_dagger_curve.npz")
fig, ax = plt.subplots(figsize=(8.5, 4.8))
ax.plot(pinned["rounds"], pinned["success_rate"], marker="o")
ax.set_ylim(-0.03, 1.05)
ax.set_xticks(pinned["rounds"])
ax.set_xlabel(LABELS["dagger_x"])
ax.set_ylabel(LABELS["success"])
ax.set_title(LABELS["dagger_title"])
ax.grid(True, alpha=0.3)
plt.show()

episodes = 130 if PROFILE == "full" else 35
eval_episodes = 24 if PROFILE == "full" else 12
start = time.perf_counter()
dqn_run = train_online_dqn(seed=5, episodes=episodes, double=False, evaluation_episodes=eval_episodes)
double_run = train_online_dqn(seed=5, episodes=episodes, double=True, evaluation_episodes=eval_episodes)
print(f"DQN pair wall time: {time.perf_counter() - start:.2f}s")
print("DQN evaluation:", dqn_run["evaluation"])
print("Double DQN evaluation:", double_run["evaluation"])

pinned = np.load(DATA / "M5_online_dqn_curve.npz")
window = 15

def moving_mean(x, window):
    x = np.asarray(x, dtype=float)
    if len(x) < window:
        return x
    return np.convolve(x, np.ones(window) / window, mode="valid")

fig, ax = plt.subplots(figsize=(9, 4.8))
ax.plot(np.arange(window, len(pinned["dqn_success"]) + 1), moving_mean(pinned["dqn_success"], window), label="DQN")
ax.plot(np.arange(window, len(pinned["double_success"]) + 1), moving_mean(pinned["double_success"], window), label="Double DQN")
ax.set_ylim(-0.03, 1.05)
ax.set_xlabel(LABELS["episode"])
ax.set_ylabel(LABELS["rolling_success"])
ax.set_title(LABELS["dqn_title"])
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

episodes_log, transitions_log = collect_offline_dataset(
    seed=41,
    episodes=80 if PROFILE == "full" else 30,
    profile="mixed",
)
actions = np.asarray([t.action for t in transitions_log], dtype=int)
counts = np.bincount(actions, minlength=4)
probabilities = counts / counts.sum()
mask = behavior_support_mask(probabilities, threshold_ratio=0.35)
print("action counts:", dict(zip(ACTION_NAMES, counts.tolist())))
print("behavior probabilities:", dict(zip(ACTION_NAMES, probabilities.round(4).tolist())))
print("support mask (relative threshold 0.35):", dict(zip(ACTION_NAMES, mask.tolist())))
print("dataset return mean:", float(np.mean([ep.return_ for ep in episodes_log])))

fig, ax = plt.subplots(figsize=(8, 4.5))
ax.bar(ACTION_NAMES, counts)
ax.set_ylabel(LABELS["count"])
ax.set_title(LABELS["log_title"])
ax.grid(True, axis="y", alpha=0.3)
plt.show()

start = time.perf_counter()
if PROFILE == "full":
    offline_run = offline_experiment(seed=41, dataset_episodes=80, train_epochs=16, evaluation_episodes=24)
else:
    offline_run = offline_experiment(seed=41, dataset_episodes=30, train_epochs=5, evaluation_episodes=12)
print(f"offline-method wall time: {time.perf_counter() - start:.2f}s")
for name, metrics in offline_run["results"].items():
    print(f"{name:>18s}: {metrics}")

pinned = np.load(DATA / "M5_offline_method_results.npz", allow_pickle=True)
names = [str(x) for x in pinned["method_names"]]
returns = pinned["mean_return"]
success = pinned["success_rate"]
fig, ax = plt.subplots(figsize=(9, 4.8))
positions = np.arange(len(names))
ax.bar(positions, returns)
ax.axhline(0.0, linewidth=1)
ax.set_xticks(positions, names, rotation=15, ha="right")
ax.set_ylabel(LABELS["mean_return"])
ax.set_title(LABELS["offline_title"])
for x, y, s in zip(positions, returns, success):
    ax.text(x, y + (0.08 if y >= 0 else -0.12), f"success={s:.2f}", ha="center", va="bottom" if y >= 0 else "top", fontsize=8)
ax.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

pinned = np.load(DATA / "M5_offline_method_results.npz", allow_pickle=True)
fqe_names = [str(x) for x in pinned["fqe_names"]]
fqe_values = pinned["fqe_values"]
true_by_name = dict(zip([str(x) for x in pinned["method_names"]], pinned["mean_return"]))
for name, value in zip(fqe_names, fqe_values):
    print(f"{name}: FQE={float(value): .4f}; direct simulator={float(true_by_name[name]): .4f}")

# A one-step target shows that evaluation averages under the fixed target policy;
# it does not maximize as Q-learning does.
policy_probs = np.array([[0.25, 0.75]])
next_values = np.array([[2.0, 4.0]])
print("one-step FQE target:", fqe_targets([-0.1], [False], next_values, policy_probs, gamma=0.9))
print(LABELS["fqe_warning"])

if PROFILE == "full":
    dt_episodes, _ = collect_offline_dataset(seed=41, episodes=120, profile="mixed")
    dt_epochs, context_steps, eval_eps = 12, 16, 6
else:
    dt_episodes, _ = collect_offline_dataset(seed=41, episodes=40, profile="mixed")
    dt_epochs, context_steps, eval_eps = 4, 12, 4

start = time.perf_counter()
with warnings.catch_warnings():
    warnings.filterwarnings("ignore", message="enable_nested_tensor")
    dt_model, dt_loss = train_decision_transformer(
        dt_episodes,
        seed=61,
        state_dim=KeyDoorGrid().observation_dim,
        action_dim=KeyDoorGrid().action_dim,
        context_steps=context_steps,
        epochs=dt_epochs,
        batch_size=96 if PROFILE == "full" else 64,
    )
smoke_dt_eval = evaluate_decision_transformer(
    KeyDoorGrid(seed=41), dt_model, target_return=0.7,
    episodes=eval_eps, seed=2000, context_steps=context_steps,
)
print(f"DT wall time: {time.perf_counter() - start:.2f}s")
print("current-profile evaluation:", smoke_dt_eval)

pinned = np.load(DATA / "M5_decision_transformer_results.npz")
fig, ax = plt.subplots(figsize=(8, 4.7))
ax.plot(pinned["target_return"], pinned["mean_return"], marker="o")
ax.plot(pinned["target_return"], pinned["target_return"], linestyle="--", label=LABELS["identity"])
ax.set_xlabel(LABELS["requested_return"])
ax.set_ylabel(LABELS["achieved_return"])
ax.set_title(LABELS["dt_title"])
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

start = time.perf_counter()
if PROFILE == "full":
    language_run = language_offline_experiment(seed=71, episodes_per_prompt=45, bc_epochs=12, iql_epochs=20)
else:
    language_run = language_offline_experiment(seed=71, episodes_per_prompt=20, bc_epochs=5, iql_epochs=8)
print(f"language offline-RL wall time: {time.perf_counter() - start:.2f}s")
print("current-profile dataset:", language_run["dataset"])
for name, row in language_run["evaluation"].items():
    print("\n", name, {"mean_reward": row["mean_reward"], "exact_match": row["exact_match"]})
    for prompt, seq, reward in zip(PROMPTS, row["sequences"], row["rewards"]):
        print("  ", prompt, "->", token_text(seq), "reward=", reward)

pinned = np.load(DATA / "M5_language_ilql_results.npz", allow_pickle=True)
fig, ax = plt.subplots(figsize=(8.8, 4.8))
names = [str(x) for x in pinned["method_names"]]
positions = np.arange(len(names))
ax.bar(positions, pinned["exact_match"])
ax.set_xticks(positions, names, rotation=12, ha="right")
ax.set_ylim(0, 1.08)
ax.set_ylabel(LABELS["exact_match"])
ax.set_title(LABELS["language_title"])
ax.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

protocol = __import__("json").loads((DATA / "M5_experiment_protocol.json").read_text(encoding="utf-8"))
print(LABELS["summary_heading"])
print("-", LABELS["summary_bc"])
print("-", LABELS["summary_dqn"])
print("-", LABELS["summary_offline"])
print("-", LABELS["summary_dt"])
print("-", LABELS["summary_llm"])
print("\nartifact scope:", protocol["artifact_scope"])