Module 11. RL for LLM agents: tools, environments, and long trajectories
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 KBThe archive preserves the source folder structure, both languages, and supporting files.
Module 11 Lab: Agent Environments, Turn-Level Credit, and Long Trajectories
Main module notebook
Download .ipynb module-11-agentic-rl.en.ipynb · 23.8 KB
Show notebook text
Module 11 Lab: Agent Environments, Turn-Level Credit, and Long Trajectories
This notebook connects the lecture’s accounting identities to a small but complete partially observed environment. The agent must inspect a room, retain a one-turn clue, choose a key, open the correct door, and finish the task. The computational route uses a recurrent actor–critic, behavior log-probabilities, policy versions, true terminals, and collector truncations.
The results belong only to the course’s synthetic environment. They are not an LLM benchmark or a claim about the universal performance of a contemporary Agentic RL method.
1. Reproducible profiles
The smoke profile runs on CPU and exercises the full pipeline. The full profile increases the number of seeds, updates, and evaluation episodes. The optional open-model branch is disabled by default.
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:
"""Locate the course root independently of the notebook working directory."""
start = start.resolve()
for candidate in (start, *start.parents):
if (candidate / "modules/module-11/module_11_reference.py").is_file():
return candidate
raise FileNotFoundError("Could not locate the RL for LLM repository root")
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 must be smoke, full, or 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"Course root: {ROOT}")
print(f"Profile: {PROFILE}; seeds={SEEDS}; updates={UPDATES}; episodes/update={EPISODES_PER_UPDATE}")
2. The environment is a contract, not scenery
The environment keeps the target door color hidden. The clue_red or clue_blue observation appears only after inspect; the recurrent policy must then retain that clue. The base reward is issued only after successful completion, while the potential term provides a controlled learning signal.
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": observation,
"action": action,
"next observation": next_observation,
"base reward": base_reward,
"potential before": info["potential_before"],
"potential after": info["potential_after"],
"shaped reward": info["shaped_reward"],
"terminated": terminated,
"truncated": truncated,
})
observation = next_observation
trace_df = pd.DataFrame(rows)
display(trace_df)
print("Sum of base rewards:", trace_df["base reward"].sum())
print("Discounted return of shaped rewards:", discounted_return(trace_df["shaped reward"], 0.99))
3. A policy mask must follow token provenance
Tool outputs can be much longer than model turns. Averaging a policy loss over the entire transcript lets environment tokens dilute the policy gradient. In a simple chat template, the assistant role is a useful approximation; once a controller inserts or rewrites content, explicit provenance is required.
segments = [
("system", 44),
("user", 63),
("assistant", 18),
("tool", 162),
("assistant", 21),
]
report = trajectory_mask_report(segments)
mask_table = pd.DataFrame(segments, columns=["role", "token count"])
mask_table["included in policy loss"] = mask_table["role"].eq("assistant")
display(mask_table)
print(report)
print(
f"If the loss is incorrectly averaged over all {report.total_tokens} tokens, "
f"the signal from {report.policy_tokens} policy tokens is diluted by a factor of {report.dilution_factor:.2f}."
)
4. True termination versus collector truncation
The visible end of a rollout buffer can represent two different mathematical situations. A true terminal has no continuation under the task definition. An external collector cut leaves the process unfinished, so the critic target normally retains a bootstrap value from the final observation.
gamma = 0.9
boundary_table = pd.DataFrame([
{
"case": "true terminal",
"reward": 0.0,
"V(next)": 0.7,
"target": collector_target(0.0, 0.7, gamma=gamma, terminated=True),
},
{
"case": "external collector truncation",
"reward": 0.0,
"V(next)": 0.7,
"target": 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({
"observation": final_observation,
"reward": reward,
"terminated": terminated,
"truncated": truncated,
"stop_reason": info["stop_reason"],
})
5. From tokens to turns and semi-Markov actions
The log-probability of the sampled action string is the sum of its token log-probabilities. Several strings may nevertheless parse to one semantic action. Tool operations also have different durations, which produces semi-Markov discounting when elapsed time is part of the objective.
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([
{"quantity": "log P(generated string)", "value": macro_action_logprob(behavior_logprobs)},
{"quantity": "turn-level current/behavior ratio", "value": turn_probability_ratio(current_logprobs, behavior_logprobs)},
{"quantity": "P(semantic action inspect)", "value": semantic_action_probability(string_probabilities, parser, "inspect")},
{"quantity": "semi-MDP return with durations 2 and 3", "value": semi_markov_return([0.0, 1.0], [2, 3], 0.9)},
])
display(macro_table)
6. TD telescoping and potential-based shaping
Turn-level TD residuals must remain consistent with the trajectory return and the final boundary value. Potential-based shaping redistributes reward across transitions; under the stated assumptions, it changes the return only by a predictable boundary term.
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"Telescoping check: {left:.6f} = {right:.6f}")
print(f"Original return={original_return:.6f}; shaped return={shaped_return:.6f}; identity={predicted_return:.6f}")
7. When a relative group contains any signal
With a binary verifier, a group supports relative normalization only when it contains both successes and failures. Extremely easy and extremely hard prompts create homogeneous groups; a larger learning rate cannot manufacture missing variation.
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("Single-rollout success probability p")
ax.set_ylabel("Probability of a mixed group")
ax.set_title("A group is informative only when it contains successes and failures")
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)
8. A recurrent actor–critic in a partially observed environment
We train the same architecture with a sparse terminal reward and with potential-based shaping. The plot reports the median and full range across several seeds rather than a single favorable run. This is a local diagnostic of the environment, not a universal superiority claim for shaping.
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)
fig, ax = plt.subplots(figsize=(8.5, 4.8))
for reward_mode, label in (("sparse", "Sparse terminal reward"), ("potential", "Potential-based shaping")):
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("Policy update (-1 is the pre-training evaluation)")
ax.set_ylabel("Greedy episode success rate")
ax.set_ylim(-0.03, 1.03)
ax.set_title("Synthetic KeyDoorAgentEnv CPU experiment")
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"])
print("Best configuration for the next diagnostic:", best_key, best_metrics)
9. The rollout log as a reproducible object
Text and a final reward are not enough. Training also needs actions, observations, a policy mask, behavior log-probabilities, policy versions, and a stop reason. The next cell builds a valid record from a fresh rollout and then deliberately breaks it.
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([
{
"turn": index,
"observation": transition.observation,
"action": transition.action,
"next_observation": transition.next_observation,
"behavior_logprob": transition.behavior_logprob,
"base_reward": transition.base_reward,
"shaped_reward": transition.shaped_reward,
"terminated": transition.terminated,
"truncated": transition.truncated,
}
for index, transition in enumerate(transitions)
])
display(trace)
print("Rollout report:", rollout_report)
contract = rollout_contract_from_transitions(
transitions,
task_id="m11-keydoor-red",
environment_version="1.0",
policy_version="m11-notebook-policy",
)
print("Violations in the valid contract:", validate_rollout_record(contract))
broken_contract = dict(contract)
broken_contract["policy_versions"] = {}
print("Violations after removing the policy version:", validate_rollout_record(broken_contract))
10. Optional branch: a stateful TRL environment
As of the factual snapshot, TRL’s official interface accepts an environment_factory in GRPOTrainer and constructs a separate environment instance for each completion. The API is experimental and mutable, so a real run must pin the resolved library versions and re-check the documentation. This branch was not executed in the release gate.
if not RUN_REAL_MODEL:
print(
"The optional branch is disabled. To run it, create an isolated environment, "
"pin the resolved trl/transformers/datasets versions, and set RUN_REAL_MODEL = True."
)
else:
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer
class TextKeyDoorEnv:
"""A minimal stateful environment for TRL’s official environment_factory contract."""
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("The trainer was constructed. trainer.train() is intentionally not called automatically.")