State Density Weighted OBD

• State Density Weighted (SDW) OBD is a method for generating compact synthetic datasets in offline reinforcement learning by adaptively reweighting states based on empirical density.
• It employs a bi-level optimization framework integrating behavioral cloning and density estimation techniques, such as Masked Autoregressive Flow, to target underrepresented states.
• Empirical results demonstrate that SDW OBD outperforms traditional methods, especially in low state diversity scenarios, leading to improved downstream policy learning.

State Density Weighted (SDW) OBD refers to "State Density Weighted Offline Behavior Distillation," a method for producing compact synthetic datasets for offline reinforcement learning (RL) by explicitly reweighting the distillation objective according to the empirical state density. SDW OBD was introduced to address the limitations of classical Offline Behavior Distillation (OBD), which can misalign synthetic dataset coverage relative to downstream policy learning, particularly when the original dataset exhibits low state diversity or contains pivotal states unevenly represented. SDW OBD adaptively upweights rare (low-density) states in the distillation process, resulting in synthetic datasets that yield improved policy performance in subsequent behavioral cloning. The theoretical motivation, algorithmic details, and empirical evidence supporting SDW OBD are discussed below (Lei et al., 7 Dec 2025).

1. Offline Behavior Distillation: Bi-Level Structure

Offline Behavior Distillation (OBD) compresses large offline RL datasets into small synthetic sets $\mathcal{D}_{\mathrm{syn}}$ suitable for efficient policy learning. OBD is formulated as a bi-level optimization problem:

• Inner Loop (Behavioral Cloning):

$$\theta^*(\mathcal{D}) = \arg\min_\theta \ell^{\mathrm{BC}}(\theta, \mathcal{D})$$

where $\ell^{\mathrm{BC}}$ is the behavioral cloning loss, typically mean-squared error over actions.

• Outer Loop (Distillation):

$$\mathcal{D}_{\mathrm{syn}}^* = \arg\min_{\mathcal{D}} \mathcal{H}(\pi_{\theta^*(\mathcal{D})}, \mathcal{D}_{\mathrm{real}})$$

Here, $\mathcal{H}$ quantifies how well the policy $\pi$ trained on synthetic data matches the "real" offline dataset $\mathcal{D}_{\mathrm{real}}$ (using metric choices such as Policy-Based Cloning: $\mathcal{H}_{\mathrm{PBC}}$, or Action-Value Weighted PBC: $\mathcal{H}_{\mathrm{Av-PBC}}$).

• Gradient Update: The synthetic set $\mathcal{D}_{\mathrm{syn}}$ is updated by backpropagating through the inner behavioral cloning fit using BPTT.

2. Empirical State Density Estimation

To reweight the distillation objective, it is necessary to estimate the empirical state density $d(s)$ for each state $s$ in $\mathcal{D}_{\mathrm{real}}$:

• Technique: Masked Autoregressive Flow (MAF) is employed to model $d(s) \approx \hat{p}(s)$, delivering tractable and differentiable log-densities for all states.
• Alternative Methods: Kernel density estimation or other density estimators may be substituted without affecting the framework's generality.

The density model provides weights which prioritize rare states and facilitate calculation of per-state contributions in subsequent optimization.

3. State Density Weighted Objective

SDW OBD modifies the outer distillation objective by introducing a density-based weight $w(s)$ for each state. For hyperparameter $\tau \geq 0$:

$w(s) = d(s)^{-\tau}$

The SDW distillation loss generalizes Av-PBC:

$$\mathcal{H}_{\mathrm{SDW}}(\pi; \mathcal{D}_{\mathrm{real}}) = \mathbb{E}_{(s,a) \sim \mathcal{D}_\text{real}} \left[ q_{\pi^*}(s, a) \cdot d(s)^{-\tau} \cdot \|\pi(s) - a\|^2 \right]$$

where $q_{\pi^*}(s, a)$ denotes the action-value of the expert policy. Setting $\tau = 0$ recovers the non-weighted Av-PBC objective; positive $\tau$ upweights low-density (rare) states.

4. Theoretical Foundation: Pivotal and Surrounding Errors

SDW OBD is motivated by a precise analysis of policy error decomposed into two terms:

• Pivotal Error ($\epsilon$): Associated with states visited by the expert policy $\pi^*$, denoted $S_e = \{s \mid d_{\pi^*}(s) > 0\}$:

$$\epsilon = \mathcal{E}_e(\hat{\pi}) = \mathbb{E}_{s \sim d_{\pi^*}} \left[ \sum_a |\hat{\pi}(a | s) - \pi^*(a | s)| \right]$$

• Surrounding Error ($\epsilon_\mu$): Refers to the probability the learned policy persists in states never visited by $\pi^*$, $S_\mu = \{s \mid d_{\pi^*}(s) = 0\}$.
• Suboptimality Bounds:
  • Expert-Only: If $\epsilon \leq \epsilon_0$, then $$|J(\pi^*) - J(\hat{\pi})| \leq \epsilon T^2 R_{\max}$$.
  • Main Bound with Surrounding Error: If errors are $\epsilon, \epsilon_\mu$, and $\pi^*$ meets mild visitation assumptions, then $$|J(\pi^*) - J(\hat{\pi})| \leq (\epsilon_\mu T + 3)\epsilon T R_{\max}$$.

This analysis demonstrates that, when pivotal error $\epsilon$ remains non-negligible (as occurs in bi-level OBD), surrounding error $\epsilon_\mu$ grows in importance. Thus, enhancing coverage—especially of sparse regions—directly impacts policy performance.

5. SDW OBD Algorithm: Workflow and Pseudocode

The SDW OBD approach operates as follows:

• Inputs: Offline dataset $\mathcal{D}_{\text{off}}$, synthetic set size $N_{\text{syn}}$, density exponent $\tau$.
• Steps:
  • a. Inner loop: train $\pi_\theta$ on $\mathcal{D}_{\text{syn}}$ for $T_\text{in}$ gradient steps.
  • b. Sample minibatch $B$ from $\mathcal{D}_{\text{off}}$.
  • c. Compute SDW loss:

  $$\mathcal{H} = (1 / |B|) \sum_{(s,a) \in B} [q_{\pi^*}(s,a) \, d(s)^{-\tau} \| \pi_\theta(s) - a \|^2]$$ - d. Update synthetic data:

  $$\mathcal{D}_{\text{syn}} \leftarrow \mathcal{D}_{\text{syn}} - \alpha_1 \nabla_{\mathcal{D}_{\text{syn}}} \mathcal{H}$$

• Output: Final synthetic set $\mathcal{D}_{\text{syn}}$.

This algorithm explicitly targets rare states in $\mathcal{D}_{\text{off}}$, mitigating surrounding error.

6. Empirical Evaluation and Results

SDW OBD demonstrates superior performance in benchmarks using D4RL datasets (MuJoCo: HalfCheetah, Hopper, Walker2D; Medium and Medium-Expert qualities):

Method	HalfC-M	HalfC-M-E	Hopper-M	Hopper-M-E	Walker-M	Walker-M-E	Avg
Rand($\mathcal{D}_{\text{off}}$)	1.8	2.0	19.2	11.6	4.9	6.7	7.7
Rand($\mathcal{D}_{\text{real}}$)	5.9	7.8	29.1	27.1	17.1	17.8	17.5
DBC	28.2	29.0	37.8	31.1	29.3	11.7	27.9
PBC	30.9	20.5	25.1	33.4	33.2	34.0	29.5
Av-PBC	36.9	22.0	32.5	38.7	39.5	42.1	35.3
SDW ($\tau=0.1$)	39.5	25.0	38.4	42.6	42.5	44.6	38.8

• Key Outcomes:
  • SDW achieves a mean improvement of $3.5\%$ over Av-PBC across all environments.
  • Highest gains observed for datasets with lowest state diversity.
  • Performance holds for $\tau$ values in the range $0.05$–$0.15$.
  • SDW-synthesized datasets generalize well across a range of policy architectures and optimizers, consistently outperforming non-SDW methods.

7. Implications and Conclusions

SDW OBD directly mitigates surrounding error by redistributing distillation focus toward underrepresented regions of the state space. The method's theoretical validity is substantiated by analysis of suboptimality bounds, and its empirical efficacy is consistently confirmed across multiple RL tasks and data regimes. A plausible implication is that future compact dataset distillation frameworks may benefit from explicit diversity-aware weighting, especially in applications where data coverage is sparse or biased. SDW OBD enables high-quality policy learning with significantly reduced data requirements in offline RL (Lei et al., 7 Dec 2025).