Offline Imitation Learning from Multiple Baselines with Applications to Compiler Optimization
Teodor V. Marinov, Alekh Agarwal, Mircea Trofin
TL;DR
The paper tackles offline imitation learning with $K$ baseline policies in a contextual finite-horizon MDP, where only trajectory-level rewards are observed. It introduces BC-Max, a simple yet effective algorithm that, for each context, imitates the actions from the highest-reward baseline trajectory and provides a regret bound Reg$\\(\\hat{\\pi}\\) \\leq O(\\epsilon H + \\frac{H^2 \\log(H|\\Pi|/\\delta)}{n})$ under a realizability assumption, with a matching lower bound establishing minimax optimality. The authors demonstrate practical value through a compiler-optimization case study, showing that iterative BC-Max can improve over an initial online RL baseline in reducing binary size, using both proprietary and Chrome-on-Android datasets. The work offers a scalable offline learning recipe with strong theoretical guarantees and real-world applicability, enabling policy improvement from multiple baselines without online interaction.
Abstract
This work studies a Reinforcement Learning (RL) problem in which we are given a set of trajectories collected with K baseline policies. Each of these policies can be quite suboptimal in isolation, and have strong performance in complementary parts of the state space. The goal is to learn a policy which performs as well as the best combination of baselines on the entire state space. We propose a simple imitation learning based algorithm, show a sample complexity bound on its accuracy and prove that the the algorithm is minimax optimal by showing a matching lower bound. Further, we apply the algorithm in the setting of machine learning guided compiler optimization to learn policies for inlining programs with the objective of creating a small binary. We demonstrate that we can learn a policy that outperforms an initial policy learned via standard RL through a few iterations of our approach.