RL for LLM5 / 13

Module 5. Learning from fixed data: imitation, deep Q-methods, and offline RL

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 382.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Module 5 lab — imitation, DQN, and learning from fixed data

Main module notebook

Download .ipynb module-05-imitation-offline-rl.en.ipynb · 21.6 KB

Show notebook text

Module 5 lab — imitation, DQN, and learning from fixed data

This notebook connects the lecture equations to one continuing KeyDoorGrid world and a small token-level pipeline. It contains complete collection, training, evaluation, and diagnostic loops rather than isolated formulas. Numerical results reused in the lecture are identified as course-generated artifacts, not external benchmarks.

The default smoke profile executes reduced CPU versions of every path and loads pinned full-profile artifacts for stable plots. Set RL_LLM_M5_PROFILE=full before starting the kernel to rebuild the longer runs.

Python

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: key, door, goal, traps, and two stopping semantics",
    "grid_x": "column",
    "grid_y": "row",
    "dagger_x": "data-aggregation round",
    "success": "successful-episode fraction",
    "dagger_title": "Pinned run: BC followed by DAgger rounds",
    "episode": "training episode",
    "rolling_success": "rolling success fraction",
    "dqn_title": "DQN and Double DQN: one pinned seed, window 15",
    "count": "transition count",
    "log_title": "Action composition of the fixed log",
    "mean_return": "mean return",
    "offline_title": "One log, several protection mechanisms",
    "fqe_warning": "FQE is another learned model; calibration does not follow from the method name.",
    "identity": "achieved = requested",
    "requested_return": "requested return-to-go",
    "achieved_return": "observed mean return",
    "dt_title": "A Decision Transformer condition is a request, not a guarantee",
    "exact_match": "exact-answer fraction",
    "language_title": "Synthetic token-level offline pipeline",
    "summary_heading": "Lab takeaways",
    "summary_bc": "BC is stable on demonstration-distributed states; DAgger shows how to label states induced by the learner.",
    "summary_dqn": "Replay and a target network stabilize online Q-learning, while Double DQN separates selection from evaluation.",
    "summary_offline": "A fixed log turns coverage into an identifiability constraint; pessimism and FQE must themselves be validated.",
    "summary_dt": "Decision Transformer trains as conditional sequence modeling, but return-to-go is not a performance promise.",
    "summary_llm": "Token-level Q/V learning is possible, yet the reduced ILQL-like implementation is not an LLM-scale reproduction.",
}

1. Begin with environment and transition contracts

KeyDoorGrid is small without being a three-state toy: collecting the key changes door accessibility, traps are true terminals, and the time limit is an external truncation. The same environment can therefore test bootstrap masks while supporting imitation, online DQN, and offline RL.

Python

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]))

2. Behavioral cloning and DAgger

The first policy sees only a few expert trajectories. DAgger then asks the expert to label states induced by that learned policy. The plot uses the pinned full protocol, while the current profile executes an actual training loop separately.

Python

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()

3. DQN and Double DQN

Both implementations use replay, a target network, $\varepsilon$ -greedy collection, and separate termination/truncation flags. Sparse reward motivates an explicit replay warm-start with noisy expert-guided trajectories. That is a documented course engineering choice, not a hidden property of vanilla DQN.

Python

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()

4. Inspect the fixed log before choosing an algorithm

Action counts and reward composition usually reveal more than an algorithm label. The support mask below illustrates one explicit contract: an action is eligible when its behavior probability is not too small relative to the most likely action.

Python

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()

5. One log, several ways to distrust extrapolation

Naive offline DQN, action restriction, CQL, and IQL protect different objects. Equal outcomes do not make the methods identical, and one favorable seed is not a global ranking.

Python

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()

6. FQE is also a model that requires validation

FQE replaces a maximizing continuation with an expectation under a fixed evaluation policy. That changes improvement into evaluation, but it does not remove approximation or coverage error.

Python

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"])

7. Decision Transformer: a return condition inside a causal sequence

A compact causal transformer is trained on fixed trajectories. The curve compares achieved return across initial return-to-go requests; the diagonal is a reference, not a theorem the model is expected to satisfy.

Python

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()

8. An ILQL-like token pipeline

The last environment is linguistic: a prefix is a state, the next token is an action, and reward arrives after the response ends. We train a behavior policy, token-level $Q/V$ estimates, and an advantage-weighted policy. The mechanics are real, but the vocabulary and model are deliberately small and do not reproduce full-scale ILQL.

Python

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()

9. What this notebook establishes

The important result is seeing where the statistical regime changes. Demonstrations, fresh interaction, and an immutable log support different claims. The final cell prints the interpretation boundaries stored with the experiment protocol.

Python

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"])