Sharp Sample Complexity of KL-Regularized Reinforcement Learning

KL-regularized reinforcement learning, the objective behind maximum-entropy RL and RL from human feedback, is variational inference in the graphical model of "control as inference": the regularized optimal policy is a posterior over trajectories, and the suboptimality of a policy is $β$ times a KL divergence between path measures. However, it is not known how many samples it takes to learn this posterior, even with a generative model. Earlier $1/ε$ guarantees scale with $H^5$ in the effective horizon $H$, and no lower bound for MDPs depends on the temperature $β$. In this work, we show that for a discounted MDP with $K$ state-action pairs, the plug-in estimator finds a policy that is $ε$-optimal from every initial state with $\widetilde{O}(KH^3/(ε\max\{ε,β\}) + KH)$ samples, and that no algorithm, even one that chooses where to sample, can do with fewer when $ε\lesssim \min\{H, H^2/β\}$…