Noise-Whitened Trajectory-KL Policy Regularization

1. Exact integration point: use this as an auxiliary loss during policy training. Let $\pi_\theta(u\mid x)$ be the trainable policy and $\pi_0(u\mid x)$ a frozen behavior policy, offline policy, or reference-dynamics controller. Estimate the environment drift as $f_\theta(x)=b(x)+B(x)\bar u_\theta(x)$ and the reference drift as $f_0(x)=b_0(x)+B(x)\bar u_0(x)$; for a learned world model, learn both drifts from transition data. 2. Pseudocode: sample a batch of transitions or short rollouts; compute $d_i=f_\theta(x_i)-f_0(x_i)$ and $a_i=\sigma(x_i)\sigma(x_i)^\top$; solve $a_i z_i=d_i$ with damping $a_i+\epsilon I$; set $L_{\mathrm{trajKL}}=\frac{1}{2N}\sum_i d_i^\top z_i\Delta t$; optimize $L=L_{\mathrm{task}}+\lambda L_{\mathrm{trajKL}}$. If only actions are available and $f=b+Bu$, use $d_i=B_i(\bar u_\theta-\bar u_0)$ rather than an action-norm penalty. 3. Computed from the paper: the quadratic drift mismatch and inverse-covariance weighting. Estimated empirically: $B$, $b$, $b_0$, $\sigma$, and the finite-horizon KL from rollouts. Clip or damp inverse covariance eigenvalues for numerical stability. 4. First cheap experiment: use Pendulum-v1 or a linear-Gaussian control task, comparing PPO or SAC with ordinary behavior regularization against this penalty, using a noisy reference controller trained on half the data. Sweep $\lambda$. The quantitative prediction is that measured trajectory KL decreases monotonically with $\lambda$, while in the local linear-quadratic regime the control deviation satisfies $\|u_\lambda-u_0\|=O(1/\lambda)$ along directions where $B^\top a^\dagger B$ is positive; the fitted local control curvature should match $R+\lambda B^\top a^\dagger B$ within approximately 20%. Increasing environmental noise in a drift direction should reduce its regularization weight in proportion to the corresponding eigenvalue of $a^{-1}$.