Is Density vs. Feature Coverage That Different?

Coverage is a fundamental concept in offline RL and measures how data distribution contains information about another policy. The most basic form of coverage is measured by the magnitude of density ratios, but it is well known that this can be relaxed to feature coverage with linear reward/value functions. Having a lenient definition of coverage can be beneficial in policy evaluation and learning: for evaluation, a target policy may be poorly covered by data in terms of density ratio but well covered in feature, and whether the analysis takes advantage of the latter will determine whether a guarantee can be established. For policy learning, pessimistic approaches allow us to compete with the best covered policy, and a lenient notion of coverage (such as feature coverage) effectively searches over a significantly larger space of policies.

So for a long time, my understanding is that feature coverage is a significant improvement and generally “black-and-white” different from density coverage, until I went down a rabbit hole recently that shook my belief: in the policy-learning example below, whenever a policy has good feature coverage, an equally performant policy can be found (up to “minor” issues) that has good density coverage.

… wait what?

---

Setup and Background

For simplicity consider linear bandits (spoiler: keep this in mind). There is an action space $\mathcal{X} \subset \mathbb{R}^d$. Upon taking action $x \in \mathcal{X}$, the reward is $r = x^\top \theta^\star + \epsilon$ with zero-mean bounded noise $\epsilon$. A policy is simply a distribution over $\mathcal{X}$. Suppose our data is sampled from policy $q \in \Delta(\mathcal{X})$. Let $\mu_q:= \mathbb{E}_q[X]$ and $\Sigma_q := \mathbb{E}_q[XX^\top]$.

For another policy under consideration, $p\in\Delta(\mathcal{X})$, the data’s feature coverage over $p$ is

$$C_p := \|\mu_p\|_{\Sigma_q^{-1}} = \sqrt{\mu_p^\top \Sigma_q^{-1} \mu_p}. \qquad (1)$$

This quantity shows up if we fit the linear reward model (let $\hat\theta$ be the learned coefficient) from data (sampled from $q$), and then use it to predict the expected reward of $p$ as $\mathbb{E}_p[X^\top \hat\theta]$. The guarantee looks like

$$|\mathbb{E}_p[r] - \mathbb{E}_p[X^\top \hat\theta]| \le C_p \cdot \textrm{statistical error on } q. \qquad (2)$$

For policy learning, by adding constraints or regularization based on $C_p$, we can learn a policy that competes with the best policy with low feature coverage.

---

The Example

And here comes the example: consider a policy $p$ with small $C_p$. We will construct a different policy $p'$ that has good density coverage from $q$ (with some caveats). The construction is pretty straightforward: let $w(x) = p'(x)/q(x)$ be

$$w(x) = \mu_p^\top \Sigma_q^{-1} x.$$

The caveat is that $p'(x) = w(x) q(x)$ may not be a valid distribution due to potentially negative weights and normalization issues, which we will discuss shortly. If we ignore these issues, we can verify that

$$\mathbb{E}_{p'}[X] = \mathbb{E}_{q}[(\mu_p^\top \Sigma_q^{-1} X) X] = \mathbb{E}_{q}[X X^\top \Sigma_q^{-1} \mu_p] = \mu_p,$$

and the size of the density ratio $w = p'/q$ measured in $\chi^2$ is:

$$\mathbb{E}_{q}[(p'/q)^2] = \mathbb{E}_q[\mu_p^\top \Sigma_q^{-1} X X^\top \Sigma_q^{-1} \mu_p] = C_p^2.$$

So $p'$ produces exactly the same mean vector $\mu_{p'}$ as $\mu_p$—thus $p'$ and $p$ have equal expected reward—and the density coverage of $p'$ is exactly equal to the feature coverage of $p$.

---

Addressing the Caveats

For me the above construction already suffices: I have the impression that there should be cases where $\mu_p$ should be impossible to reproduce however we reweight $q$, even if we allow negative and unnormalized weights (turns out I am right! see the next section), which is why I described the caveats as “minor” at the beginning. That said, there caveats can be serious out of the scope of this post, and ChatGPT (ahem) provides some interesting ideas to address it. (This part is orthogonal to the main point of the post, so feel free to skip!)

The negative weights are harder to address, as truncation can introduce errors in the mean vector. The normalization issue has an interesting solution: we may use a “difference” variant of feature coverage, defined as

$$\tilde{C}_p:= \sqrt{(\mu_p - \mu_q)^\top \tilde{\Sigma}_q^{-1} (\mu_p-\mu_q)},$$

where $\tilde{\Sigma}_q$ is the covariance matrix instead of second moment. One can control $\tilde{C}_p$ by $C_p$ as long as $\|\mu_q\|_{\Sigma_q^{-1}} < 1$, and obtain guarantees involving $\tilde{C}_p$ by centering $X$ around $\mu_q$ and $r$ around $\mathbb{E}_q[r]$ before doing regression. Then, a weight that always normalizes is (it may still suffer from negativity issues):

$$\tilde{w}(x) = (\mu_p - \mu_q)^\top \tilde{\Sigma}_q^{-1} (x - \mu_q),$$

where $\mathbb{E}_q[\tilde{w}] = 1$ and $\mathbb{E}_q[\tilde{w}^2] = \tilde{C}_p$.

---

Reconciling the Dilemma

We have two contradicting claims:

1. Sometimes $\mu_p$ can never be reproduced by reweighting $q$, even with negative and unnormalized weights.

2. $\mu_p$ can always be reproduced by the weights given above.

What gives?

The crux is in the simplified setting of linear bandits we adopted at the beginning: our analysis is technically correct, but Claim #2 only holds in linear bandits when there is no context, whereas Claim #1 is true in contextual bandits and MDPs, so there is no contradiction. For readers familiar with MDPs, the crucial observation that resolves the dilemma is that an arbitrary reweighting of data distribution is generally illegal, as any distribution that can be realized as the occupancy of a policy must respect the transition dynamics.

In fact, the mechanism that invalidates arbitrary reweighting in MDPs is already present in the simpler model of contextual bandits, where a classic example will make it instantly clear what’s going on: consider a contextual bandit (let’s say context is $s$, action $a$, linear feature $\phi(s,a)$) where behavior policy $\pi_q: \mathcal{S} \to \Delta(\mathcal{A})$ is deterministic. If context has enough variations, it is still possible that $\Sigma_q = \mathbb{E}_{\pi_q}[\phi\phi^\top]$ is full-rank, which implies feature coverage over all policies. However, $\pi_q$ clearly provides no density coverage over any policy other than itself. If we construct the weight from the earlier analysis, it corresponds to altering the context distribution, which is not allowed.¹

So that’s it! I guess the lesson is that contextless and contextual settings can sometimes be very different. When thinking about RL in LLMs, I frequently simplify the problem to ignore prompt when brainstorming, assuming that the contextual case will be similar. This is not the first time I run into problems with such simplifications—I once asked in a talk why pass@k rate is not directly determined by pass@1 via order statistics, only to realize later that this is not true due to heterogeneity of performance across prompts.

Footnote Remarks