Tom Zahavy
← Home

Rewards are Gradients

A general theory of what an RL agent can be asked to do. When the goal is a convex function of the states an agent visits — exploration, imitation, diversity, constraints — the reward that solves it is the gradient of that objective.

Reward is enough for convex MDPs ↗All publications →
A cartoon: an RL agent asks “Where do rewards come from?” and the cost player answers “Rewards are gradients!”

We can describe the standard RL problem as a linear function: the inner product between the state-action occupancy and the reward vector. The occupancy is the distribution over the states and actions that an agent visits when following a policy, and the reward defines a priority over those state-action pairs.

Sometimes a good objective function is all we need — predicting the next word in a sentence turned out to be transformative. I have been interested in deriving more general objectives of behaviour: non-linear, unsupervised, convex and non-convex objectives that control the distribution an agent visits.

Objectives beyond a reward vector

Maximum entropy exploration

Visit all the states equally. The objective is the negative entropy of the occupancy — as far from linear as it gets.

Apprenticeship learning

Visit states similar to those of another policy, by minimising a distance between your occupancy and the expert's.

Diversity

Visit states that other policies do not. The objective is defined over a whole set of occupancies at once.

Constraints

Maximise reward while keeping some other statistic of the visited states within a budget.

And more: GAIL and state-marginal matching, unsupervised skill discovery, inverse RL in contextual MDPs, adversarial MDPs. All of them are convex functions of the occupancy; none of them is a reward vector.

The main result

My main result is that we can reformulate these problems as convex-concave zero-sum games, and derive a non-stationary intrinsic reward for solving them. The reward turns out to be very simple and general: it is the gradient of the objective with respect to the state occupancy.

RL
maxd⟨r, d⟩
Convex MDP
mind f(d)
The reward
rk = −∇f(dk)

d is the state-action occupancy of the policy, constrained to the set of occupancies the MDP admits. Standard RL is the special case f(d) = −⟨r, d⟩, whose gradient is the constant −r — the stationary reward we are used to.

A stationary reward cannot do it. For any fixed reward vector there is always a deterministic optimal policy — but for objectives like maximum-entropy exploration no deterministic policy is optimal. So no stationary reward has the convex objective’s solution as its optimum, and the reward has to move as the agent learns.

Fenchel duality gives the game. It converts the convex MDP into a two-player zero-sum game between the agent (the policy player) and an adversary producing rewards (the cost player). From the agent’s point of view the game is bilinear, so for a fixed reward the problem is just standard RL with a non-stationary reward — meaning any algorithm that solves standard RL can be used as a subroutine to solve the convex MDP, with the average of the policies it produces converging to the solution.

It unifies a lot of the literature. Choosing specific no-regret algorithms for the two players recovers disparate branches of RL as instances of a single meta-algorithm: apprenticeship learning, constrained MDPs, pure exploration, GAIL, adversarial MDPs, and mutual-information skill discovery such as DIAYN and VIC.

What the papers ask

  • 1In what sense are general-utility RL problems different from RL problems?
  • 2Can we solve them with similar techniques?
  • 3Are there interesting objectives we can now easily solve using this approach?
  • 4In particular — Quality-Diversity objectives.

Papers

Reward is enough for convex MDPs ↗

NeurIPS 2021 (spotlight)

We study non-linear and unsupervised objectives defined over the state occupancy of an RL agent — apprenticeship learning, diverse skill discovery, constrained MDPs, pure exploration. We show that maximising the gradient of such an objective, as an intrinsic reward, solves the problem efficiently, and propose a meta-algorithm that explains many existing algorithms as instances of it.

Tom Zahavy, Brendan O'Donoghue, Guillaume Desjardins, Satinder Singh

A simulated walker performing one of the skills discovered by the worst-case policy iteration algorithm.

Discovering a set of policies for the worst case reward ↗

ICLR 2021 (spotlight)

We propose a method for discovering a set of policies that perform well with respect to the worst-case reward when composed together. The sets it finds turn out to be diverse — distinct locomotion skills emerge without ever asking for them.

Tom Zahavy, Andre Barreto, Daniel J. Mankowitz, Shaobo Hou, Brendan O'Donoghue, Iurii Kemaev, Satinder Singh

Constraint value against episodes: the standard method oscillates wildly across the constraint threshold, while ReLOAD sits flat on it.

ReLOAD: Reinforcement Learning with Optimistic Ascent-Descent for Last-Iterate Convergence in Constrained MDPs ↗

ICML 2023

Gradient descent-ascent converges on average, but the current policy may never converge — in practice it oscillates between satisfying the constraint and maximising the reward. ReLOAD is a constrained-RL method with guaranteed last-iterate convergence, and empirically it removes the oscillation.

Ted Moskovitz, Brendan O'Donoghue, Vivek Veeriah, Sebastian Flennerhag, Satinder Singh, Tom Zahavy

Apprenticeship Learning via Frank-Wolfe ↗

AAAI 2020

We show that the well-known apprenticeship learning algorithm of Abbeel and Ng (2004) can be understood as a Frank-Wolfe method, and propose ways to accelerate it.

Tom Zahavy, Alon Cohen, Haim Kaplan, Yishay Mansour

Online Apprenticeship Learning ↗

AAAI 2021

The first apprenticeship learning algorithm that does not require solving an MDP at every iteration — both players play no-regret — with a regret analysis.

Lior Shani, Tom Zahavy, Shie Mannor

Inverse Reinforcement Learning in Contextual MDPs ↗

Machine Learning Journal 2021

Recovering a reward that explains an expert's behaviour when the task itself varies with a context — another convex objective over occupancies, solved with the same descent-ascent recipe.

Stav Belogolovsky, Philip Korsunsky, Shie Mannor, Chen Tessler, Tom Zahavy

The same ideas power AlphaZerodb, where quality-diversity objectives produce a league of chess agents that think in different ways. Filter the full publication list by Intrinsic reward or Diversity for everything else.