KL-regularized RL gets a sharp sample count: with a generative model, samples scale linearly in the effective horizon and temperature no longer slows learning
Synopsis
KL-regularized reinforcement learning, the objective behind maximum-entropy RL and RL fine-tuning of language models from human feedback, previously had guarantees whose sample counts grew with a worse power of the effective horizon, and no MDP lower bound depended on the temperature. For tabular discounted MDPs with a generative model, this work shows that a plug-in estimator, which draws a fixed number of samples per state-action pair and solves the empirical MDP, reaches KL suboptimality at the temperature scale with samples linear in the effective horizon, and that no algorithm, even one choosing where to sample, can do with fewer.
Interpretation
For tabular discounted MDPs with known rewards and a generative model that returns independent draws for any queried state-action pair, the samples needed to learn a KL-regularized optimal policy are pinned down up to logarithmic factors, and the count grows linearly in the effective horizon, as in unregularized RL. Earlier guarantees came from mirror-descent methods and were no better than a worse power of the effective horizon; the plug-in estimator inherited a looser rate through a sup-norm bound, and no MDP lower bound depended on the temperature, leaving the optimal rate open. The upper and lower bounds match within the same temperature range: the lower bound builds two MDPs with five state-action pairs such that any estimator using a fixed number of samples per pair fails on one of them, and an adaptive-sampling lower bound uses a family of MDPs with at most a bounded number of pairs.
The rate is attained by a plug-in estimator: sample every state-action pair equally often, replace the true kernel with the empirical one, and solve the empirical MDP exactly for its soft-optimal policy. The proof does not sum per-step worst cases; it bounds the across-action variance of the error along the learner's own trajectories with a martingale identity, removing an extra factor. Simulations on small MDPs show the risk falling with temperature at different powers in the fast and slow regimes, and after rescaling the four temperature curves fall on one line.
Beyond the rate, the asymptotic constant of the plug-in estimator is given exactly as a Fisher-weighted variance of the estimated soft Q-function. Earlier work gave rates but no constant; this constant is locally asymptotically minimax among methods that sample every pair equally often. On seven instances whose constants span four orders of magnitude, including gadgets, random Garnet MDPs and a 100-state ring, the theoretical constant agrees with the measured risk to within a few percent.
Perspective
The results cover tabular discounted MDPs with known rewards and data from a generative model that returns independent draws for any state-action pair; this setting isolates the cost of learning the transition kernel and leaves out exploration and function approximation. Simulations are on small MDPs only, including a two-action gadget, random Garnet instances and a 100-state ring.
The lower bound leaves a window uncovered when the temperature falls below a threshold, where the upper bound is set by the burn-in sample count. The optimal allocation of Theorem 4 needs weights that depend on the unknown instance and would have to be estimated. The constants of Theorem 1 are far from tight, and its logarithmic factor includes an extra term from the decoupling argument, so most simulated sample sizes lie below the burn-in it requires.