Occupancy-Regularized Policy Optimization

• ORPO is a reinforcement learning framework that regularizes the global state-action occupancy distribution to mitigate reward hacking and ensure robust performance.
• It employs divergence measures like chi-squared and KL to constrain occupancy discrepancies, enhancing safety and adaptability across dynamics shifts.
• The method integrates density ratio estimation with policy gradients, demonstrating empirical improvements in safe RL and transfer learning tasks.

Occupancy-Regularized Policy Optimization (ORPO) is a class of reinforcement learning (RL) algorithms in which the policy optimization objective is augmented or constrained using global divergences between occupancy measures. ORPO is driven by the observation that regularizing the entire state-action visitation distribution—as opposed to local, per-state action probabilities—provides critical robustness benefits: it tightly controls worst-case true reward loss, mitigates reward hacking, and facilitates adaptation across distributional or dynamical shifts. ORPO unifies a range of algorithms including $\chi^2$-regularization for safe RL, state-occupancy regularization in transfer and off-policy RL, and Bregman-divergence-penalized optimal-transport formulations.

1. Mathematical Foundations and Key Concepts

Let $(\mathcal{S},\mathcal{A},P,\gamma)$ denote a discounted infinite-horizon Markov decision process. The occupancy measure of a stationary policy $\pi$ is defined as

$$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$

By construction, $\sum_{s,a} \mu_\pi(s,a) = 1$, and the expected discounted return for reward $r$ is $J(\pi, r) = \sum_{s,a} \mu_\pi(s,a) r(s,a)$. ORPO applies a divergence $D(\mu_\pi\,\|\,\mu_{\text{ref}})$ (for some reference $\mu_{\text{ref}}$, typically induced by a “safe” policy $\pi_{\text{ref}}$ or optimal occupancy) as a regularizer or constraint in policy search.

A canonical choice in recent literature is the $(\mathcal{S},\mathcal{A},P,\gamma)$0-divergence:

$(\mathcal{S},\mathcal{A},P,\gamma)$1

or equivalently, $(\mathcal{S},\mathcal{A},P,\gamma)$2 (Laidlaw et al., 2024).

Alternative formulations, such as KL-regularization over state marginals or optimal-transport divergences, appear in robust transfer RL and dynamics shift adaptation (Xue et al., 2023, Givchi et al., 2021). These variants differ in the choice of divergence and marginalization but are unified by the principle of shaping the global visitation distribution of the learned policy.

2. Motivation: Reward Hacking, Robustness, and Dynamics Shift

ORPO addresses several fundamental challenges in reinforcement learning:

• Reward hacking: In RL for complex objectives, proxy rewards $(\mathcal{S},\mathcal{A},P,\gamma)$3 are typically used. Optimizing $(\mathcal{S},\mathcal{A},P,\gamma)$4 without global distributional regularization can yield policies that exploit statistical artifacts, producing high proxy return but low true return $(\mathcal{S},\mathcal{A},P,\gamma)$5. This occurs when the state-action regions exploited by $(\mathcal{S},\mathcal{A},P,\gamma)$6 are not heavily visited by a reference (safe) policy $(\mathcal{S},\mathcal{A},P,\gamma)$7, causing breakdown in the correlation between proxy and true reward (Laidlaw et al., 2024).
• Distributional and dynamics shifts: ORPO is also principled for robust transfer or data reuse where RL data are collected under multiple dynamics models (e.g., varying physical parameters). Across such “homomorphous” MDPs, optimal policies often induce similar occupancy measures, even if action choices diverge substantially. Regularizing toward the global or cross-dynamics optimal occupancy fosters adaptation and reuse (Xue et al., 2023).
• Marginal shaping: In problems with prescribed or safety-critical distributions over states or actions, ORPO enables rigorous enforcement or penalization of both state and action marginals. The framework supports both hard constraints and soft penalties, leveraging Bregman divergences and optimal transport relaxations (Givchi et al., 2021).

3. Core ORPO Objective and Algorithmic Implementation

The general ORPO objective augments the standard RL reward maximization with a penalty or constraint on the discrepancy between the learned policy’s occupancy $(\mathcal{S},\mathcal{A},P,\gamma)$8 and a target reference measure $(\mathcal{S},\mathcal{A},P,\gamma)$9:

$\pi$0

where $\pi$1 may denote $\pi$2, KL, or another divergence.

Practical estimation: Since $\pi$3 is not known in closed form, density ratio estimation techniques are used. For $\pi$4-ORPO, a discriminator $\pi$5 is trained to satisfy $\pi$6 via a loss

$\pi$7

(Laidlaw et al., 2024).

Policy gradient integration: The total policy-gradient update is given by

$\pi$8

where the factor 2 arises by differentiating the penalty $\pi$9 with respect to $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$0.

Pseudocode summary (Laidlaw et al., 2024):

Step	Description
1	Collect trajectories from both current policy $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$1 and $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$2
2	Update discriminator $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$3 to estimate density ratios
3	Estimate per-sample occupancy penalty $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$4
4	Policy gradient step maximizing advantage under proxy reward, subtracting $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$5-weighted $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$6 penalty
5	(Optional) Value-function update using penalty-augmented rewards

Variants targeting state-only regularization (Xue et al., 2023) replace $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$7 with $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$8, and employ a [GAN-style] classifier to estimate $$\mu_\pi(s,a) = (1-\gamma)\sum_{t=0}^\infty \gamma^t\Pr_\pi(s_t=s, a_t=a).$$9, where $\sum_{s,a} \mu_\pi(s,a) = 1$0 is the learned cross-dynamics optimal state occupancy.

4. Theoretical Guarantees and Optimality Properties

ORPO admits sharp theoretical guarantees regarding its ability to control performance degradation and induce desired behaviors.

• Worst-case reward gap: For bounded true reward $\sum_{s,a} \mu_\pi(s,a) = 1$1, the return difference is tightly bounded:

$\sum_{s,a} \mu_\pi(s,a) = 1$2

This inequality is tight and holds even under worst-case proxy alignment, ensuring no catastrophic “reward hacking” as long as the $\sum_{s,a} \mu_\pi(s,a) = 1$3 penalty is controlled (Laidlaw et al., 2024).

• Occupancy versus action-distribution regularization: Per-state action Kullback-Leibler regularizers, as in standard RLHF or “safe” RL, provide no such guarantee. It is possible for the KL to remain small while the occupancy measure diverges dramatically, resulting in severe true-return degradation. Thus, KL-regularization is not a predictive or robust safeguard in complex MDPs with cascading effects (Laidlaw et al., 2024).
• Transfer RB for homomorphous MDPs: If all MDPs in a domain share the same reachability graph (homomorphous class), and the per-state KL divergence between current and optimal occupancy is bounded, then the sub-optimality gap enjoys a lower-bound of the form

$\sum_{s,a} \mu_\pi(s,a) = 1$4

where $\sum_{s,a} \mu_\pi(s,a) = 1$5 are Lipschitz constants, $\sum_{s,a} \mu_\pi(s,a) = 1$6 bounds the dynamics shift, and $\sum_{s,a} \mu_\pi(s,a) = 1$7 bounds the occupancy divergence (Xue et al., 2023).

• Convergence: In Bregman-divergence ORPO via Dykstra’s algorithm, the primal iterates converge to the unique solution under standard convexity and feasibility assumptions. With vanishing regularization, value monotonicity and global optimality are restored (Givchi et al., 2021).

5. Connections and Variants: State-Regularized and OT-Based ORPO

ORPO encompasses various methodological instantiations, including:

• State-regularized policy optimization (SRPO) (Xue et al., 2023): Generalizes occupancy regularization to focus solely on state distributions, facilitating adaptation to environmental shifts where optimal action policies may diverge but state visitation patterns persist. SRPO estimates target occupancy via “real vs. fake” state classification, and incorporates a $\sum_{s,a} \mu_\pi(s,a) = 1$8-density ratio penalty into the reward. This approach improves sample efficiency and lower-bound performance in both online and offline RL with shifting dynamics.
• Distributionally-constrained policy optimization via unbalanced optimal transport (Givchi et al., 2021): Formulates ORPO as an optimal-transport problem, using Bregman divergences to penalize deviations from both state and action marginals. Dykstra projection (cyclic Bregman projections) enables efficient solution. In large-scale settings, actor-critic algorithms are derived leveraging dual representations and off-policy samples.

Variant	Regularization	Reference Distribution	Primary Application
$\sum_{s,a} \mu_\pi(s,a) = 1$9-ORPO	$r$0 over $r$1-occupancy	Safe policy	Reward hacking mitigation, RLHF
SRPO	KL over state occupancy	Cross-dynamics optimum	Dynamics shift, transfer RL
OT-ORPO	Generic Bregman (state/action)	Prescribed marginals	Structured occupancy shaping

6. Empirical Validation and Observed Properties

Experiments on benchmark tasks reinforce ORPO’s benefits relative to conventional regularization:

• On reward-hacking MDPs and RLHF-style language modeling, ORPO achieves near-safe policy true return while allowing proxy return improvement. KL-regularized methods frequently suffer from either under- or over-regularization, leading to degraded true return or learning failures (Laidlaw et al., 2024).
• In dynamics-shifted MuJoCo domains, SRPO-augmented algorithms outperform pure context-based or baseline methods, demonstrating higher data efficiency and robust transfer as environmental variability increases. The SRPO regularizer is robust to hyperparameter selection and “plug-in” compatible with standard actor-critic methods (Xue et al., 2023).
• Unbalanced-OT ORPO enables precise control over both state and action distributions in tabular tasks, facilitating exact behavioral shaping that is unattainable with standard entropy or per-state KL regularizers (Givchi et al., 2021).

A key empirical observation is that the global occupancy divergence, particularly $r$2, correlates tightly with true return loss, whereas per-state KL shows little predictive value for off-distribution return.

7. Extensions, Limitations, and Current Directions

ORPO’s theoretical and empirical strengths position it as a foundational methodology for robust, safe, and transferable reinforcement learning. However, practical considerations remain:

• Density ratio estimation is central to all ORPO instantiations. Discriminator design and training stability are critical for accurate penalty estimation, especially in high-dimensional domains.
• The choice of divergence (e.g., $r$3, KL, Bregman class) and the scope of regularization (state vs. state-action) should be matched to task structure and transfer objectives.
• In RLHF and reward-misalignment applications, ORPO is comparatively resistant to subtle reward hacking, but tuning the regularization weight $r$4 remains essential to strike an optimal bias-variance tradeoff.
• Ongoing research explores scalable discriminators, nonparametric density ratio estimation, and integration with large-scale off-policy and model-based RL backends.

ORPO represents a fundamental paradigm shift from local to global distributional regularization, with broad implications for safe RL, reward alignment, and data-efficient learning under uncertainty (Laidlaw et al., 2024, Xue et al., 2023, Givchi et al., 2021).