Module 11. RL for LLM agents: tools, environments, and long trajectories

def build_policy_mask(roles: list[str], lengths: list[int]) -> list[int]:
    allowed = {"system", "user", "assistant", "tool"}
    if len(roles) != len(lengths):
        raise ValueError("roles and lengths must have equal lengths")

    result: list[int] = []
    for role, length in zip(roles, lengths, strict=True):
        if role not in allowed:
            raise ValueError(f"unknown role: {role}")
        if isinstance(length, bool) or not isinstance(length, int) or length < 0:
            raise ValueError("lengths must be non-negative integers")
        result.extend([1 if role == "assistant" else 0] * length)
    return result

import math


def td_residuals(
    rewards: list[float],
    values: list[float],
    gamma: float,
) -> list[float]:
    if len(values) != len(rewards) + 1:
        raise ValueError("values must have length len(rewards)+1")
    if not math.isfinite(gamma) or not 0.0 <= gamma <= 1.0:
        raise ValueError("gamma must lie in [0, 1]")
    numbers = [*rewards, *values]
    if not all(math.isfinite(float(x)) for x in numbers):
        raise ValueError("inputs must be finite")

    return [
        float(rewards[t]) + gamma * float(values[t + 1]) - float(values[t])
        for t in range(len(rewards))
    ]

import math


def informative_group_probability(p: float, group_size: int) -> float:
    p = float(p)
    if not math.isfinite(p) or not 0.0 <= p <= 1.0:
        raise ValueError("p must lie in [0, 1]")
    if isinstance(group_size, bool) or not isinstance(group_size, int) or group_size < 1:
        raise ValueError("group_size must be a positive integer")
    if p in (0.0, 1.0):
        return 0.0

    at_least_one_success = -math.expm1(group_size * math.log1p(-p))
    all_success = math.exp(group_size * math.log(p))
    return max(0.0, min(1.0, at_least_one_success - all_success))

import math


def potential_shaping(
    rewards: list[float],
    potentials: list[float],
    gamma: float,
) -> list[float]:
    if len(potentials) != len(rewards) + 1:
        raise ValueError("potentials must have length len(rewards)+1")
    if not math.isfinite(gamma) or not 0.0 <= gamma <= 1.0:
        raise ValueError("gamma must lie in [0, 1]")
    if not all(math.isfinite(float(x)) for x in [*rewards, *potentials]):
        raise ValueError("inputs must be finite")

    return [
        float(rewards[t])
        + gamma * float(potentials[t + 1])
        - float(potentials[t])
        for t in range(len(rewards))
    ]

import math
from collections.abc import Mapping


def validate_rollout_record(record: dict) -> list[str]:
    required = {
        "task_id", "environment_id", "environment_version",
        "token_ids", "policy_mask", "behavior_logprobs",
        "policy_versions", "turn_ids",
        "action_ids", "observations", "rewards", "stop_reason",
        "terminated", "truncated", "infra_error",
    }
    errors: list[str] = []
    missing = sorted(required - set(record))
    if missing:
        return ["missing fields: " + ", ".join(missing)]

    token_ids = record["token_ids"]
    mask = record["policy_mask"]
    turns = record["turn_ids"]
    logprobs = record["behavior_logprobs"]

    if not all(isinstance(item, list) for item in (token_ids, mask, turns, logprobs)):
        return ["token_ids, policy_mask, turn_ids and behavior_logprobs must be lists"]
    if not (len(token_ids) == len(mask) == len(turns)):
        errors.append("token_ids, policy_mask and turn_ids must have equal lengths")
    if not all(x in (0, 1) and not isinstance(x, bool) for x in mask):
        errors.append("policy_mask must contain integer zeros and ones")
    if sum(x == 1 for x in mask) != len(logprobs):
        errors.append("logprob count must equal the number of policy tokens")
    try:
        valid_logprobs = all(math.isfinite(float(x)) and float(x) <= 0.0 for x in logprobs)
    except (TypeError, ValueError):
        valid_logprobs = False
    if not valid_logprobs:
        errors.append("logprobs must be finite and non-positive")

    versions = record["policy_versions"]
    if not isinstance(versions, Mapping):
        errors.append("policy_versions must be a mapping")
    elif len(token_ids) == len(mask) == len(turns):
        policy_turns = {str(t) for t, m in zip(turns, mask, strict=True) if m == 1}
        missing_turns = policy_turns - {str(t) for t in versions}
        if missing_turns:
            errors.append("missing versions for policy turns")

    actions = record["action_ids"]
    observations = record["observations"]
    rewards = record["rewards"]
    if not all(isinstance(item, list) for item in (actions, observations, rewards)):
        errors.append("action_ids, observations and rewards must be lists")
    elif not (len(actions) == len(rewards) and len(observations) == len(actions) + 1):
        errors.append("observations must have one more item than actions and rewards")

    status_count = sum(bool(record[name]) for name in ("terminated", "truncated", "infra_error"))
    if status_count > 1:
        errors.append("at most one terminal status may be true")
    if status_count == 0 and record["stop_reason"] is not None:
        errors.append("stop_reason requires a stopped rollout")
    if status_count == 1 and not isinstance(record["stop_reason"], str):
        errors.append("a stopped rollout must have a string stop_reason")
    return errors

def run_keydoor_episode(
    actions: list[str],
    *,
    key_color: str,
    door_color: str,
    max_steps: int = 6,
) -> dict:
    colors = {"red", "blue"}
    if key_color not in colors or door_color not in colors:
        raise ValueError("unknown color")
    if isinstance(max_steps, bool) or not isinstance(max_steps, int) or max_steps < 1:
        raise ValueError("max_steps must be positive")

    inspected = False
    carrying = None
    opened: set[str] = set()
    invalid = 0
    terminated = False
    success = False
    steps = 0

    for action in actions[:max_steps]:
        steps += 1
        valid = True
        if action == "inspect":
            inspected = True
        elif action in {"pickup_red", "pickup_blue"}:
            color = action.removeprefix("pickup_")
            if not inspected or color != key_color:
                valid = False
            else:
                carrying = color
        elif action in {"open_red", "open_blue"}:
            color = action.removeprefix("open_")
            if carrying != color:
                valid = False
            else:
                opened.add(color)
        elif action == "finish":
            terminated = True
            success = door_color in opened
        else:
            valid = False

        if not valid:
            invalid += 1
        if terminated:
            break

    truncated = not terminated and steps >= max_steps
    return {
        "reward": 1.0 if success else 0.0,
        "success": success,
        "terminated": terminated,
        "truncated": truncated,
        "steps": steps,
        "invalid_actions": invalid,
        "carrying": carrying,
        "opened": sorted(opened),
    }

assert run_keydoor_episode(
    ["inspect", "pickup_red", "open_red", "finish"],
    key_color="red", door_color="red",
)["success"]

assert run_keydoor_episode(
    ["finish"], key_color="red", door_color="red",
)["terminated"]

assert run_keydoor_episode(
    ["inspect", "pickup_blue", "finish"],
    key_color="red", door_color="red",
)["invalid_actions"] == 1

assert run_keydoor_episode(
    ["inspect"] * 3,
    key_color="blue", door_color="blue", max_steps=3,
)["truncated"]