Matrix Estimation for Offline Evaluation in Reinforcement Learning with Low-Rank Structure
Abstract
We consider offline Reinforcement Learning (RL), where the agent does not interact with the environment and must rely on offline data collected using a behavior policy.
Previous works provide policy evaluation guarantees when the target policy to be evaluated is covered by the behavior policy, that is, state-action pairs visited by the target policy must also be visited by the behavior policy.
We show that when the MDP has a latent low-rank structure, this coverage condition can be relaxed.
Building on the connection to weighted matrix completion with nonuniform observations, we propose an offline policy evaluation algorithm that leverages the low-rank structure to estimate the values of uncovered state-action pairs.
Our algorithm does not require a known feature representation, and our finite-sample error bound involves a novel discrepancy measure quantifying the discrepancy between the behavior and target policies in the spectral space.
We provide concrete examples where our algorithm achieves accurate estimation while existing coverage conditions are not satisfied.
Building on the above evaluation algorithm, we further design an offline policy optimization algorithm and provide non-asymptotic performance guarantees.
Introduction
Reinforcement Learning (RL) has achieved significant empirical success in the online setting, where the agent continuously interacts with the environment to collect data and improve its performance.
However, online exploration is costly and risky in many applications, such as healthcare [4] and autonomous driving [16], in which case it is preferable to learn from a pre-collected observational dataset from doctors or human drivers using their own policies.
Due to lack of on-policy interaction with the environment, offline RL faces the fundamental challenge of distribution shift [7].
A standard approach for handling distribution shift is importance sampling [12, 11].
More sophisticated approaches have been proposed to alleviate the high variance of importance sampling [2, 21].
Recent works [17, 10, 23] consider estimating the state marginal importance ratio, a more tractable problem.
Existing work on offline RL requires the dataset to have sufficient coverage.
A standard measure for coverage is the concentrability coefficient [19]: C π = max s,a d π ( s,a ) ρ ( s,a ) , which is the ratio between the state-action occupancy measure of a policy π of interest and the (empirical) occupancy measure ρ of the behavior policy generating the offline dataset.
However, this can be restrictive as the support of ρ must contain that of d π in order for C π to be finite.
Earlier work such as the Fitted Q-iteration (FQI) algorithm [9] requires full coverage, i.e. C π < ∞ for all policies π .
More recent works [19, 13, 8] requires a more relaxed partial coverage condition C π ∗ < ∞ for the optimal policy π ∗ .
Partial coverage is still a fairly strong requirement: the behavior policy must visit every state the optimal policy would visit, and take every action the optimal policy would take.
In this paper, we seek to relax the coverage condition for offline policy evaluation in settings where the Markov decision process (MDP) has a latent low-rank structure.
Similarly to [15, 14], we view the Q function as a matrix and exploit its low-rank structure to infer the entries that were not observed in the offline data.
Unlike typical results from the low-rank matrix completion literature, our setting requires completing the matrix under non-uniform sampling, as in [3, 6]; moreover, the error is evaluated under a different distribution or weighted norm, leading to the fundamental challenge of distribution shift.
By leveraging techniques from weighted and non-uniform matrix completion, we develop a new offline policy evaluation algorithm, which alternates between Q iteration and matrix estimation.
For both the infinite and finite sample settings, we show that the evaluation error can be bounded in terms of a novel discrepancy measure between the behavior and target policies.
In contrast to the standard concentrability coefficient, our discrepancy measure may remain finite even when the behavior policy does not cover the support of the target policy.
We present a concrete example where the concentrability coefficient is infinite but our method achieves a meaningful error bound.
Building on the above evaluation algorithm, we further design an offline policy optimization algorithm with provable performance guarantees.
Problem Setup
Consider an MDP M =
( S , A , H, P, r, µ 1 ) with finite state space S , finite action space A , horizon
H , transition kernel P = { P t } t ∈
[ H ] , bounded reward function r = { r t : S ×
A → [0 , 1] } t ∈ [ H ] , and initial state distribution µ 1 ∈ ∆( S ) .
Let S = |S| and A = |A| .
For each policy π = { π t : S → ∆( A ) } t ∈ [ H ] , the Q function Q πt : S ×
A → R is defined as Q πt ( s, a ) =
E π [ (cid:80) Hi =
t r i ( s i , a i ) | s t = s, a t = a ] , and the total expected reward is J π
= E π
[ (cid:80)
Ht =1 r t ( s t , a t ) | s 1 ∼ µ 1 ] .
Let d πt : S ×
A → [0 , 1] denote the state-action occupancy measure at time t ∈
[ H ] under policy π .
Given a dataset generated by the behavior policy π β , our goal is to estimate J π θ for a target policy π θ .
We assume that the MDP has the following low-rank structure, which implies that for any policy π , its Q function (viewed as an S -by- A matrix) is at most rank d .
Assumption.
For all t , r t ∈ [0 , 1] S × A has rank at most d/ 2 , and P t admits the decomposition
P t ( s ′ | s, a ) = (cid:80) d/ 2 i =1 u t,i ( s ′ , s ) w t,i ( a ) or P t ( s ′ | s, a )
= (cid:80) d/ 2 i
=1 u t,i ( s ) w t,i ( s ′ , a ) , ∀ s ′ , s, a.
The above low-rank model is different from the Low-rank MDP model considered in previous works [1, 20].
In Low-rank MDPs, the transition kernel P is assumed to have a factorization of the form P ( s ′ | s, a ) = (cid:80) di =1 u i ( s ′ ) w i ( s, a ) , where the factors u i ( · ) and w i ( · , · ) are unknown. Closely related is the Linear MDP model [5, 22], where the feature maps w i ( · , · ) are known.
In these models, the low-rank/linear structures are with respect to the relationship between the originating state-action pair ( s, a ) and the destination state s ′ ; they do not imply that Q function is low-rank when viewed as a matrix.
In contrast, our model stipulates that the transition kernel can be factorized either between (i) a and ( s, s ′ ) or (ii) s and ( s ′ , a ) , both of which imply a low dimensional relationship between the current state s and the action a taken at that state, resulting in a low-rank Q function.
For a matrix M , let ∥ M ∥ ∗ denote its nuclear norm (sum of singular values), ∥ M ∥ op its operator norm (maximum singular value), ∥ M ∥ ∞ =
max i,j | M ij | its entrywise ℓ ∞ norm, and supp( M ) =
{ ( i, j ) :
M ij ̸ = 0 } its support.
The indicator matrix 1 M is a binary matrix encoding the position of the support of M .
The entrywise product between two matrices M and M ′ is denoted by M ◦ M ′ .
We propose a novel discrepancy measure defined below, and show that it can replace the role of the concentrability coefficient in our error bound under the low-rank assumption.
Definition 1 (Operator discrepancy) .
The operator discrepancy between two probability distributions p, q ∈ ∆( S × A ) is defined as
[Equation]
Note that Dis( p ∥ q ) ≤ ∥ p − q ∥ op is always finite, and Dis( p ∥ q ) =
if and only if supp( q ) ⊆ supp( p ) .
To gain intuition for the above definition, assume the minimizer in (1) is g ∗ .
By generalized
Hölder’s inequality, we have
[Equation]
If the nonzero singular values of M are of the same scale, then the RHS of (2) is of order Dis( p ∥ q ) · rank ( M ) .
Therefore, Dis( p ∥ q ) measures the distribution shift between p and q in terms of preserving the expectation of low-rank matrices.
Note that Dis( p ∥ q ) only depends on the support of p : if supp( p ) =
supp( p ′ ) , then Dis( p ∥ q ) =
Dis( p ′ ∥
q ) for all q .
Moreover, thanks to the minimization in the definition (1), Dis( p ∥ q ) can be significantly smaller than ∥ p − q ∥ op .
For instance, if p is the uniform distribution on S × A , then g ∗ =
q and hence Dis( p ∥ q ) = 0 for all q .
The operator discrepancy shares similarity with the parameter Λ in [6].
Note that the operator discrepancy is not symmetric.
The error bounds for our proposed off-policy evaluation algorithm will be a function of Dis( d π β t ∥
d π θ t ) , even when one is given infinite samples from the behavior policy.
The operator discrepancy only depends on the support of the behavior policy and not the exact distribution, which is expected under the infinite sample setting.
As such, the operator discrepancy highlights the inherent error induced by distribution shift.
In the finite sample setting, our error will have additional terms which represent the empirical approximation error on the support of the behavior policy.
Algorithm
Our algorithm alternates between two steps: applying Q-value iteration on the support of d π β t , and using low-rank matrix completion to infer the Q values off support.
The algorithm takes as input an offline dataset D
= { ( s kt , a kt , r kt ) } t ∈
[ H ] ,k ∈
[ K ] , which contains K independent trajectories generated from the behavior policy π β .
Over state-action pairs in the support of π β , we use the data to construct unbiased empirical estimates of the immediate reward, transition kernel and occupancy measure of the behavior policy, denoted by (cid:98) r t , (cid:98) P t and (cid:98) d π β t , respectively.
Let (cid:98) B
π θ t denote the target policy’s empirical Bellman operator, which is given by
[Equation] [Equation]
) for all f : S ×
A → R .
Note that we can evaluate ( (cid:98) B πt f )( s, a ) only over ( s, a ) ∈ supp( (cid:98) d π β t ) .
With these notations, our algorithm is given below. Algorithm
Data: dataset D , π Line 3 of Algorithm 1 involves a weight matrix ρ t with supp( ρ t ) ⊆ supp( (cid:98) d π β t ) .
One may simply take ρ t to be (cid:98) d π β t , or use other weights.
Our theorems below quantify the performance of different ρ t .
Line 4 of the algorithm uses a matrix estimation subroutine ME ( · ) , which is given by the following nuclear norm minimization program with L t :
= H − t + 1 and a tuning parameter κ ≥ 0 :
[Equation]
Analysis
We present evaluation error bounds under both the infinite sample setting K → ∞ and the finite sample setting K < ∞ .
Define the population Bellman operator B π θ t , which is given by equation (3) with (cid:98) r t and (cid:98) P t replaced by r t and P t .
Define the matrix Y t ∈ R S ×
A via Y t ( s, a ) =
( B π θ t (cid:98) Q π θ t +1 )
( s, a ) , which is the population version of Z t computed in Algorithm 1.
In the infinite sample setting, we have (cid:98) d π β t ( s, a ) → d π β t ( s, a ) , (cid:98) r t ( s, a ) → r t ( s, a ) and (cid:98) P t ( s, a ) →
P t ( s, a ) for all ( s, a ) ∈ supp( d π β t ) .
Consequently, both (cid:98) B π θ t and Z t converge to their population versions B π θ t and Y t , respectively.
We set κ = 0 in equation (4), which implies that [-
-]1 ρ t ◦ M =ρ t ◦ Y t .
We have the following guarantee.
Theorem 1 (Infinite samples) .
In the infinite sample setting, under Algorithm 1 with ρ t = d π β t and κ = 0 , the output estimator (cid:98) J satisfies
[Equation]
Next consider the setting with a finite dataset D = { ( s kt , a kt , r kt ) } t ∈
[ H ] ,k ∈
[ K ] .
Let n t ( s, a ) :
= (cid:80) k ∈ [ K ] 1 ( s kt ,a kt )=( s,a ) be the visitation count of each state-action pair.
Accordingly, the empirical occupancy of π β is given by (cid:98) d π β t ( s, a ) =
n t ( s, a ) /K .
The following guarantee, which generalizes
Theorem 1, holds for any weight matrices ρ
= ( ρ t ) t ∈ [ H ] , for which we define n min ( ρ ) :
= min { n t ( s, a ) : t ∈ [ H ] , ( s, a ) ∈ supp( ρ t ) } .
Theorem 2 (Finite samples) .
Consider the finite sample setting under Algorithm 1 with ρ t satisfying supp( ρ t ) ⊆ supp( (cid:98) d π β t ) and κ = ∥ 1 ρ t ◦ ( Y t − Z t ) ∥ ∞ .
There exists an absolute constant C >such that with probability at least 1 − δ , we have
[Equation]
Remark 1 .
The finite sample error bound in (6) decomposes into infinite sample error and statistical error.
The infinite sample error quantifies the intrinsic hardness of distribution shift.
The statistical error is due to empirical estimation of the reward and transition kernel.
We provide the flexibility of using ρ t different from (cid:98) d π β t , which would allow us to balance the infinite sample error and the statistical error.
If n min ( (cid:98) d π β ) is small, the statistical error can be large.
In this case, one may choose a ρ t that puts zero weight on the pairs ( s, a ) with small count n t ( s, a ) .
The resulting n min ( ρ ) can be significantly larger than n min ( (cid:98) d π β ) , which reduces the statistical error while potentially increasing the infinite sample error.
The proofs of Theorem 1 and 2 are deferred to Appendix A.1 and A.6, respectively.
Our bounds (5) and (6) scale with the operator discrepancy (1), which is always finite, given any behavior policy π β and target policy π θ , in contrast to the concentrability coefficient C π = max s,a d π ( s,a ) ρ ( s,a ) .
Suppose that there exists some ( s, a ) such that ρ t ( s, a ) = [-
-]0 and d π θ t ( s, a ) > 0 .
Then, C π θ = ∞ whereas
Dis( ρ t ∥
d π θ t ) is finite and meaningful.
We present a concrete example in the infinite samples setting showcasing the effectiveness of our algorithm.
Assume S =
A = n .
Consider the simple setting where the transition is uniform over all state-action pairs.
For each s and t , assume π θt ( ·| s ) is supported on m actions, and the locations of these actions are a realization of uniform random sampling over [ n ] .
We assume π βt is generated from the same model independently.
Note that the support of d π β t and d π θ t will be mostly disjoint, making the concentrability coefficient infinite with high probability.
Using Theorem 1, we derive the following error bound, the proof of which is deferred to Appendix A.2.
Corollary 1.
Under the aforementioned setting, there exists an absolute constant C > 0 such that when n ≥ C , with probability at least 1 − 1 n , we have (cid:12)(cid:12) (cid:98) J − J π θ (cid:12)(cid:12) ≤ C (cid:113) dH 3 log( n ) m .
If m satisfies m ≳ dH log n ϵ, then we have | (cid:98) J − J π θ | ≤ ϵH .
Suppose m = n/
In this setting, the behavior and target policies both randomize over half of the actions, but their actions may have little overlap.
Our bound gives | (cid:98) J − J π θ | ≲ (cid:113) dH 3 log n n , which can be vanishingly small when n is large.  Policy Optimization and Guarantee
In this section, we build on our policy evaluation methode to design an offline policy optimization algorithm.
Given a dataset D generated by a behavior policy π β , we can use Algorithm 1 to obtain an value estimate (cid:98) J π for each policy π . We optimize over policies for which the above estimate is reliable. Specifically, we consider the following set of candidate policies [Equation]
where the parameter B ≥ 0 controls how close the occupancy distribution of the candidate policy is to the data distribution, in terms of operator discrepancy defined in (1).
It is easy to see that when policy π satisfies supp(
{+d πt ) ⊆ supp( ρ t ) , we have Dis( d πt ∥ ρ t ) = 
+}
and as a result, π ∈ Π B for any B ≥ 0 .
Consequently, if we take ρ t = (cid:98) d π β t , then Π B includes the policies that are covered by the offline dataset.
In other words, all policies with finite concentrability coefficients are in Π B .
Importantly, when B > 0 , the set Π B contains other policies with infinite concentrability coefficients, as demonstrated in the example at the end of last section.
With a bigger B , the set Π B includes more policies, though we have weaker evaluation guarantees for these policies.
Among all candidate policies in Π B , we maximize the estimated values obtained by Algorithm 1 to get
[Equation]
We present the following guarantee for (cid:98) π , the proof of which can be found in Appendix A.7.
Theorem 3.
Suppose π β ∈ Π B .
We obtain (cid:98) π by solving (7).
There exists an absolute constant C >such that with probability at least 1 − δ , we have
[Equation]
The above bound shows that we are able to find a policy (cid:98) π with a nearly optimal value, compared to other policies in Π B .
How close (cid:98) π is to the optimal policy in Π B depends on how accurately we can evaluate all policies in Π B .
According to Theorem 2, the estimations are accurate if B is small (policies are close to behavior) and n min ( ρ ) is large (dataset is large), which is reflected in the bound (8).
Similarly as before, the two error terms in (8) quantifies the fundamental difficulty of distribution shift and finite sample noise, respectively.
If we let ρ t = (cid:98) d π β t and B = 0 , we are essentially optimizing over all policies with finite concentrability coefficients.
In this case, our bound reduces to
[Equation]
Compared to the PAC bound of CPPO proposed in [19] (Section 5.1), our bound has a better polynomial dependence on the time horizon H .
We propose a novel algorithm for efficient offline evaluation when low-rank structure is present in the
MDP.
Our algorithm is a combination of Q iteration and low-rank matrix estimation, which is easy to implement.
We show that the proposed operator discrepancy measure better captures the difficulty of policy evaluation in the offline setting, compared to the traditional concentrability coefficient. We also combine the evaluation algorithm with policy optimization and provide performance guarantee.
We believe that this work is a first step in exploiting the benefit of low-rank structure in the Q function in offline RL.
In the future, we hope to develop a more efficient policy optimization algorithm, with a better estimation accuracy
A Proofs
Let µ πt : S → [0 , 1] denote the state occupancy measure at time
t ∈
[ H ] under policy π .
A.1 Proof of Theorem
We present two lemmas before analyzing the evaluation error.
The first one analyzes the error incurred at the matrix estimation step.
The proof is deferred to Appendix A.4.
Lemma 1.
For arbitrary real matrices A, B, P, W ∈ R m × n , we have (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:88) i
[Equation]
Remark.
Under the matrix estimation framework, we can interpret matrix P as the sampling pattern and W as the weights for evaluation.
Next, we introduce a lemma decomposing the evaluation error as a summation of the matrix estimation accuracy from future timesteps.
The proof can be found in Appendix A.5.
Lemma 2.
For the Q function and its estimator Q π θ t , (cid:98) Q π θ t ∈ R S
× A , we have
[Equation]
and consequently
[Equation]
Based on Lemma 1 and 2, we derive the following error bound.
For each t ∈ [ H ] and an arbitrary g ∈ ∆( S × A ) with supp( g )
⊆ supp( d π β t ) , we have (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:88) [Table]
[Equation]
where the first step follows from Lemma 1, and the second equality follows from the constraints in (4) with κ = 0 .
Then, combining with the decomposition in Lemma 2, we obtain the desired bound by minimizing over all such g .
A.2 Proof of Corollary
For simplicity, define b := mn .
Since the transition is uniform, the state occupancy µ πt ( · ) is uniform under any policy π , i.e. µ πt ( s ) = 1 n .
By the way the policies are generated, d π θ t ( · , · ) = µ π θ t ( · ) π θt ( ·|· ) ∈
R n × n is supported on mn entries whose locations are realization of random sampling, and on these entries d π θ t ( s, a ) =
1 mn .
Specifically, all d π θ t ( s, a ) are i.i.d.
Bernoulli random variables that take the value 1 with probability b .
The behavior policy π β is generated independently via the same process.
Let M :=
d π θ t − d π β t
and we have
[Equation]
independently across all entries ( i, j ) .
By matrix Bernstein inequality, we obtain the following result, the proof of which is deferred to Appendix A.3.
Lemma 3.
There exists an absolute constant C > 0 such that when n ≥ C , with probability at least 1 − 1 n , we have
[Equation]
Since Y t is at most rank- d and has bounded entries, we have ∥ Y t ∥ ∗ ≤ √ d ∥
Y t ∥ F ≤ n √ dH .
Combining Theorem 1 and (10), we get
[Equation]
where the first upper bound is obtained by plugging d π β t into the objective of (1).
A.3 Proof of Lemma
We apply matrix Berstein’s inequality (Theorem 6.1.1 in [18]).
Let S k
= M ij e i e ⊤ j , for all k ∈ [ n 2 ] .
Since | M ij | ≤ 1 mn , we derive that ∥ S k ∥ op ≤ 1 mn .
We calculate that
[Equation]
As a result, we have
[Equation]
By symmetry, we also have (cid:13)(cid:13)(cid:80) k E
[ S ⊤ k S k ] [Equation] op ≤ 2 n 2 m .
Hence, we get
[Equation]
Letting the RHS be upper bounded by 1 n yields the desired result.
A.4 Proof of Lemma
The proof uses the following result, which holds for any pairs of dual norms.
In this paper, we only consider using ∥·∥ ∗ and ∥·∥ op .
Lemma 4.
For a real matrix M ∈ R m × n and two weight matrices P, W ∈ R m × n , we have that
[Equation]
Proof.
We can rewrite (cid:80) i,j P ij M ij − (cid:80) i,j W ij M ij as
[Equation]
where ⟨· , ·⟩ denotes the trace inner product between matrices.
Applying Hölder’s inequality, we obtain
[Equation]
Substituing M ij =
A ij − B ij in Lemma 4, we immediately obtain the desired results in Lemma 1.
A.5 Proof of Lemma
Recall that
[Equation]
For each ( s, a ) ∈ S × A , we have
[Equation]
where the last step follows from equations (12) and (11).
Multiplying both sides by d π θ t ( s, a ) and summing over ( s, a ) , we obtain
[Equation]
thereby proving the first equation in the lemma.
Continuing the above recursion yields the second equation.
Fix t ∈ [ H ] .
For g t ∈ ∆( S × A ) satisfying supp( g t ) ⊆ supp( ρ t ) , we have
[Equation]
≤
[Equation]
s
[Equation]
by Lemma 1 and the fact that (cid:13)(cid:13)(cid:13) (cid:98) Q π
[Equation]
t (cid:13)(cid:13)(cid:13) ∗ ≤ ∥
Y t ∥ ∗ by construction.
Applying triangle inequality, we get (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:88) s,a g [Equation] s, [Equation]
Q
[Equation]
s, a
[Equation]
Y
[Equation]
≤
[Equation]
s
[Equation]
g
[Equation]
Q
[Equation]
) −
[Equation]
) (cid:17)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
where the last step follows from the constraint in (4).
Recall that
[Equation]
Using the above expressions [Equation] we obtain Z
[Equation]
a ) . Note that (cid:98) r t and (cid:98) F t are both the average of n t ( s, a ) bounded random variables.
Additionally, | (cid:98) r t | is bounded by 1 and (cid:12)(cid:12)(cid:12) (cid:98) F t (cid:12)(cid:12)(cid:12) is bounded by H .
A standard application of Hoeffding’s inequality followed by union bound yields
[Equation]
Setting the probability to be lower bounded by 1 − δ gives
[Equation]
Similar techniques yield
[Equation]
Applying the high probability bound and Lemma 2, we obtain
H (cid:88) t =
[Equation]
Minimizing the RHS over all g t yields the desired result.
Remark 2 .
Note that the proof does not rely on the specific distribution of ρ t .
Instead, we only use the support of ρ t and upper bound the statistical error by the maximum deviation on the support.
It will be an interesting future direction to incorporate the exact distribution of ρ t into the error bound. A.7 Proof of Theorem
Invoking Theorem 2, we have
[Equation]
n
[Equation]
with probability at least 1 − δ .
Since we assume Dis( ρ t ∥ d πt ) ≤ B for all t ∈
[ H ] and π ∈ Π B , the discrepancies can be upper bounded by B .
Combining with the fact that ∥ Y πt ∥ ∗ ≤ √ d ∥
Y πt ∥ F ≤ √
SAHd , we get the desired result.