Unbiased Policy Divergence Estimation for Regularized Optimization

In reinforcement learning and policy optimization algorithms, it is crucial to prevent a learned policy from deviating too drastically from a stable reference distribution. This is achieved by incorporating a divergence penalty into the objective function, which regularizes the optimization process and ensures smooth, controlled policy updates.

The unbiased divergence estimator computes a per-sample measure of how different the current policy distribution (\pi_\theta) is from a reference policy distribution (\pi_{\text{ref}}). This measure is particularly important in algorithms like Group Relative Policy Optimization (GRPO) and Proximal Policy Optimization (PPO), where staying close to the reference policy prevents catastrophic forgetting and maintains training stability.

Mathematical Formulation:

For each sample, the divergence estimator is computed using the following formula:

$$D_i = r_i - \log(r_i) - 1$$

where the probability ratio (r_i) is defined as:

$$r_i = \frac{\pi_{\text{ref}}(a_i)}{\pi_\theta(a_i)}$$

Here:

• (\pi_\theta(a_i)) is the probability assigned by the current policy to action (a_i)
• (\pi_{\text{ref}}(a_i)) is the probability assigned by the reference policy to the same action

Intuition Behind the Formula:

This estimator has several elegant mathematical properties:

1. Non-negative: The divergence (D_i \geq 0) for all valid probability ratios
2. Zero at equality: When (\pi_\theta = \pi_{\text{ref}}), we get (r = 1), and (D = 1 - \log(1) - 1 = 0)
3. Convex penalty: The function (f(r) = r - \log(r) - 1) is convex, penalizing both over-confident and under-confident deviations

Your Task:

Implement a function that computes the per-sample policy divergence values. For each element in the input arrays, calculate the divergence using the formula above and return the results rounded to 4 decimal places.