Learn AI Series):Before we flip reinforcement learning on its head, let's clear last episode's three exercises. They lean on the DataCenterEnv, SafeRLAgent classes from episode #113, plus a DQNAgent from episode #107, so I'm assuming those are imported and in scope.
Exercise 1: Train a small DQN on the DataCenterEnv over a handful of discretized cooling levels, log average power and the fraction of steps spent outside the safe band, then crank the out-of-band penalty from 10.0 to 100.0 and watch the behaviour shift.
The one wrinkle is that episode #113's DataCenterEnv hard-codes the penalty at 10.0 inside step, so to change it we subclass and expose it as a knob:
import numpy as np
# Assumes DataCenterEnv from episode #113 and DQNAgent from episode #107.
COOLING_LEVELS = np.array([0.0, 0.25, 0.5, 0.75, 1.0])
class TunableDataCenterEnv(DataCenterEnv):
"""Same environment, but the out-of-band penalty is a parameter."""
def __init__(self, penalty=10.0):
super().__init__()
self.penalty = penalty
def step(self, action):
server_heat = np.random.uniform(2.0, 5.0)
self.temperature += server_heat - action * 8.0 \
+ 0.1 * (self.outside_temp - self.temperature)
self.power_usage = action
reward = -self.power_usage
lo, hi = self.target_range
if self.temperature > hi:
reward -= self.penalty * (self.temperature - hi)
elif self.temperature < lo:
reward -= self.penalty * (lo - self.temperature)
obs = np.array([self.temperature, self.outside_temp,
server_heat, self.power_usage], dtype=np.float32)
return obs, reward, False
def train_cooling(penalty, episodes=2000, steps=48):
env = TunableDataCenterEnv(penalty=penalty)
agent = DQNAgent(state_dim=4, n_actions=len(COOLING_LEVELS))
power_hist, oob_hist = [], []
for ep in range(episodes):
env.temperature = 22.0 # reset the room each episode
obs = np.array([env.temperature, env.outside_temp, 0.0, 0.0],
dtype=np.float32)
ep_power, ep_oob = 0.0, 0
for t in range(steps):
a = agent.choose_action(obs)
nobs, r, _ = env.step(COOLING_LEVELS[a])
agent.store(obs, a, r, nobs, t == steps - 1)
agent.learn()
obs = nobs
ep_power += env.power_usage
lo, hi = env.target_range
if not (lo <= env.temperature <= hi):
ep_oob += 1
power_hist.append(ep_power / steps)
oob_hist.append(ep_oob / steps)
return np.mean(power_hist[-100:]), np.mean(oob_hist[-100:])
for pen in (10.0, 100.0):
p, o = train_cooling(pen)
print(f"penalty={pen:5.0f}: avg power={p:.3f}, out-of-band fraction={o:.3f}")
The lesson lives entirely in that last print. With penalty=10.0 the agent happily accepts the occasional excursion out of the safe band because a little overheating is cheaper than running the fans harder -- so it spends less power but drifts out of range more often. Bump the penalty to 100.0 and the trade-off inverts: now leaving the band is agony, so the agent spends more power to stay comfortably inside it, and the out-of-band fraction collapses toward zero. Same environment, same algorithm, one number changed -- and you've just watched a reward weight literally encode a priority. That is the whole art of reward design in a single experiment, and remember it well, because in a few minutes we're going to admit we often can't write that number down at all.
Exercise 2: Reproduce the reality gap. Train a policy on a SimHand with friction=0.5, test it on a RealHand with friction=0.35, then train a second policy with domain randomization and see it survive the jump far better.
I'll use a tiny cross-entropy-method optimiser here rather than full PPO -- it's completely self-contained, it's robust, and it keeps the exercise about transfer rather than about tuning a policy-gradient:
import numpy as np
class SlideEnv:
"""1D 'push the block to x=1.0'. Higher friction damps each push."""
def __init__(self, friction=0.5):
self.friction = friction
def reset(self):
self.x, self.v, self.t = 0.0, 0.0, 0
return np.array([self.x, self.v], dtype=np.float32)
def step(self, action):
force = float(np.clip(action, -1.0, 1.0))
self.v = (1.0 - self.friction) * self.v + 0.1 * force
self.x += self.v
self.t += 1
return (np.array([self.x, self.v], dtype=np.float32),
-abs(self.x - 1.0), self.t >= 50)
def rollout(env_factory, theta, n=5):
"""Mean return of the linear policy a = tanh(theta . obs)."""
total = 0.0
for _ in range(n):
env = env_factory()
obs, done = env.reset(), False
while not done:
obs, r, done = env.step(np.tanh(theta @ obs))
total += r
return total / n
def cem_train(env_factory, dim=2, iters=30, pop=40, elite=8):
mean, std = np.zeros(dim), np.ones(dim)
for _ in range(iters):
pops = mean + std * np.random.randn(pop, dim)
scores = np.array([rollout(env_factory, th) for th in pops])
elite_idx = scores.argsort()[-elite:]
mean, std = pops[elite_idx].mean(0), pops[elite_idx].std(0) + 1e-6
return mean
real = lambda: SlideEnv(friction=0.35)
theta_sim = cem_train(lambda: SlideEnv(friction=0.5)) # trained on ONE sim
print("sim-only -> real:", round(rollout(real, theta_sim, n=100), 3))
theta_dr = cem_train(lambda: SlideEnv(friction=np.random.uniform(0.2, 0.8)))
print("randomized -> real:", round(rollout(real, theta_dr, n=100), 3))
The single-sim policy overfits to friction 0.5 -- it has learned exactly how hard to shove assuming that one damping, and on the slippier real hand it consistently overshoots the target and pays for it in the return. The domain-randomized policy never got to assume a friction at all; it was forced to work across the whole [0.2, 0.8] range, and 0.35 is just one more value it already handles. You'll typically see it land a better (less negative) return on real. Nota bene: this task is deliberately gentle, so don't expect a dramatic gulf -- the direction is the point, and it's the same direction that got OpenAI's cube-solving hand out of the simulator and onto real hardware.
Exercise 3: Build the missing safety_checker for SafeRLAgent on a 1D cliff world -- position clamped to [0, 10] -- wrap a deliberately reckless policy, and confirm the agent never once leaves the safe interval, logging how often the safety layer had to step in.
import numpy as np
# Assumes SafeRLAgent from episode #113.
class CliffSafetyChecker:
"""Position must stay in [low, high]. Actions are 1D steps."""
def __init__(self, low=0.0, high=10.0, moves=(-1.0, 0.0, 1.0)):
self.low, self.high = low, high
self.moves = [np.array([m]) for m in moves]
def is_safe(self, state, action):
return self.low <= state[0] + action[0] <= self.high
def get_safe_actions(self, state):
# 'stay put' (0.0) is always safe, so this never returns empty.
return [m for m in self.moves if self.is_safe(state, m)]
class RecklessPolicy:
"""No idea safety exists -- proposes a random step every call."""
def __call__(self, state):
return np.array([np.random.uniform(-1.0, 1.0)])
checker = CliffSafetyChecker()
agent = SafeRLAgent(RecklessPolicy(), checker)
state, overrides = np.array([5.0]), 0
for _ in range(500):
proposed = agent.policy(state) # what the reckless policy wants
if checker.is_safe(state, proposed):
action = proposed
else:
overrides += 1 # safety layer had to intervene
action = agent._nearest_safe_action(proposed, checker.get_safe_actions(state))
state = state + action
assert 0.0 <= state[0] <= 10.0 # the guarantee we care about
print(f"survived 500 steps, safety layer overrode {overrides} times")
The assert never fires -- that's the whole exercise. No matter how badly the underlying policy misbehaves, the wrapper draws a fence the agent physically cannot cross, and the override counter tells you exactly how often the reckless policy tried to. In a real deployment that override count is a number a safety engineer watches like a hawk: a policy that has to be corrected constantly isn't safe, it's just being held safe, and the day the checker has a bug is the day someone gets hurt. Right -- with the safety scaffolding cleared, let's turn RL completely inside out.
For thirteen episodes we've assumed the reward function is handed to us -- score in a game, distance to a goal, energy saved, profit earned -- and the whole job was finding the policy that maximises it. Episode #113 ended on an uncomfortable admission, though: we've been hand-designing those rewards the entire time, and for genuinely complex behaviour that's nearly impossible to do. So today we flip the arrow around.
Standard RL has a reward function and finds the optimal policy. Inverse RL (IRL) has a policy -- observed behaviour -- and finds the reward function. It answers a very human question: "given that an expert acts this way, what must they be optimising for?"
This reversal is much more than a cute theoretical curiosity. Designing good reward functions is brutally hard. In episode #112 we saw reward shaping for games, and in #113 we shaped a pick-and-place reward by hand -- and even those toy examples were fiddly. But for the behaviours we actually care about -- driving a car, performing surgery, writing a genuinely helpful answer -- specifying what "good" means as a clean mathematical function is somewhere between miserable and impossible. What you can do, quite easily, is show examples of good behaviour. Inverse RL extracts the implicit reward hiding inside those examples.
And if that setup sounds oddly familiar, it should. RLHF -- the technique from episode #61 that made assistants genuinely useful -- is essentially inverse RL pointed at language models. Human preferences define what's "good", a reward model learns from those preferences, and RL optimises the model against that learned reward. The intellectual lineage runs directly from dusty IRL papers to modern alignment. Having said that, let's build up to that connection properly, starting with the awkward fact that makes IRL hard in the first place.
Here's the catch that keeps IRL researchers up at night: given a set of expert demonstrations, many different reward functions could explain the exact same behaviour. An expert driver who stops at red lights might be optimising for "follow traffic rules", or "don't hit other cars", or "minimise fuel consumption", or simply "arrive alive". Every one of those reward functions produces identical behaviour at a red light.
This ambiguity is fundamental, not a bug you can engineer away. The behaviour underdetermines the reward -- there's genuinely not enough information in "what they did" to pin down "why they did it". Worse, a reward function of all zeros makes every policy optimal, so it trivially "explains" any demonstration you like. Useless, obviously -- but technically a perfect fit ;-) We need extra principles to pick a sensible reward out of the infinite pile that fit. That's what the rest of this episode is really about.
The classic first attempt (Abbeel & Ng, 2004) makes a simplifying assumption: the reward is a linear combination of features, and we hunt for the weights that make the expert's behaviour look optimal.
import numpy as np
class LinearIRL:
"""Feature-based Inverse RL: recover reward weights from demonstrations."""
def __init__(self, n_features, gamma=0.99):
self.n_features = n_features
self.gamma = gamma
def feature_expectations(self, trajectories, feature_fn):
"""Expected discounted feature counts over a set of trajectories."""
mu = np.zeros(self.n_features)
for trajectory in trajectories:
for t, (state, action) in enumerate(trajectory):
mu += (self.gamma ** t) * feature_fn(state, action)
return mu / len(trajectories)
def learn_reward(self, expert_trajectories, feature_fn,
rl_solver, n_iterations=50):
"""Apprenticeship-learning projection algorithm."""
mu_expert = self.feature_expectations(expert_trajectories, feature_fn)
w = np.random.randn(self.n_features)
w /= np.linalg.norm(w)
best_margin, best_w = -np.inf, w.copy()
for i in range(n_iterations):
# Reward under the CURRENT guess is R(s,a) = w . phi(s,a).
reward_fn = lambda s, a: np.dot(w, feature_fn(s, a))
# Solve for the optimal policy under that reward...
policy = rl_solver.solve(reward_fn)
mu_learned = self.feature_expectations(
rl_solver.generate_trajectories(policy), feature_fn)
# ...then push the weights toward wherever the expert differs.
diff = mu_expert - mu_learned
margin = np.dot(w, diff)
if margin > best_margin:
best_margin, best_w = margin, w.copy()
w = diff / (np.linalg.norm(diff) + 1e-8)
print(f"iteration {i}: margin = {margin:.4f}")
return best_w
The loop alternates between two moves. First, treat the current weights as the reward and solve for the best policy under it (that inner rl_solver is doing full forward RL, the stuff we've built all series). Second, nudge the weights so the expert's feature expectations look better than that learned policy's. It converges when no policy can beat the expert under the recovered reward -- which is a rather elegant definition of "I've explained what the expert wanted". The obvious wart: there's a complete RL solve inside every iteration, which is expensive. That cost is exactly what the next ideas try to dodge.
MaxEnt IRL (Ziebart et al., 2008) resolves the reward ambiguity with a genuinely principled tie-breaker: among all the reward functions consistent with the demonstrations, prefer the one that assumes the least. Formally it picks the distribution over trajectories with maximum entropy subject to matching the expert's feature expectations -- the probabilistic version of "don't commit to anything the data didn't force you to".
The model that falls out says the probability of a trajectory is proportional to the exponential of its total reward:
P(tau) proportional to exp( sum_t R(s_t, a_t) )
Higher-reward trajectories are more probable, but every trajectory keeps some probability. That's not a fudge -- it's a feature. Real experts are noisy; they don't take the identical path twice, and MaxEnt bakes that variability straight into the model in stead of pretending the expert is a flawless robot.
import torch
import torch.nn as nn
class MaxEntIRL:
"""Maximum Entropy Inverse RL with a neural reward function."""
def __init__(self, state_dim, hidden=64, lr=1e-3):
self.reward_net = nn.Sequential(
nn.Linear(state_dim, hidden), nn.ReLU(),
nn.Linear(hidden, hidden), nn.ReLU(),
nn.Linear(hidden, 1),
)
self.optimizer = torch.optim.Adam(self.reward_net.parameters(), lr=lr)
def reward(self, states):
return self.reward_net(states).squeeze(-1)
def update(self, expert_states, policy_states):
"""Push expert-state rewards UP and current-policy-state rewards DOWN."""
expert_r = self.reward(torch.FloatTensor(np.array(expert_states)))
policy_r = self.reward(torch.FloatTensor(np.array(policy_states)))
# Maximise expert log-likelihood == expert reward above policy reward.
loss = -(expert_r.mean() - policy_r.mean())
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
The reward_net is where MaxEnt earns its modern keep: swap the linear reward for a neural net and you get deep MaxEnt IRL, able to recover reward functions no hand-picked feature set could express. The training still alternates -- use the current reward to train a policy, then update the reward so expert states outscore that policy's states -- but now the reward can be arbitrarily expressive. Look closely at that loss and you may feel a flicker of recognition: "make the expert score higher than the current policy" is a discriminative game between two distributions. Hold that thought, because the next section makes it explicit.
GAIL (Ho & Ermon, 2016) is the result that connected IRL to the GANs we built in episode #55, and it's genuinely clever. Instead of laboriously recovering an explicit reward function, GAIL trains a discriminator to tell expert behaviour apart from agent behaviour -- and then hands the discriminator's output to the policy as the reward.
class GAIL:
"""Generative Adversarial Imitation Learning."""
def __init__(self, state_dim, action_dim, hidden=128):
# Discriminator: is this (state, action) from the expert or the agent?
self.discriminator = nn.Sequential(
nn.Linear(state_dim + action_dim, hidden), nn.ReLU(),
nn.Linear(hidden, hidden), nn.ReLU(),
nn.Linear(hidden, 1), nn.Sigmoid(),
)
self.disc_optimizer = torch.optim.Adam(
self.discriminator.parameters(), lr=3e-4)
# The policy (generator) is trained with PPO from episode #109.
self.policy = PolicyNetwork(state_dim, action_dim)
self.policy_optimizer = torch.optim.Adam(
self.policy.parameters(), lr=3e-4)
def discriminator_reward(self, state, action):
"""Turn the discriminator into a reward: 'fool me and get paid'."""
sa = torch.cat([torch.FloatTensor(state),
torch.FloatTensor(action)]).unsqueeze(0)
with torch.no_grad():
d = self.discriminator(sa)
return -torch.log(1.0 - d + 1e-8).item()
def update_discriminator(self, expert_sa, agent_sa):
expert_t = torch.FloatTensor(np.array(expert_sa))
agent_t = torch.FloatTensor(np.array(agent_sa))
expert_pred = self.discriminator(expert_t) # want -> 1 (expert)
agent_pred = self.discriminator(agent_t) # want -> 0 (agent)
loss = (-torch.log(expert_pred + 1e-8).mean()
- torch.log(1.0 - agent_pred + 1e-8).mean())
self.disc_optimizer.zero_grad()
loss.backward()
self.disc_optimizer.step()
return loss.item()
def train_step(self, env, expert_demonstrations):
"""One round: collect agent data, sharpen the critic, improve the policy."""
agent_trajectories = self.collect_trajectories(env, n_episodes=10)
expert_sa = [(s, a) for traj in expert_demonstrations for s, a in traj]
agent_sa = [(s, a) for traj in agent_trajectories for s, a, r in traj]
disc_loss = self.update_discriminator(expert_sa, agent_sa)
# Policy improves via PPO using discriminator_reward for each transition.
self.update_policy_with_ppo(agent_trajectories)
return disc_loss
The beauty is that the discriminator implicitly is the reward function, and ordinary GAN training handles the alternating optimisation for free -- no explicit reward solve, no inner RL loop per iteration. The policy learns to behave indistinguishably from the expert, and crucially it does so by matching the whole distribution of state-action pairs, not by parroting individual actions. That distinction is the difference between imitation that generalises and imitation that shatters the moment the world drifts -- which brings us neatly to the connection I promised, and then to why naive copying fails.
Line up RLHF for language models (episode #61) next to IRL and they're the same skeleton wearing different clothes:
The reward model in RLHF is doing precisely what MaxEnt IRL does: inferring what humans value from observed choices. The policy -- the language model itself -- is then optimised against that inferred reward.
class RLHFAsIRL:
"""RLHF, viewed deliberately through the IRL lens."""
def __init__(self, language_model, reward_model):
self.lm = language_model # the policy / generator
self.rm = reward_model # the INFERRED reward (the IRL product)
def generate_and_score(self, prompt):
response = self.lm.generate(prompt) # policy acts
reward = self.rm.score(prompt, response) # inferred reward judges
return response, reward
# Training maximises reward while a KL penalty keeps the policy near the
# base model -- standard PPO with a reference policy, straight from #109.
DPO (Direct Preference Optimization) takes this one step further by collapsing the reward-learning and policy-optimisation stages into a single loss -- no separate reward model at all -- but the conceptual frame is still pure IRL: figure out what's wanted from preferences, then act on it. When people say alignment is "just" fine-tuning, this is the bit they're skating over: at its core it's reward inference, the exact problem IRL has been chewing on since 2004.
Now, why go to all this trouble? Isn't there a lazier way to imitate an expert? There is, and it's called behavioral cloning (BC): plain supervised learning that maps states straight to the expert's actions.
import torch.nn as nn
class BehavioralCloning(nn.Module):
"""Supervised imitation: just predict the expert's action from the state."""
def __init__(self, state_dim, action_dim, hidden=128):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden), nn.ReLU(),
nn.Linear(hidden, action_dim),
)
def forward(self, state):
return self.net(state)
def train_step(self, states, expert_actions, optimizer):
pred = self.net(torch.FloatTensor(states))
loss = nn.functional.mse_loss(pred, torch.FloatTensor(expert_actions))
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss.item()
BC is fast, simple and often astonishingly effective -- and then it falls off a cliff. The failure mode is compounding errors: the agent makes one small mistake, drifts into a state the expert never visited (and therefore never demonstrated), has no idea what to do there, makes a bigger mistake, drifts further... and the whole thing snowballs. A self-driving BC policy that's a hair too far right has never seen "a hair too far right", so it doesn't know to correct, so it ends up in the ditch.
IRL sidesteps this by learning the underlying reward rather than the surface actions. Recover what the expert was optimising for and the agent can use ordinary RL to behave well even in states nobody demonstrated -- because it knows what to optimise, not merely what was done. The reward transfers; a lookup table of actions does not.
| Behavioral Cloning | IRL / GAIL | |
|---|---|---|
| What's learned | action = f(state) | reward = g(state, action) |
| Training | single supervised pass | iterative RL + reward learning |
| Compounding errors | yes (drifts off the expert's path) | no (reward generalises) |
| Sample efficiency | very high | lower (needs an RL inner loop) |
| Novel situations | fails | can generalise via the reward |
IRL is the right tool when:
IRL is overkill when:
The honest engineering take: try behavioral cloning first because it's cheap, and reach for GAIL or MaxEnt IRL only when BC's brittleness actually shows up in testing. Most forementioned "we need IRL" instincts evaporate the moment someone checks whether plain cloning was good enough ;-)
Exercise 1: Build the tiniest possible IRL demo on a 5x5 gridworld. Hand-pick a "true" reward (say +1 in one corner, 0 elsewhere), solve it with value iteration from episode #104 to get an expert policy, and roll out a handful of expert trajectories. Now throw the true reward away and run the LinearIRL projection loop above, using one-hot cell indicators as features, to recover reward weights from those trajectories alone. Plot the recovered reward as a heatmap next to the true one. You won't get the true reward back exactly -- reward ambiguity guarantees that -- but the optimal policy under your recovered reward should match the expert's. That gap between "same reward" and "same behaviour" is the entire lesson of this episode, made visible.
Exercise 2: Feel the compounding-error cliff for yourself. On CartPole-v1, train an expert with PPO (episode #109), collect 50 expert episodes, and fit a BehavioralCloning policy on the (state, action) pairs. Measure the BC policy's average episode length. Then deliberately perturb the start state a little more than the expert ever saw and measure again -- watch performance sag as the agent wanders off the demonstrated manifold. For bonus insight, plot how BC performance scales as you feed it 5, 10, 25 then 50 demonstration episodes; more data pushes the cliff back but never removes it.
Exercise 3: Turn the MaxEntIRL reward-update into something you can actually watch converge. On the same gridworld as Exercise 1, alternate: (a) train a simple policy with your current reward_net, (b) call update with expert states versus your policy's states. Every few rounds, print the correlation between the recovered reward and the true reward across all cells. You should see the reward network gradually line up with the truth (up to a constant offset and scale -- remember, adding a constant to every reward changes nothing). Then sabotage it: give the expert demonstrations that only ever visit half the grid, and observe the recovered reward go wild in the unvisited region. That's underdetermination biting -- the demonstrations simply never said what happens over there.
Get at least the gridworld in Exercise 1 running before the next episode -- recovering a reward from nothing but behaviour, and watching the policy match even though the numbers don't, is the single best intuition for why "learning what someone wants" is both possible and slippery. We've spent this whole arc assuming we can freely collect fresh experience by acting in the environment. But what if you can't? What if all you have is a fixed, frozen log of somebody else's interactions, and no way to try anything new? That constraint changes the rules more than you'd expect -- and it's where we head next.
@scipio