Learn AI Series (#115) - Offline Reinforcement Learning

avatar

Learn AI Series (#115) - Offline Reinforcement Learning

variant-c-11-teal.png

What will I learn

  • You will learn why learning from a fixed, frozen dataset -- with no environment to poke at -- is both wildly useful and genuinely hard;
  • distribution shift, the one problem that sits underneath every offline RL headache;
  • Conservative Q-Learning (CQL) -- buying safety with a dose of pessimism about actions you've never tried;
  • Implicit Q-Learning (IQL) -- dodging the whole overestimation trap by never asking the awkward question in the first place;
  • Decision Transformer -- the moment RL quietly turned into sequence modelling and started borrowing GPT's homework (episodes #52-53, #58);
  • and where all of this actually matters -- hospitals, self-driving fleets, robots -- plus the bigger picture: offline RL is the bridge from "agent that learns by trial and error" to "agent pre-trained like a foundation model".

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.11+) distribution with NumPy and PyTorch;
  • You've followed the RL arc -- especially episodes #107 (DQN), #109 (advanced policy optimisation) and the transformer two-parter #52-53 -- because today leans on all three at once, and it lands far better once you've met each one properly.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#115) - Offline Reinforcement Learning

Solutions to Episode #114 Exercises

Before we freeze the environment solid, let's clear last episode's three exercises. They build on the LinearIRL, MaxEntIRL and BehavioralCloning classes from episode #114, so I'm assuming those are imported and in scope.

Exercise 1: the tiniest possible IRL demo on a 5x5 gridworld -- hand-pick a true reward, solve it with value iteration (episode #104) to get an expert, roll out trajectories, then throw the true reward away and recover it with the projection loop.

The trick is keeping the gridworld small enough that value iteration is basically instant, so the whole thing runs in a fraction of a second:

import numpy as np
# Assumes LinearIRL from episode #114.

N = 5
N_STATES = N * N
ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]   # up, down, left, right
GAMMA = 0.9
GOAL = (4, 4)


def step(s, a):
    r, c = divmod(s, N)
    dr, dc = ACTIONS[a]
    nr, nc = np.clip(r + dr, 0, N - 1), np.clip(c + dc, 0, N - 1)
    return nr * N + nc


def value_iteration(reward, iters=200):
    V = np.zeros(N_STATES)
    for _ in range(iters):
        Q = np.array([[reward[step(s, a)] + GAMMA * V[step(s, a)]
                       for a in range(4)] for s in range(N_STATES)])
        V = Q.max(axis=1)
    return Q.argmax(axis=1)                     # greedy policy per state


def rollout(policy, start=0, length=12):
    s, traj = start, []
    for _ in range(length):
        a = policy[s]
        traj.append((s, a))
        s = step(s, a)
    return traj


# --- true reward: +1 in the goal cell, 0 elsewhere ---
true_reward = np.zeros(N_STATES)
true_reward[GOAL[0] * N + GOAL[1]] = 1.0
expert_policy = value_iteration(true_reward)
expert_trajs = [rollout(expert_policy, start=s) for s in (0, 2, 10, 20)]

# --- one-hot cell features, and an RL solver the projection loop can call ---
def feature_fn(s, a):
    phi = np.zeros(N_STATES)
    phi[step(s, a)] = 1.0                        # feature = which cell you land in
    return phi


class GridSolver:
    def solve(self, reward_fn):
        r_vec = np.array([reward_fn(s, 0) for s in range(N_STATES)])
        return value_iteration(r_vec)
    def generate_trajectories(self, policy):
        return [rollout(policy, start=s) for s in (0, 2, 10, 20)]


irl = LinearIRL(n_features=N_STATES, gamma=GAMMA)
w = irl.learn_reward(expert_trajs, feature_fn, GridSolver(), n_iterations=30)

recovered = w.reshape(N, N)
print("recovered reward (rounded):")
print(np.round(recovered - recovered.min(), 2))
print("expert action in each state matches recovered-policy?",
      np.array_equal(expert_policy, value_iteration(w)))

The heatmap (I'll leave the matplotlib call to you) will not be a clean copy of the true reward -- the recovered numbers slope gently up toward the goal in stead of putting all the mass in one corner. And that's exactly the point episode #114 hammered on: reward ambiguity means you don't get the true reward back, you get a reward that produces the same behaviour. The final array_equal check is the one that matters -- same policy, different numbers.

Exercise 2: feel the compounding-error cliff. Train an expert on CartPole-v1, clone it, then nudge the start state past anything the expert ever saw and watch behavioral cloning fall apart.

I'll keep the "expert" honest but cheap -- a short policy-gradient train from episode #108 is plenty for CartPole, and the whole lesson lives in the two evaluate numbers at the end:

import gymnasium as gym
import numpy as np
import torch
# Assumes BehavioralCloning from episode #114 and a trained `expert` policy.

def collect(expert, n_episodes=50):
    env = gym.make("CartPole-v1")
    states, actions = [], []
    for _ in range(n_episodes):
        obs, _ = env.reset()
        done = False
        while not done:
            a = expert.act(obs)                  # expert picks the action
            states.append(obs)
            actions.append(a)
            obs, _, term, trunc, _ = env.step(a)
            done = term or trunc
    return np.array(states, dtype=np.float32), np.array(actions)


def evaluate(policy_fn, n=50, jolt=0.0):
    env = gym.make("CartPole-v1")
    lengths = []
    for _ in range(n):
        obs, _ = env.reset()
        obs = obs + np.random.uniform(-jolt, jolt, size=obs.shape)  # push off-manifold
        done, steps = False, 0
        while not done and steps < 500:
            obs, _, term, trunc, _ = env.step(policy_fn(obs))
            done, steps = term or trunc, steps + 1
        lengths.append(steps)
    return np.mean(lengths)


states, actions = collect(expert, n_episodes=50)
bc = BehavioralCloning(state_dim=4, action_dim=2)
opt = torch.optim.Adam(bc.parameters(), lr=1e-3)
for epoch in range(300):
    logits = bc(torch.FloatTensor(states))
    loss = torch.nn.functional.cross_entropy(logits, torch.LongTensor(actions))
    opt.zero_grad(); loss.backward(); opt.step()

bc_act = lambda o: bc(torch.FloatTensor(o)).argmax().item()
print("BC, normal starts :", evaluate(bc_act, jolt=0.0))
print("BC, jolted starts :", evaluate(bc_act, jolt=0.15))

On calm starts the cloned policy looks great -- it happily balances for hundreds of steps because those states look exactly like the demonstrations. Jolt the start state past what the expert ever visited, though, and the average length sags: the agent has genuinely never seen "leaning this far" and has no clue how to recover, so one mistake snowballs into a fall. Nota bene, that's not a bug in the code -- it's the whole failure mode of imitation without a reward, and it's precisely why we spent episode #114 learning to recover the reward in stead of memorising actions.

Exercise 3: watch MaxEntIRL converge on the gridworld from Exercise 1, then sabotage it by only ever showing the expert half the grid.

import numpy as np, torch
# Assumes MaxEntIRL from episode #114, plus N, N_STATES, expert_policy, rollout, step.

def onehot(s):
    v = np.zeros(N_STATES); v[s] = 1.0; return v

def policy_states(policy, starts):                # cells the current policy visits
    return [onehot(s) for st in starts for (s, _) in rollout(policy, start=st)]

def greedy_from_reward(reward_net):
    r = np.array([reward_net(torch.FloatTensor(onehot(s))).item()
                  for s in range(N_STATES)])
    from_ex1 = value_iteration(r)                 # reuse Exercise 1's solver
    return from_ex1, r

irl = MaxEntIRL(state_dim=N_STATES)
all_starts = list(range(N_STATES))
expert_visited = [onehot(s) for st in all_starts for (s, _) in rollout(expert_policy, st)]

for rnd in range(40):
    policy, r_now = greedy_from_reward(irl.reward_net)
    irl.update(expert_visited, policy_states(policy, all_starts))
    if rnd % 8 == 0:
        corr = np.corrcoef(r_now, true_reward)[0, 1]
        print(f"round {rnd}: corr(recovered, true) = {corr:.3f}")

# --- sabotage: expert only ever seen in the top half of the grid ---
half = [s for s in range(N_STATES) if s // N < N // 2]
expert_half = [onehot(s) for st in half for (s, _) in rollout(expert_policy, st)]
for _ in range(40):
    policy, _ = greedy_from_reward(irl.reward_net)
    irl.update(expert_half, policy_states(policy, half))
r_final = np.array([irl.reward_net(torch.FloatTensor(onehot(s))).item()
                    for s in range(N_STATES)]).reshape(N, N)
print("recovered reward, bottom (unvisited) half is unconstrained garbage:")
print(np.round(r_final, 2))

With the full grid, the correlation between the recovered and true reward climbs round after round (up to a constant offset and scale -- adding a constant to every reward changes nothing, remember). Starve the algorithm of the bottom half and the recovered reward there goes wild: the demonstrations simply never said what happens down there, so MaxEnt has nothing to pin it to. That's underdetermination again -- and hold onto that feeling, because "the data never covered this region" is about to become the entire story of today's episode.

The setup nobody warned you about

For thirteen episodes of RL we've quietly assumed one luxury: whenever the agent wants to learn something, it can just act and see what happens. Collect experience, learn from it, collect better experience, learn more. This works brilliantly in a simulator where you can run millions of episodes overnight and nobody gets hurt.

But step outside the simulator and that luxury evaporates. A hospital has years of patient records -- observations, the treatments doctors chose, the outcomes -- and would very much like a better treatment policy. It cannot "explore" by trying random treatments on real patients to see what happens (for hopefully obvious reasons). A company has logs of millions of customer interactions and wants a better recommendation policy -- but it can't A/B-test bad recommendations on real users just to fill in the gaps in its data.

Offline RL (also called batch RL) is RL under exactly that constraint: learn entirely from a fixed dataset of previously-collected interactions. No environment access. No exploration. No new data. Just a static log and the challenge of squeezing the best possible policy out of it. Having said that, this innocent-sounding change breaks nearly everything we've built -- and understanding why is the fastest way to understand modern RL's flirtation with foundation models (episode #58).

Why offline RL is hard: distribution shift

Here's the whole problem in one sentence. The dataset was collected by some behaviour policy we'll call beta. Your learned policy pi is going to visit different states and pick different actions than beta ever did -- but you can only trust your value estimates for the state-action pairs that actually appear in the data. Everything else is a guess dressed up as knowledge.

Let's watch the naive approach detonate. Just run ordinary DQN (episode #107) on the fixed dataset, no changes:

import torch
import torch.nn as nn
import numpy as np

class NaiveOfflineDQN:
    """Naive approach: just run DQN on the fixed dataset. (This fails.)"""
    def __init__(self, state_dim, n_actions, lr=1e-3, gamma=0.99):
        self.gamma = gamma
        self.q_net = nn.Sequential(
            nn.Linear(state_dim, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, n_actions),
        )
        self.optimizer = torch.optim.Adam(self.q_net.parameters(), lr=lr)

    def train_step(self, batch):
        states, actions, rewards, next_states, dones = batch

        # Current Q-values
        q_values = self.q_net(states).gather(1, actions.unsqueeze(1)).squeeze()

        # Target Q-values
        with torch.no_grad():
            next_q = self.q_net(next_states).max(dim=1)[0]
            targets = rewards + self.gamma * next_q * (1 - dones)

        loss = nn.functional.mse_loss(q_values, targets)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss.item()

This fails catastrophically, and the culprit is that innocent-looking max. Q-Learning's target uses the highest Q-value in the next state. But for actions that never appeared in the dataset, the Q-network was never trained on them -- its outputs there are essentially random. And random numbers from an untrained network are, purely by chance, sometimes very high. The max operation seeks out exactly those erroneously high values and bakes them into the targets, which propagate through training, inflating the Q-estimates for out-of-distribution actions to absurd levels.

Then the learned policy, being a good little maximiser, happily picks those out-of-distribution actions because they now have the highest (hallucinated) Q-values. The result is behaviour completely unlike anything in the training data -- and, almost always, completely terrible. In online RL this self-correcting: try the bad action, get a bad reward, fix the estimate. Offline, there's no "try" -- the delusion never gets corrected. That, in one paragraph, is why offline RL needed its own toolbox.

Conservative Q-Learning (CQL)

CQL (Kumar et al., 2020) fixes the overestimation with a beautifully blunt idea: be pessimistic. Add a penalty that pushes Q-values down for actions you haven't actually seen, while keeping the estimates honest for actions that are in the data.

class CQLAgent:
    """Conservative Q-Learning: pessimistic about unseen actions."""
    def __init__(self, state_dim, n_actions, lr=3e-4, gamma=0.99,
                 cql_alpha=1.0):
        self.gamma = gamma
        self.n_actions = n_actions
        self.cql_alpha = cql_alpha  # conservatism weight

        self.q_net = nn.Sequential(
            nn.Linear(state_dim, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, n_actions),
        )
        self.target_net = nn.Sequential(
            nn.Linear(state_dim, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, n_actions),
        )
        self.target_net.load_state_dict(self.q_net.state_dict())
        self.optimizer = torch.optim.Adam(self.q_net.parameters(), lr=lr)

    def train_step(self, batch):
        states, actions, rewards, next_states, dones = batch

        # Standard Bellman loss (same as DQN)
        q_values = self.q_net(states)
        q_selected = q_values.gather(1, actions.unsqueeze(1)).squeeze()

        with torch.no_grad():
            next_q = self.target_net(next_states).max(dim=1)[0]
            targets = rewards + self.gamma * next_q * (1 - dones)

        bellman_loss = nn.functional.mse_loss(q_selected, targets)

        # CQL regularizer: push down Q-values for ALL actions,
        # pull UP Q-values for actions IN the dataset
        logsumexp_q = torch.logsumexp(q_values, dim=1).mean()
        dataset_q = q_selected.mean()
        cql_loss = logsumexp_q - dataset_q

        # Total loss
        loss = bellman_loss + self.cql_alpha * cql_loss

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return loss.item(), bellman_loss.item(), cql_loss.item()

The penalty logsumexp(Q) - Q(data_actions) is doing two jobs at once:

  1. logsumexp(Q) is a soft max -- it presses down on Q-values across all actions, and hardest on the biggest ones (exactly the delusions we're worried about);
  2. -Q(data_actions) shoves back up specifically for the actions that really occurred in the dataset.

Net effect: in-distribution Q-values stay sensible, out-of-distribution Q-values get squashed. The policy turns conservative -- it would rather stick with behaviour it has evidence for than gamble on an action whose outcome is a shrug. That cql_alpha knob is the dial between "trust the data blindly" and "trust nothing" -- and tuning it is most of the practical work with CQL ;-)

Implicit Q-Learning (IQL)

IQL (Kostrikov et al., 2022) is my favourite kind of fix -- it doesn't wrestle the overestimation problem, it simply refuses to create it. The whole trouble came from querying max Q(s', a') over actions that might not be in the data. So IQL just... never does that. In stead it learns a separate value network V(s) via expectile regression and uses that:

class IQLAgent:
    """Implicit Q-Learning: avoids querying OOD actions entirely."""
    def __init__(self, state_dim, n_actions, lr=3e-4, gamma=0.99, tau=0.7):
        self.gamma = gamma
        self.tau = tau  # expectile (>0.5 focuses on higher values)

        # Q-network
        self.q_net = nn.Sequential(
            nn.Linear(state_dim, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, n_actions),
        )

        # Value network: V(s) approx expectile of Q(s,a) under data distribution
        self.v_net = nn.Sequential(
            nn.Linear(state_dim, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, 1),
        )

        self.q_optimizer = torch.optim.Adam(self.q_net.parameters(), lr=lr)
        self.v_optimizer = torch.optim.Adam(self.v_net.parameters(), lr=lr)

    def expectile_loss(self, diff, tau):
        """Asymmetric L2 loss: weight positive errors more (tau > 0.5)."""
        weight = torch.where(diff > 0, tau, 1 - tau)
        return (weight * diff.pow(2)).mean()

    def train_step(self, batch):
        states, actions, rewards, next_states, dones = batch

        # Value update: V(s) trained with expectile regression on Q(s,a)
        with torch.no_grad():
            q_values = self.q_net(states).gather(1, actions.unsqueeze(1)).squeeze()

        v_pred = self.v_net(states).squeeze()
        v_loss = self.expectile_loss(q_values - v_pred, self.tau)

        self.v_optimizer.zero_grad()
        v_loss.backward()
        self.v_optimizer.step()

        # Q update: use V(s') instead of max Q(s', a')
        q_pred = self.q_net(states).gather(1, actions.unsqueeze(1)).squeeze()
        with torch.no_grad():
            next_v = self.v_net(next_states).squeeze()
            targets = rewards + self.gamma * next_v * (1 - dones)

        q_loss = nn.functional.mse_loss(q_pred, targets)

        self.q_optimizer.zero_grad()
        q_loss.backward()
        self.q_optimizer.step()

        return q_loss.item(), v_loss.item()

The clever bit is that tau > 0.5 in the expectile loss. A plain mean regression (tau = 0.5) would make V(s) the average Q-value under the data. Crank tau up toward 1 and it leans toward the upper end -- it approximates the maximum Q-value among actions that actually appear in the dataset, without ever evaluating Q on an action that doesn't. No out-of-distribution query, no hallucinated maximum, no overestimation. It's the difference between "what's the best I could imagine doing here?" and "what's the best I've got receipts for?" -- and offline, receipts are all you've got.

Decision Transformer: RL as sequence modelling

Now for the plot twist that reorganised the whole field. Decision Transformer (Chen et al., 2021) looks at offline RL, notices it's "learn from a big static dataset", notices that's exactly what language models do, and asks the obvious-in-hindsight question: why not just use a transformer (episodes #52-53)?

No value function. No Bellman equation. No bootstrapping. Just: model the sequence of returns, states and actions, and predict the next action.

class DecisionTransformer(nn.Module):
    """RL as sequence modeling: predict actions conditioned on desired return."""
    def __init__(self, state_dim, action_dim, hidden=128, n_heads=4,
                 n_layers=3, max_timestep=1000, max_ep_len=200):
        super().__init__()
        self.hidden = hidden

        # Embeddings for returns, states, actions, timesteps
        self.return_embed = nn.Linear(1, hidden)
        self.state_embed = nn.Linear(state_dim, hidden)
        self.action_embed = nn.Linear(action_dim, hidden)
        self.timestep_embed = nn.Embedding(max_timestep, hidden)

        # Layer norm
        self.ln = nn.LayerNorm(hidden)

        # GPT-style transformer decoder
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden, nhead=n_heads, batch_first=True
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, n_layers)

        # Output: predict action
        self.action_head = nn.Linear(hidden, action_dim)

    def forward(self, returns_to_go, states, actions, timesteps):
        """
        Input sequence: (R_1, s_1, a_1, R_2, s_2, a_2, ...)
        Predict: a_t given (R_1, s_1, a_1, ..., R_t, s_t)
        """
        batch_size, seq_len = states.shape[:2]

        # Embed each modality
        return_embeds = self.return_embed(returns_to_go.unsqueeze(-1))
        state_embeds = self.state_embed(states)
        action_embeds = self.action_embed(actions)
        time_embeds = self.timestep_embed(timesteps)

        # Add timestep embeddings
        return_embeds = return_embeds + time_embeds
        state_embeds = state_embeds + time_embeds
        action_embeds = action_embeds + time_embeds

        # Interleave: R_1, s_1, a_1, R_2, s_2, a_2, ...
        stacked = torch.stack(
            [return_embeds, state_embeds, action_embeds], dim=2
        ).reshape(batch_size, 3 * seq_len, self.hidden)

        stacked = self.ln(stacked)

        # Causal transformer (can only attend to past tokens)
        mask = nn.Transformer.generate_square_subsequent_mask(3 * seq_len)
        output = self.transformer(stacked, stacked, tgt_mask=mask)

        # Extract action predictions (every 3rd token, offset by 1)
        action_preds = output[:, 1::3, :]  # state positions predict actions
        return self.action_head(action_preds)

The genuinely mind-bending part is what happens at inference. You feed the model a desired return -- the "returns-to-go" token. Want a total reward of 100? You literally tell it 100, and it generates the actions that, based on everything it saw in training, are consistent with hitting that number. Ask for a high return, get expert-like behaviour; ask for a low one, get deliberately mediocre behaviour. The reward stops being something you optimise and becomes something you condition on, like a prompt (episode #62).

@torch.no_grad()
def generate_action(model, target_return, states, actions, timesteps):
    """At test time: condition on the return you WANT, read off the action."""
    rtg = torch.full((1, states.shape[1]), float(target_return))
    preds = model(rtg, states, actions, timesteps)
    return preds[:, -1]          # the action for the most recent state

Training it is almost anticlimactic after all the Bellman machinery we've built -- it's just supervised sequence learning, the exact same loop that trains a GPT (episode #58), only predicting actions in stead of tokens:

def train_decision_transformer(model, dataset, n_epochs=100,
                                batch_size=64, context_len=20, lr=1e-4):
    """Train DT on an offline dataset of trajectories."""
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)

    for epoch in range(n_epochs):
        total_loss = 0
        n_batches = 0

        for batch in sample_batches(dataset, batch_size, context_len):
            returns_to_go, states, actions, timesteps = batch

            # Predict actions
            pred_actions = model(returns_to_go, states, actions, timesteps)

            # Loss: MSE between predicted and actual actions
            # Only on non-padded positions
            loss = nn.functional.mse_loss(pred_actions[:, :-1], actions[:, 1:])

            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 0.25)
            optimizer.step()

            total_loss += loss.item()
            n_batches += 1

        print(f"Epoch {epoch}: loss = {total_loss / n_batches:.4f}")


def sample_batches(dataset, batch_size, context_len):
    """Sample random subsequences from trajectories."""
    trajectories = dataset['trajectories']
    for _ in range(len(trajectories) // batch_size):
        batch_returns, batch_states, batch_actions, batch_times = [], [], [], []

        for _ in range(batch_size):
            traj = trajectories[np.random.randint(len(trajectories))]
            start = np.random.randint(max(1, len(traj['states']) - context_len))
            end = start + context_len

            batch_states.append(traj['states'][start:end])
            batch_actions.append(traj['actions'][start:end])
            batch_times.append(np.arange(start, end))

            # Compute returns-to-go
            rewards = traj['rewards'][start:end]
            rtg = np.cumsum(rewards[::-1])[::-1]
            batch_returns.append(rtg)

        yield (torch.FloatTensor(np.array(batch_returns)),
               torch.FloatTensor(np.array(batch_states)),
               torch.FloatTensor(np.array(batch_actions)),
               torch.LongTensor(np.array(batch_times)))

And -- this still slightly amazes me -- it works, often beating the careful value-based methods on diverse offline datasets. No pessimism penalty, no expectile trick, no target networks. Just "predict the next action given what a high-return trajectory tends to look like". Sometimes the boring, general tool wins.

Where offline RL actually earns its keep

Healthcare: treatment decisions from historical patient records. You cannot randomise cancer treatments to "explore", full stop. Offline RL learns from the outcomes already on file while respecting that each decision reshapes the future options -- something plain supervised learning on "what the doctor did" (that's behavioral cloning from episode #114, and it inherits all of BC's brittleness) doesn't capture.

Autonomous driving: millions of hours of logged human driving. Offline RL extracts good policies without the rather severe safety cost of online exploration on public roads.

Robotics: teleoperation demonstrations. A human drives the robot, generating a dataset of successful (and gloriously unsuccessful) manipulation trajectories, and offline RL learns a policy that can end up better than the average demonstrator -- because it's stitching together the good bits.

Dialogue and recommendation: conversation and interaction logs. Improve the policy offline, evaluate it carefully, and only then let it near real users -- in stead of experimenting on them live.

The common thread: whenever a mistake during learning is expensive, dangerous or irreversible, the ability to learn from a frozen log rather than fresh trial-and-error stops being a nice-to-have and becomes the only option on the table.

The bigger picture: foundation models meet RL

Decision Transformer let a rather large cat out of the bag: RL and sequence modelling are converging. If an agent's experience is "just another sequence", then a big pre-trained model can serve as a prior for RL the same way ImageNet pre-training serves vision. Gato (DeepMind, 2022) drove this home -- a single transformer that plays Atari, controls a robot arm, captions images and chats, all from one set of weights trained on a wildly diverse offline dataset.

The trajectory (pun very much intended) is clear: collect enormous datasets of agent interactions across many tasks, pre-train one large transformer on all of it, then fine-tune -- online or offline -- for whatever specific task you care about. It's the pre-train-then-adapt recipe that already rewired NLP and vision, now aimed at decision-making. Offline RL is the piece that makes it possible, because without it "learn from a giant static log" is a non-starter.

That was a lot - here's the essence

  • Offline RL learns entirely from a fixed dataset with no environment access -- essential wherever exploration is dangerous, expensive or simply impossible;
  • distribution shift is the whole ballgame: naive Q-Learning inflates Q-values for out-of-distribution actions via the max, and the delusion never gets corrected because there's no "try it and see";
  • CQL buys safety with pessimism -- a penalty that squashes Q-values for unseen actions while protecting the in-distribution ones;
  • IQL sidesteps the problem entirely by never querying out-of-distribution actions, using expectile regression on V(s') in place of max Q(s', a');
  • Decision Transformer reframes RL as sequence modelling -- condition on the return you want and let a GPT-style model generate the actions, no Bellman equation in sight;
  • offline RL powers learning-from-history in healthcare, driving, robotics and dialogue, and its convergence with foundation models points straight at general-purpose, pre-trained agents.

Exercises

Exercise 1: Build the "naive fails, conservative saves it" experiment end to end. On a small discrete environment (a gridworld, or CartPole-v1 with discretised observations), collect an offline dataset from a deliberately mediocre behaviour policy (say, 70% random / 30% a half-decent heuristic). Train both NaiveOfflineDQN and CQLAgent on that same frozen dataset, then evaluate each by actually rolling it out in the environment. Log the maximum Q-value each method predicts as training goes on -- you should watch the naive agent's Q-values blow up toward the sky while CQL's stay grounded, and see that reflected in the rollout returns.

Exercise 2: Explore the cql_alpha knob. Take the setup from Exercise 1 and sweep cql_alpha across something like [0.0, 0.1, 1.0, 5.0, 20.0] -- noting that cql_alpha = 0.0 is exactly the naive agent again. Plot final rollout return against cql_alpha. You're looking for a U-shape: too little conservatism and overestimation wrecks you, too much and the policy is so timid it just clings to the (mediocre) behaviour policy and never improves. Somewhere in the middle is the sweet spot -- and finding it by hand is a good, humbling taste of why offline RL is fiddly in practice.

Exercise 3: Give Decision Transformer the "return dial" test. Train the DecisionTransformer above on a dataset containing a real mix of good and bad trajectories (easiest: collect from a policy at several stages of training, so returns vary a lot). Then at inference, condition on a sweep of target returns -- low, medium, high -- and measure the actual return you get for each. If it's working, higher requested returns should produce genuinely higher achieved returns, at least up to the best the data contains. Push the target beyond anything in the dataset and see what happens -- does it keep scaling, or does it top out? (That ceiling tells you something honest about what the model can and cannot invent.)

Get at least Exercise 1 running before next time -- watching the naive agent's Q-values sail off into fantasy while CQL keeps its feet on the ground is the single clearest picture of what offline RL is really fighting against. We've now spent an entire arc building the pieces -- value functions, policy gradients, model-based rollouts, imitation, and today learning from frozen logs. Next time we stop collecting techniques and start building: we're going to take everything from this reinforcement-learning arc and point it at a single, concrete goal, and actually train an agent to play a game from the ground up. Bring your GPU, or at least your patience ;-)

Bedankt en tot de volgende keer!

@scipio



0
0
0.000
0 comments