Reward is enough for convex MDPs
Tom Zahavy, Brendan O'Donoghue, Guillaume Desjardins, Satinder Singh
cs.AI, cs.LG, stat.ML
2021-06-02
Fenchel duality recasts convex MDPs as a min-max game any RL solver can play; low-regret players yield O(1/√K) averaged occupancy.
Ordinary RL maximises the inner product of a reward vector with a policy's occupancy measure, the long-run visit frequency over state-action pairs. That formulation covers any goal that can be written as a stationary Markov reward. Many goals cannot. Apprenticeship learning asks for occupancy close to an expert. Constrained MDPs put inequalities on occupancy. Pure exploration maximises occupancy entropy. Those objectives are convex functions of occupancy, not linear ones.
Zahavy, O'Donoghue, Desjardins and Singh study convex MDPs: minimise a convex f over the polytope of feasible occupancies. Standard RL is the special case where f is linear. Two results follow. Some convex objectives admit no stationary reward that recovers the same optimum, because ordinary RL always has a deterministic optimal policy and negative entropy does not. Even so, any algorithm that solves ordinary RL can be used as a subroutine for the convex problem.
The conversion uses Fenchel duality. A convex f equals maxλ (λ · d − f(λ)), so the original problem becomes a convex-concave saddle point between a policy player and a cost player. For a fixed cost λ the policy side is ordinary RL with reward −λ.
An optimal dual variable λ is still not enough to recover the solution with a single RL call. The occupancy that minimises f is often not the occupancy of any deterministic policy. The algorithm therefore returns a sequence of policies and averages their occupancies.
Each round of the meta-algorithm, adapted from Abernethy and Wang, does two things:
Theorem 1: if the two players have average regret at most εK and δK, the averaged occupancy is within εK + δK of the convex optimum. Both regrets O(1/√K) give a gap of O(1/√K).
Convex constraints gi(d) ≤ 0 expand the Lagrangian with extra dual variables. The occupancy still enters linearly, so the same policy players apply. Non-convex f can run the same game and only yields weak-duality bounds.
The headline numbers are sample-complexity bounds, not leaderboard scores.
| Policy player | Cost player | Samples to ε-optimality |
| PAC approximate best response | oracle with regret O(1/K) | O(tmix² SA / ε³) (average reward, Lemma 2) |
| PAC approximate best response | O(1/√K) oracle | O(tmix² SA / ε⁴) |
| UCRL2 with known non-stationary rewards | O(1/√K) oracle | O(D² S² A / ε²) (Lemma 4) |
Hazan et al. need O(1/ε³) samples for approximate Frank-Wolfe on pure exploration. The UCRL2 route improves the ε exponent, with two caveats the paper states: if the cost player estimates occupancy from scratch each round, the total sample count falls back to O(1/ε³); the failure-probability dependence is 1/δ², and the authors expect a log(1/δ) bound is possible.
Table 1 slots existing algorithms into the meta-algorithm. FTL plus best response is Frank-Wolfe, which recovers Abbeel and Ng's projection method for apprenticeship learning and Hazan's pure exploration. OMD plus best response is MWAL. OMD plus RL covers constrained MDPs, adversarial MDPs and Wasserstein GAIL. FTL plus RL covers GAIL and state-marginal matching. Abbeel and Ng's max-margin variant is Fully Corrective Frank-Wolfe, which has a linear rate on polytopes, matching the faster empirical behaviour they reported.
GAIL and DIAYN share the same KL between per-skill occupancies and their mixture. GAIL minimises it and is a convex MDP. DIAYN maximises it, so the convex guarantees do not apply. The FTL gradient of that KL is the familiar intrinsic reward log p(z|s) − log p(z).
The only experiment sits in an appendix. On bsuite Deep Sea they run entropy-constrained RL: maximise environmental reward subject to occupancy entropy at least half the environment maximum. Cost player is FTL, policy player is Impala. Vanilla Impala fails this hard-exploration task. The entropy-constrained version tracks Bootstrapped DQN, which handles uncertainty with an ensemble. No numeric returns are reported, only Figure 2.
The portable recipe is: take the gradient of a convex occupancy objective, feed it as a non-stationary reward to an off-the-shelf RL algorithm, and average policies. Apprenticeship, constraints, exploration and GAIL become different (cost player, policy player) pairs inside one game.
This is a unification paper, not a new leaderboard agent. Deep Sea shows the heuristic can work. One environment is not evidence for production use. In a real domain the objective has to be convex in occupancy, and occupancy has to be estimable; a bad occupancy estimate poisons the cost player's rewards.
Lemma 4 treats the cost player as an oracle. Occupancy must be estimated from samples; without reuse across rounds the O(1/ε²) bound drops back to O(1/ε³). The δ dependence is polynomial, not logarithmic. Convergence is in the average occupancy, so the last policy need not be optimal, which sits poorly with deep RL where one network is kept at the end of training.
The empirical section is a single Deep Sea plot against vanilla Impala and Bootstrapped DQN, with no tabulated scores and no ablation of the cost player. DIAYN appears in Table 1 even though the text states it maximises a convex function. The analysis assumes finite state-action spaces; there is no regret theory under function approximation.