Horizon-Compensated Natural Policy Updates
Implementation & benchmark of arXiv:2607.22982 — Finite-Time Analysis of the Natural Policy Gradient in Finite-Horizon Markov Decision Processes
Source paper: Finite-Time Analysis of the Natural Policy Gradient in Finite-Horizon Markov Decision Processes arXiv:2607.22982 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace the constant policy natural-gradient step size in an episodic RL algorithm with the paper's geometrically increasing schedule, eta_t = eta_0 (H/(H-1))^t for horizon H >= 2. The schedule is intended to offset horizon-dependent decay of policy-improvement signals and produce geometric rather than sublinear convergence in finite-horizon policy optimization.
Formulas
Mathematical statement
The paper studies finite-horizon policies pi = (pi^1,...,pi^H), with value function V^{pi,h}(s), action value Q^{pi,h}(s,a) = R^h(s,a) + sum_{s'} P^h(s'|s,a)V^{pi,h+1}(s'), and advantage A^{pi,h}(s,a) = Q^{pi,h}(s,a) - V^{pi,h}(s). Here H is the episode horizon, h is the decision time, s is a state, a is an action, R^h is the reward, and P^h is the transition kernel. With constant step size eta_t = eta, the paper proves an O(H^2/t) convergence rate. With increasing steps, it proves an error rate O((1 - 1/vartheta_rho)^t), where vartheta_rho > 1 is a problem-dependent concentrability constant determined by the initial-state distribution rho and optimal-policy visitation distributions. The robust schedule eta_t = eta_0(H/(H-1))^t achieves this geometric rate without knowing vartheta_rho. The implementation uses the exact tabular natural-gradient policy update pi_{t+1}^h(a|s) proportional to pi_t^h(a|s) exp(eta_t A_t^h(s,a)); for a neural policy, this becomes a Fisher-preconditioned gradient update with an optional KL cap.
Implementation notes
1. Integration point: modify the policy-update step of a finite-horizon PPO, vanilla natural policy gradient, or tabular softmax actor. Leave rollout collection, critic targets, and GAE unchanged initially. For a neural policy with logits z_theta(h,s,a), compute the policy-gradient vector g_t and an approximate Fisher matrix F_t, then use the damped natural-gradient direction d_t = (F_t + lambda I)^(-1)g_t. Multiply this direction by the scheduled scalar eta_t. Enforce an empirical KL cap to prevent the theoretically unbounded schedule from destabilizing noisy neural updates.
2. Pseudocode:
input H, eta0, policy theta, damping lambda, KL_max
for t = 0,...,T-1:
collect complete H-step trajectories
estimate Ahat[h] and policy gradient g
define Fisher-vector product Fv from grad log pi_theta(a|s)
solve (F + lambda I)d = g with conjugate gradient
eta = eta0 * (H/(H-1))**t
theta_trial = theta + eta*d
kl = mean_KL(pi_theta, pi_theta_trial)
if kl > KL_max:
theta_trial = theta + sqrt(KL_max/kl)*eta*d
theta = theta_trial
For a tabular test, directly use the exponentiated update in the fourth formula. Compute from the paper's mathematics the multiplier H/(H-1), schedule, and predicted geometric-versus-sublinear comparison. Estimate advantages, Fisher-vector products, KL, and empirical contraction factors from successive optimality gaps. First run a tabular 15-state, 5-action MDP with H in {5,10,20}, exact dynamic-programming Q values, and transitions P(s'|s,a) = 0.7 1{s'=s} + 0.3 1{s'=s+1 mod |S|}. Compare constant-step NPG, increasing-step NPG, and fixed-KL PPO at equal environment interactions. Plot optimality gap and KL on a log scale. The expected signal is a faster log-linear decrease for the increasing schedule, particularly for larger H, without KL explosion. Repeat with a two-layer MLP policy and learned GAE; success means lower regret at a fixed rollout budget and no collapse in entropy or numerical instability.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.