Implementing Inverse Q-Learning

Reinforcement learning (RL) agents can need millions of interactions, the equivalent of years of experience, to learn a simple task, while people can often pick something up from a few examples. Imitation learning tries to close some of that gap by learning from demonstrations. For this project my partner and I reimplemented IQ-Learn, an imitation learning algorithm, and tested it in a small grid world.

In 1998, Stuart Russell posed inverse reinforcement learning (IRL) as a problem: given examples of an agent’s behavior, can we work out what reward function it’s optimizing? Early methods needed nested optimization. They guessed a reward, solved for the best policy under it, compared that policy to the expert, and updated the reward.

Behavioral cloning takes a simpler route and copies expert actions directly with supervised learning. It tends to fail once the agent drifts away from the expert’s trajectories, because it has no data on how to recover. IQ-Learn instead learns a single Q-function, and both the policy and the reward can be read off of it.

Reinforcement Learning

Reinforcement learning is about learning to make good decisions through trial and error. Take chess: each move leads to a new board position, and eventually to a win or a loss. To get better, you need to figure out which moves in which positions tend to lead to wins.

More formally, we model this as a Markov Decision Process (MDP) with:

States $s$ (like the board position in chess)
Actions $a$ (like possible moves)
A reward function $r(s,a)$ (like +1 for winning, -1 for losing)
State transitions $P(s'|s,a)$ (how the board changes after a move)
A discount factor $\gamma$ (making future rewards worth less than immediate ones)

The goal is to find a policy $\pi(a|s)$ that maximizes expected future reward. Q-learning does this by learning a Q-function $Q(s,a)$ , the expected sum of future rewards from taking action $a$ in state $s$ and then acting optimally:

$Q(s,a) = r(s,a) + \gamma \mathbb{E}_{s'\sim P(s'|s,a)}[\max_{a'} Q(s',a')]$

This is the Bellman equation. It says the value of an action now is the immediate reward plus the discounted value of the best you can do afterward. Once you know the true Q-function, the optimal policy is to pick the action with the highest Q-value in each state.

The Inverse Problem

In many real situations we can’t easily write down what makes an action “good”, but we do have experts who know what they’re doing. That’s the setting for inverse reinforcement learning.

Traditional IRL methods:

Guess a reward function
Run RL to find the optimal policy for those rewards
Compare that policy to the expert
Adjust the rewards and repeat

This is expensive and can be unstable. IQ-Learn skips the loop by learning the Q-function directly from expert behavior and extracting both the policy and the rewards from it.

The IQ-Learn Algorithm

IQ-Learn changes the Q-learning objective so that it learns from expert demonstrations without any rewards. Roughly, it pushes up the Q-values of the actions the expert took relative to other actions:

$\max_Q \mathbb{E}_{(s,a,s')\sim\rho_E}[Q(s,a) - \gamma V(s')] - (1-\gamma)\mathbb{E}_{s\sim\rho_0}[V(s)]$

where $V(s) = \max_a Q(s,a)$ is the value function, $\rho_E$ is the expert’s state-action distribution, and $\rho_0$ is the initial state distribution.

The authors show that:

Optimizing this objective is equivalent to minimizing a statistical distance between the expert’s behavior and the learned behavior
The learned Q-function implicitly defines a reward function
The policy can be recovered directly as $\pi(a|s) \propto \exp(Q(s,a))$

Our Implementation

The original paper tested IQ-Learn on MuJoCo and Atari. We wanted to understand the core algorithm, so we used a much simpler setting:

Only the base algorithm, with total variation distance
Discrete actions instead of continuous control
A configurable GridWorld environment we built for testing

GridWorld

This is the environment the agents learned in. Use the arrow keys to move the red square (you) to the yellow square (the goal) without running into the black walls.

The checkboxes control the difficulty:

Random Walls moves the walls each time you reach the goal
Random Start puts you somewhere random instead of the top-left corner
Random Goal puts the goal somewhere random instead of the bottom-right corner

You can turn on several at once. Try it with none first and then add them one at a time to see how much harder it gets.

Click the grid or Tab to it before using the arrow keys. When you reach the goal, the grid resets based on the options you’ve picked.

Random WallsRandom StartRandom Goal

Focus the grid, then use the arrow keys to move

The Environment

We trained on three versions of GridWorld:

Static Walls: fixed walls and goal
Random Walls: random walls, fixed goal
Random Everything: random walls and goal

Players (or agents) move in four directions and have to reach the goal without hitting walls.

Implementation Details

The Q-Network

class Model:
    def __init__(self, input_size, hidden_size, output_size):
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size)
        )

The IQ-Learn Update

def iq_learn_update(self, expert_states, expert_actions):
    current_Q = Q_vals.gather(2, action.unsqueeze(1)).squeeze()
    current_V = get_max_Q(Q_vals)
    next_V = get_max_Q(self.model.get_Q(next_state))
    y = (1 - done) * self.gamma * next_V
    loss = iq_loss(current_Q, current_V, y)

State Processing To avoid “dead neurons” during training, we added a little Gaussian noise to the state representations:

def preprocess_state(state):
    noise = np.random.normal(0, 0.01, state.shape)
    return state + noise

Training

We compared IQ-Learn against DQN, a standard Q-learning agent that gets the true rewards.

IQ-Learn usually converged within 200 training steps, and DQN needed a lot more, even though it was the one getting the true rewards. Both got worse as the environment got more random, and both had trouble avoiding walls in the harder versions.

DQN’s value maps had a clear gradient toward the goal, which is what you’d draw if someone asked you to sketch a value map. IQ-Learn’s looked less structured, yet its policies still reached the goal.

The expert demonstrations clustered near the optimal paths when the walls were static, spread out more when the walls were random, and concentrated in the middle of the grid when the goal was random.

To score paths, we divided each path’s length by the shortest possible distance to the goal, so 1.0 is optimal.

Limitations

The simple MLP struggled with spatial relationships, and neither algorithm learned to avoid walls in a way that generalized. A CNN would probably capture the grid’s spatial structure better.

IQ-Learn was sensitive to how the states were preprocessed. Batch size had a big effect on convergence, and we needed a learning rate schedule to keep training stable.

Compared to the Original Paper

We only covered part of what the paper did. We used the basic algorithm with no extensions, tested only in simple environments with discrete actions, and looked at learning dynamics instead of benchmarking against other methods. We didn’t try continuous control, the other statistical distances, or the actor-critic variants the paper uses for high-dimensional problems.

The obvious next steps would be a convolutional Q-network, the other distances, and continuous actions. We’d also want better visualizations of the learned policies and a closer look at how well the recovered rewards match the real ones.

The code and documentation are in the project repository. Special thanks to my project partner Celine Aubuchon. Without Celine there would be no game, no analysis of the agents’ behavior, and I wouldn’t have understood the original work nearly as well.