Scaling Marginalized Importance Sampling to High-Dimensional State-Spaces via State Abstraction
Abstract
We consider the problem of off-policy evaluation (OPE) in reinforcement learning (RL), where the goal is to estimate the performance of an evaluation policy, π e , using a fixed dataset, D , collected by one or more policies that may be different from π e .
Current OPE algorithms may produce poor OPE estimates under policy distribution shift i.e., when the probability of a particular state-action pair occurring under π e is very different from the probability of that same pair occurring in D
Voloshin et al.
Fu et al.
In this work, we propose to improve the accuracy of OPE estimation by projecting the ground state-space into a lowerdimensional state-space using concepts from the state abstraction literature in RL.
Specifically, we consider marginalized importance sampling (MIS) OPE algorithms which compute distribution correction ratios to produce their OPE estimate.
In the original state-space, these ratios may have high variance which may lead to high variance OPE.
However, we prove that in the lower-dimensional abstract state-space the ratios can have lower variance resulting in lower variance OPE.
We then highlight the challenges that arise when estimating the abstract ratios from data, identify sufficient conditions to overcome these issues, and present a minimax optimization problem whose solution yields these abstract ratios.
Finally, our empirical evaluation on difficult, high-dimensional state-space OPE tasks shows that the abstract ratios can make MIS OPE estimators achieve lower mean-squared error and more robust to hyperparameter tuning than the ground ratios.
Introduction
One of the key challenges when applying reinforcement learning (RL) Sutton and Barto (2018) to real-world tasks is the problem of off-policy evaluation (OPE)
Fu et al.
; Voloshin et al.
The goal of OPE is to evaluate a policy of interest by leveraging offline data generated by possibly different policies.
Solving the OPE problem would enable us to estimate the performance of a potentially risky policy without having to actually deploy it.
This capability is especially important for sensitive real-world tasks such as healthcare and autonomous driving.
The core OPE problem is to produce accurate policy value estimates under policy distribution shift.
This problem is particularly difficult on tasks with high-dimensional state-spaces Voloshin et al.
Fu et al.
For example, consider the AntUMaze problem illustrated on the left side of Figure 1.
In this task, an ant-like robot with a high-dimensional state representation moves in a U-shaped maze and receives a reward only for reaching a specific 2 D coordinate goal location.
The state-space of this task includes information such as 2 D location, ant limb angles, torso orientation etc., resulting in a 29 -dimensional state-space.
The OPE task is to evaluate the performance of a particular policy’s ability to take the ant to the 2 D goal location using data that may be collected by different policies.
Policy distribution shift is common in this type of high-dimensional task since the chances of different policies inducing similar limb angles, torso orientations, paths traversed etc. are incredibly slim.
Notice, however, that while different policies may induce different body configurations, they may traverse similar 2 D paths since all (successful) policies must move the ant through roughly the same path to reach the goal.
Moreover, the only critical information needed from the state-space to determine the ant’s per-step reward are its 2 D coordinates.
Motivated by this observation, we propose to improve the accuracy of
OPE algorithms on high-dimensional state-space tasks by projecting the high-dimensional state-space into a lower-dimensional space.
This idea is illustrated on the right side of Figure 1 where the ant is reduced to a 2 D point-mass.
With this general motivation in mind, in this paper, we leverage concepts from the state abstraction literature Li et al.
(2006) to improve the accuracy of marginalized importance sampling (MIS) OPE algorithms which estimate state-action density correction ratios to compute a policy value estimate
Liu et al.
(2018a);
Xie et al.
Due to the low chances of similarity between states of policies in high-dimensional state-spaces, current MIS algorithms can produce high variance state-action density ratios, resulting in high variance OPE estimates.
However, if we are given a suitable state abstraction function, we can project the high-dimensional ground state-space into a lower-dimensional abstract state-space.
The projection step increases the chances of similarity between these lower-dimensional states, resulting in low variance density ratios and OPE estimates.
To the best of our knowledge, this work is the first to leverage concepts from state abstraction to improve OPE.
We make the following contributions:
Theoretical analyses showing: (a) the variance of abstract state-action ratios is at most that of ground state-action ratios; and (b) that our abstract MIS OPE estimator is unbiased, strongly consistent, and can have lower variance than its ground equivalent.
An algorithm, based on a popular class of MIS algorithms, to estimate the abstract stateaction ratios.
An empirical analysis of our estimator on a variety of high-dimensional state-space tasks.
Preliminaries
In this section, we discuss relevant background information.
2.1 Notation and Problem Setup
We consider an infinite-horizon Markov decision process (MDP), M = ⟨S , A , R , P, γ, d 0 ⟩ , where
S is the state-space, A is the action-space, R : S ×
A → ∆([0 , ∞ )) is the reward function, P :
S × A → ∆( S ) is the transition dynamics function, γ ∈ [0 , 1) is the discount factor, and d 0 ∈ ∆( S ) is the initial state distribution.
The agent acting, according to policy π , in the MDP generates a trajectory: s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , ... , where s 0 ∼ d 0 , a t ∼ π ( ·| s t ) , r t ∼ R ( s t , a t ) , and s t
+1 ∼ P ( ·| s t , a t ) for t ≥ 0 .
We define r ( s, a ) :
=
E r ∼R ( s,a ) [ r ] and the agent’s discounted state- action occupancy measure under policy π as d π :
[Equation]
where d π ( s t , a t ) is the probability the agent will be in state s and take action a at time-step t under policy π .
In the infinite-horizon setting, we define the performance of policy π to be its average reward, ρ ( π ) :
=
E ( s,a ) ∼ d π ,r ∼R ( s,a ) [ r ] .
Note that E ( s,a ) ∼ d π ,r ∼R ( s,a ) [ r ] = (1 − γ ) E s 0 ∼ d 0 ,a 0 ∼ π [ q π
( s
[Equation]
, a 0 )] where q π ( s, a ) :
=
E [ (cid:80) ∞ k =0 γ t r k + t | s t = s, a t = a ] is the action-value function which satisfies q π ( s, a )
=
r ( s, a ) + γ
[Equation]
s ′ ∼ P ( s,a ) ,a ′ ∼ π
( s ′ )
[ q π ( s ′ , a ′ )] .
2.2 Off-Policy Evaluation (OPE)
In behavior-agnostic off-policy evaluation (OPE), the goal is to estimate the performance of an evaluation policy π e given only a fixed offline data set of transition tuples, D :=
{ ( s i , a i , s ′ i , r i ) } mTi =1 , where ( s i , a i ) ∼ d D , m is the batch size (number of trajectories), and T is the fixed length of each trajectory, generated by unknown and possibly multiple behavior policies.
The difficulty in OPE is to estimate ρ ( π e ) under d π e given samples only from d D .
We define the average-reward in dataset D to be ¯ r D := E ( s,a ) ∼ d D ,r ∼R ( s,a ) [ r ] .
As in prior OPE work, we assume that if d π e ( s, a ) > 0 then d D ( s, a ) > 0 .
Empirically, we measure the accuracy of an estimate ˆ ρ ( π e ) by generating M datasets and then computing the relative mean-squared error (MSE):
MSE (ˆ ρ ( π e )) :
= 1 M (cid:80) Mi
=1 ( ρ ( π e ) −
ˆ ρ
i ( π e )) 2 ( ρ ( π e ) − ¯ r D i ) 2 ,
where ˆ ρ i ( π e ) is computed using dataset D i and
¯ r D
i is the average reward in D i .
2.2.1 Marginalized Importance Sampling (MIS)
In this work, we focus on MIS methods, which evaluate π e by using the ratio between d π e and d D .
That is, MIS methods evaluate π e by estimating ρ ( π e ) :
=
E ( s,a ) ∼ d D ,r ∼R ( s,a ) [ ζ ( s, a ) r ] , where ζ ( s, a ) :
= d πe ( s,a ) d D ( s,a ) is the state-action density ratio for state-action pair ( s, a ) and d π ( s, a ) =
d π ( s ) π ( a | s ) .
When the true ζ is known, the empirical estimate of ρ ( π e ) is:
[Equation]
where N is the number of samples.
In practice, however, ζ is unknown and must be estimated.
One set of ζ -estimation algorithms, which have also shown potential for good OPE performance
Voloshin et al.
(2021), is the DICE family Yang et al.
While there are many variations, the general DICE optimization problem is:
[Equation]
where the solution to the optimization, ζ ∗ ( s, a ) , are the true ratios.
The estimator we present in
Section 3 builds upon the DICE framework.
2.3 State Abstractions
We define a state abstraction function as a mapping ϕ : S → S ϕ , where S is called the ground state-space and
S ϕ is called the abstract state-space.
We consider state abstraction functions that partition the ground state-space into disjoint sets.
We can use ϕ to project the original MDP into a new abstract MDP with the same action-space A and reward and transition dynamics functions defined as:
[Equation]
where w : S → [0 , 1] is a ground state weighting function where for each s ϕ , (cid:80) s ∈ ϕ
[Equation]
1 ( s ϕ
[Equation]
w ( s
[Equation]
=
Li et al.
Similarly a policy can be transformed into its abstract equivalent as:
[Equation]
We define the following state-weighting function for an arbitrary policy π :
[Equation]
and only consider abstractions that satisfy:
Assumption 1 (Reward distribution equality) .
∀ s 1 , s 2 ∈ S such that s 1 , s 2 ∈ s ϕ
[Equation]
=
R ( s 2 , a ) .
Assumption 1 implies that, regardless of the choice of the state-weighting function, for given action a , ∀ s ∈ s ϕ , R ϕ ( s ϕ , a )
=
R ( s, a ) i.e. the reward distribution of an abstract state equals that of the ground states within that abstract state.
Abstract Marginalized Importance Sampling
Marginalized IS methods may suffer from high variance in high-dimensional state-spaces.
To potentially reduce this high variance, we propose to first use ϕ to project D into the abstract statespace to obtain: D ϕ :
= { ( s ϕ , a, r ϕ , s ′ ϕ ) } where s ϕ =
ϕ ( s ) and r ϕ ( s, a ) =
r ( s, a ) ∀ s ∈ s ϕ , and then use the following estimator on D ϕ to estimate π ϕe :
Definition 1 (Abstract estimator) .
We define our estimator of ρ ( π ϕe ) as follows:
[Equation]
where N is the number of samples, d π ϕ ( s ϕ , a )
=
d π
[Equation]
s ϕ )
π ϕ
[Equation]
a | s ϕ ) with d π ϕ ( s ϕ )
=
(cid:80) s
∈ ϕ − 1 ( s ϕ )
d π ( s ) and π ϕ constructed using w π .
In the remainder of this section, we first give an example to build intuition for why the abstract ratios d πϕe ( s ϕ ,a ) d D ϕ ( s ϕ ,a ) can be lower variance than the ground ratios and then show theoretically that the OPE estimator given in Equation (3) is strongly consistent and can produce lower variance OPE estimates of ρ ( π e ) than the ground equivalent.
Finally, we identify sufficient conditions under which accurate estimation is possible and adapt an existing DICE algorithm to estimate the abstract ratios.
3.1 A Hard Example for Ground MIS Ratios
We present a hard example for ground MIS ratios in Figure 2 that provides intuition for why the abstract MIS ratios can have lower variance ratios than the ground ratios.
Consider two symmetric policies, π e and π D , executed in the ground MDP (left side).
In this example, the high variance of the true state-action density ratios d πe ( s 1 ,a 0 )
d π D
( s 1 ,a 0 )
≈
and d πe ( s 2 ,a 0 ) d π D ( s 2 ,a 0 )
≈ 100 can lead to high variance estimates of ρ ( π e ) .
Notice, however, that states s 1 and s 2 are essentially equivalent i.e. r ( s 1 , a ) =
r ( s 2 , a ) ∀ a ∈ A and can be aggregated together into a single state, s ϕ 1 (Assumption 1).
We find that the state-action density ratio in this abstract MDP (right side) d πϕe ( s ϕ 1 ,a 0 ) d
πϕ D
[Equation]
s ϕ 1 ,
[Equation] [Equation]
is of low variance, which can lead to low variance estimates of ρ ( π e ) .
In fact, with Theorem 1 we prove (proof in Appendix A.3) that the abstract ratios are guaranteed to have variance at most that of the ground ratios:
Theorem 1.
Var
[Equation]
[Figure]
where equality holds only if ϕ is the identity function i.e. ϕ ( s )
=
s, ∀ s ∈ S
and/or if ∀ s 1 , s 2 ∈ S such that ∀ s 1 , s 2 ∈ s ϕ and for a given action a , d πe ( s 1 ,a ) d D ( s 1 ,a )
=
d πe ( s 2 ,a ) d D ( s 2 ,a ) , ∀ s ϕ ∈ S ϕ , a ∈ A .
Through this example, we have shown how projecting S → S ϕ can lower the variance of density ratios.
3.2 MIS OPE with True Abstract Ratios
We now present the statistical properties of our estimator assuming it has access to the true abstract state-action ratios.
Due to space constraints, we defer proofs to Appendix A.3.
We prove our abstract estimator is unbiased (Theorem 4 in Appendix A.3) and strongly consistent (Theorem 2 and Corollary 1):
Theorem 2.
Our estimator, ˆ ρ ( π ϕe ) , given in Equation 3 is an asymptotically consistent estimator of ρ ( π e ) in terms of MSE: lim N →∞ E
[(ˆ ρ ( π ϕe ) − ρ ( π e )) 2 ] = 0 .
We also compare the variances between our abstract estimator and the ground equivalent.
If we assume that D is i.i.d as done in previous work Sutton et al.
; Uehara et al.
; Nachum et al.
; Zhang et al.
(2020a)
, then Var (ˆ ρ ( π ϕe ))
≤ Var (ˆ ρ ( π e )) , where equality holds only under the same conditions as described above for Theorem 1.
In general, however, variance reduction depends on the covariances between the weighted per-step rewards Liu et al.
Theorem 3.
If
Assumption and if for any fixed
[Equation] [Equation]
[Equation] [Equation]
k , a k ) (cid:17) hold then Var (ˆ ρ ( π ϕe ))
≤ Var (ˆ ρ ( π e )) .
3.3 Estimating the Abstract Ratios
Thus far, we have assumed access to the true abstract ratios.
However, in practice, these ratios are unknown and must be estimated from D ϕ .
In this section, we highlight the challenges in estimating the abstract ratios and identify sufficient conditions on ϕ that allow for accurate estimation.
Once ϕ satisfies these conditions, any off-the-shelf state-action density estimation method can be used to estimate the abstract ratios.
In this work, we focus on showing that DICE estimates the true abstract ratios.
We note that the following new conditions on ϕ are only needed for unbiased estimation of the abstract ratios; accurate OPE on the abstracted MDP can still be realized under only Assumption 1.
We first observe that evaluating π e using D is equivalent to evaluating π ϕe using D ϕ if D ϕ is generated from an abstract MDP with transition dynamics constructed according to w π e .
This equivalence is given by the following proposition:
Proposition 1.
If Assumption 1 holds, the average reward of ground policy π executed in ground
MDP M , ρ ( π ) , is equal to the average reward of abstract policy π ϕ executed in abstract MDP M ϕ constructed with w π , ρ ( π ϕ ) .
That is, ρ ( π ) =
ρ ( π ϕ ) .
Proposition 1 suggests that we can estimate ρ ( π e ) by applying any OPE algorithm to evaluate π ϕe using D ϕ .
Unfortunately, there are two challenges which prevent us from doing so in the general case.
Fortunately, there are special cases where unbiased estimation of the abstract ratios is still possible using only existing MIS algorithms.
Challenge 1: Transition dynamics distribution shift .
In practice, Proposition 1 cannot be immediately applied since w π e is unknown.
Additionally, our off-policy data D ϕ := { ( s ϕ , a, r ϕ , s ′ ϕ ) } is distributed in the following way:
( s ϕ , a ) ∼ d D ϕ , r ϕ ∼ R ϕ , s ′ ϕ ∼ P ϕw D ( s ϕ , a ) where P ϕw D are the transition dynamics of the abstract MDP constructed with w D as the state-weighting function.
Thus, in addition to the original policy distribution shift problem, we also encounter a transition dynamics distribution shift problem due to the projection where we want to evaluate π ϕe in an MDP with P ϕw πe , but we only have samples of data generated in an MDP with P ϕw D .
Since the transition dynamics are unknown, we cannot correct its distribution shift as we would correct for policy distribution shift.
However, one condition on ϕ that will avoid the transition distribution shift is:
Assumption 2 (Transition dynamics similarity) .
∀ s 1 , s 2 ∈ S such that s 1 , s 2 ∈ s ϕ , ∀ a ∈ A , x ∈ S ϕ , we have (cid:80)
[Equation]
Together with Assumption 1, ϕ is now the so-called bisimulation abstraction Ferns et al.
Castro (2020).
This property of ϕ eliminates the transition dynamics distribution shift since now
P ϕw
[Equation]
(applying Assumption 2 and definition of P ϕ from Section 2.3).
Thus, we have D ϕ distributed with
P ϕw πe =
P ϕ , which then allows us to apply any MIS algorithm to compute the abstract state-action ratios using D ϕ .
Challenge 2: Inaccessible π ϕe .
To the best of our knowledge, all existing MIS algorithms require access to the evaluation policy to estimate the density ratios.
However, in the abstract MDP, π ϕe is inaccessible since w π e is unknown.
To overcome this issue, we identify the following condition on ϕ :
Assumption 3 ( π e action-distribution equality) .
∀ s 1 , s 2 ∈ S such that s 1 , s 2 ∈ s ϕ , π e ( s 1 ) =
π e ( s 2 )
Assumption 3 then gives us π ϕe ( s ϕ ) =
π e ( s ) , ∀ s ∈ s ϕ (applying Assumption 3 and definition of π ϕ from Section 2.3), which allows us to simulate sampling from π ϕe ( ·| ϕ ( s )) by just sampling from π e ( ·| s ) .
Given a ϕ with these properties, we can compute the abstract ratios needed to estimate ρ ( π e ) by applying a suitable MIS algorithm to D ϕ .
We use BestDICE Yang et al.
(2020), and call our algorithm
AbstractBestDICE, which solves the following optimization problem:
[Equation]
(4) where, ν : S ϕ ×
A → R , λ ∈ R , and ζ : S ϕ ×
A → R ≥ 0 .
The solution to the optimization, ζ ∗ ( s ϕ , a ) =
d π ϕe ( s ϕ , a ) /d
D ϕ
( s ϕ , a ) is the true abstract ratios.
Since the derivation of AbstractBestDICE follows the same steps as BestDICE, we defer it to Appendix A.4.
We note that it may be difficult to validate Assumptions 2 and 3 in practice, which may result in loss of the consistency guarantee of Theorem 2.
Nevertheless, in Section 4 we show that AbstractBestDICE leads to accurate OPE in high-dimensional state-spaces even when assumptions may not hold.
Empirical Study
We will now show how projecting S → S ϕ can produce more accurate OPE estimates in practice.
We answer the following questions:
Do the true abstract ratios produce lower variance OPE estimates than the true ground ratios?
Does AbstractBestDICE: (a) compute the true ratios when Assumptions 1, 2, and 3 are satisfied and (b) produce data-efficient and stable estimates of ρ ( π e ) even when Assumptions 2 and 3 fail to hold?
4.1 Empirical Setup
In this section, we describe the algorithms and domains of our empirical study.
Due to space constraints, we defer details regarding policies, hyperparamer tuning, etc. to the appendix (A.5 and
A.6).
4.1.1 Algorithms
We compare AbstractBestDICE to ground BestDICE Yang et al.
As also reported by Yang et al.
Fu et al.
, we found in preliminary experiments that BestDICE performed much better than other variants of DICE such as DualDICE Nachum et al.
(2019), GradientDICE Zhang et al.
(2020b), etc.
4.1.2 Domains
We focus on high-dimensional state-space tasks, which have been known to be particularly challenging for DICE methods Fu et al.
For each environment below we specify a fixed ϕ .
For Reacher,
Walker2D, Pusher, and AntUMaze, ϕ satisfies only Assumption 1.
• TwoPath MDP .
This MDP is pictured in Figure 2.
In this domain, Assumptions 1, 2, andare satisfied.
We also run the same experiments for when these assumptions are violated (see Appendix A.5).
All trajectories are 100 time-steps long.
• Reacher Brockman et al.
A robotic arm tries to move to a goal location.
Here, s ∈ R 11 and a ∈ R 2 .
Since the reward function is the Euclidean distance between the arm and goal, ϕ extracts only the arm-to-goal vector, and angular velocities from the ground state, resulting in s ϕ ∈ R 4 .
All trajectories are 200 time-steps long.
• Walker2D Brockman et al.
A bi-pedal robot tries to move as fast as possible.
Here, s ∈ R 18 and a ∈ R 6 .
We use the Euclidean distance from the start location as the reward function and use a ϕ that extracts x and z coordinates and top angle of the walker’s body, resulting in s ϕ ∈ R 3 .
All trajectories are 500 time-steps long.
• Pusher Brockman et al.
A robotic arm tries to push an object to a goal location.
Here, s ∈ R 23 and a ∈ R 7 .
Since the reward function is the Euclidean distance between object and arm and object and goal, ϕ extracts only object-to-arm and object-to-goal vectors, resulting in s ϕ ∈ R 6 .
All trajectories are 300 time-steps long.
• AntUMaze Fu et al.
This sparse-reward task requires an ant to move from one end of the U-shaped maze to the other end.
Here, s ∈ R 29 and a ∈ R 8 .
We use the “play" version where the goal location is fixed.
Since the reward function is +1 only if the 2 D location of the ant is at a certain Euclidean distance from the 2 D goal location, ϕ extracts only the 2 D coordinates of the ant, resulting in s ϕ ∈ R 2 .
All trajectories are 500 time-steps long.
4.2 Empirical Results
In this section, we describe our main empirical results; additional experiments can be found in appendix A.6.
4.2.1 True Ratios for OPE
We conduct an experiment on the TwoPath MDP to estimate ρ ( π e ) where we apply the ground estimator given in Equation (1) and our abstract estimator given in Equation (3), assuming both have access to their respective true ratios .
The results of this experiment are illustrated in Figure 3(a).
We can observe that the abstract estimator with the true abstract ratios produces substantially more data-efficient and lower variance OPE estimates for different batch sizes compared to the ground equivalent.
[Figure]
4.2.2 True Ratio Estimation
To verify if AbstractBestDICE accurately estimates the true ratios we conduct the following experiment on the TwoPath MDP.
We give AbstractBestDICE data of batch size 300 to estimate the abstract ratios ˆ ζ ϕ
and then use ˆ ζ ϕ to estimate the abstract state-action densities of π ϕe , ˆ d π ϕe ( s ϕ , a ) =
ˆ ζ ϕ ( s ϕ , a ) d π ϕ D ( s ϕ , a ) , where we have access to the true d π ϕ D ( s ϕ , a ) .
We then compare ˆ d π ϕe to the true d π ϕe , which we compute using a batch size of 300 trajectories collected from π ϕe roll-outs, using a correlation plot shown in Figure 3(b).
From the figure we can see AbstractBestDICE accurately estimates the abstract state-action density ratios.
When assumptions are violated, however, ratio estimation accuracy can reduce (see Appendix A.5)
[Figure]
4.2.3 Data-Efficiency
Figure 4 shows the results of our (relative) MSE vs. batch size experiment for the function approximation case.
For a given batch size, we train each algorithm for 100 k epochs with different hyperparameters sets, record the (relative) MSE on the last epoch by each hyperparameter set, and plot the lowest MSE achieved by these hyperparameter sets.
We find that AbstractBestDICE is able to achieve lower MSE than BestDICE for a given batch size.
We note that while hyperparameter tuning is difficult in OPE, in this experiment, we aim to evaluate the performance of each algorithm assuming both had favorable hyperparameters.
4.2.4 Hyperparameter Robustness
Finally, we study the robustness of these algorithms to hyperparameters tuning.
In practical OPE, hyperparameter tuning with respect to MSE is impractical since the true ρ ( π e ) is unknown Fu et al.
; Paine et al.
Thus, we want OPE algorithms to be as robust as possible to hyperparameter tuning.
The main hyperparameters for DICE are the learning rates of ζ and ν , α ζ and α ν .
For these experiments, we focus on very small batch sizes, where we would expect high sensitivity.
The results of this study are in Figure 3(c).
We find that our algorithm has a less volatile
MSE than BestDICE (also see appendix A.6 for more similar results).
In a related experiment, we also find AbstractBestDICE can be more stable than BestDICE during training (see appendix A.6).
Related Work
MIS and Off-Policy Evaluation .
There have been broadly three families of MIS algorithms in the OPE literature to estimate state-action density ratios.
One is the DICE family, which includes: minimax-weight learning Uehara et al.
, DualDice Nachum et al.
(2019), GenDICE Zhang et al.
(2020a), GradientDICE Zhang et al.
(2020b), and BestDICE Yang et al.
In our work, we adapt BestDICE to estimate the abstract ratios.
The second family of MIS algorithms is the COP-TD algorithm Hallak and Mannor (2017);
Gelada and Bellemare (2019), which learns the state density ratios with an online TD-styled update.
The third family is the variational power method Wen et al.
(2020) algorithm which generalizes the power iteration method to estimate density ratios.
While our focus has been on MIS algorithms, there are many other OPE algorithms such as model-based methods Zhang et al.
(2021b)
; Hanna et al.
; Liu et al.
(2018b)
, fitted-Q evaluation Le et al.
(2019), doubly-robust methods Jiang and Li (2015); Thomas and Brunskill (2016), and IS Precup et al.
; Thomas (2015); Hanna et al.
; Thomas et al.
State Abstraction and Representation Learning .
The literature on state abstraction is extensive
Singh et al.
; Dietterich (1999); Ferns et al. (2011); Li et al. (2006); Abel (2020).
However, much of this work has been exclusively focused on building a theory of abstraction and on learning optimal policies.
A related topic to state abstraction is representation learning.
Recently, there has been much work showing the importance of good representations for offline RL Wang et al.
Yin et al.
; Geng et al.
; Zhan et al.
; Chen and Jiang (2022).
To the best of our knowledge, no work has leveraged state abstraction techniques to improve the accuracy of OPE algorithms.
Summary and Future Work
In this work, we showed that we can improve the accuracy of OPE estimates by projecting the original ground state-space into a lower-dimensional abstract state-space using state abstraction and performing OPE in the resulting abstract Markov decision process.
Our theoretical results proved that: 1) abstract state-action ratios have variance at most that of the ground ratios; and 2) the abstraction MIS OPE estimator is unbiased, strongly consistent, and can have lower variance than the ground equivalent.
We then highlighted the challenges that arise when estimating the abstract ratios from data, identified sufficient conditions to overcome these issues, and adapted BestDICE into
AbstractBestDICE to estimate the abstract ratios.
In our empirical results, we obtained more accurate estimates with added hyperparameter robustness on difficult, high-dimensional state-space tasks.
There are several directions for future work.
First, Assumptions 2 and 3 are strict.
Further investigation is needed to see if these assumptions can be relaxed.
Second, we assumed the abstraction function was given.
It would be interesting to leverage existing ideas Gelada et al.
; Zhang et al.
(2021a) to learn ϕ .
Finally, we want to emphasize that this work instantiates the general abstraction + OPE direction.
While this work focused exclusively on MIS algorithms, a promising direction will be to apply abstraction techniques to model-based, trajectory IS, and value-function based OPE.
Acknowledgements
Support for this research was provided by the Office of the Vice Chancellor for Research and
Graduate Education at the University of Wisconsin — Madison with funding from the Wisconsin
Alumni Research Foundation.
The authors thank the anonymous reviewers, Nicholas Corrado, Ishan Durugkar, and Subhojyoti Mukherjee for their helpful comments in improving this work.
A Appendix A.1 Preliminaries
This section provides the supporting lemmas and definitions that we leverage to prove our lemmas and theorems.
Definition 2 (Almost Sure Convergence) .
A sequence of random variables, ( X n ) ∞ n =1 , almost surely converges to the random variable, X if
[Equation]
We write X n a.s.
→ X to denote that the sequence ( X n ) ∞ n =1 converges almost surely to X .
Definition 3 ((Strongly) Consistent Estimator) .
Let θ be a real number and (ˆ θ n ) ∞ n =1 be an infinite sequence of random variables.
We call ˆ θ n a (strongly) consistent estimator of θ if and only if ˆ θ n a.s.
→ θ .
Lemma 1.
If ( X i ) ∞ i =1 is a sequence of uniformly bounded real-valued random variables, then
X n a.s.
→ X if and only if lim N →∞ E
[( X n − X ) 2 ] = 0 .
Proof
See Lemma 3 in Thomas and Brunskill (2016).
A.2 Assumptions and Definitions
In the main paper, we provided the major assumptions required for our theoretical and empirical work relevant to abstraction and OPE.
Here we provide supporting assumptions typically used in the OPE literature used for the theoretical analysis.
Assumption 4 (Coverage) .
For all ( s, a ) ∈ S × A , if π e ( a | s ) > 0 then π b ( a | s ) > 0 .
Assumption 5 (Non-negative reward) .
We assume that the reward function is bounded between [0 , ∞ ) .
Definition 4 (Ground state normalized weightings) .
For a given policy π , each ground state s ∈ s ϕ , has a state aggregation weight, w π ( s ) =
[Equation]
s
′∈ ϕ − 1( sϕ ) d π ( s ′ ) , where d π ( s ) is the discounted stateoccupancy measure of π .
A.3 Proofs
In the proofs below, we denote the collection of behavior policies that generated D with π D .
That is, π D is the conditional probability of an action occurring in a given state in the data.
Similarly, we also have π ϕ D .
These minor changes give us d D = d π D and d D ϕ =
d π ϕ D .
Lemma 2.
For an arbitrary function, f , E s ϕ ∼ d πϕ ,a ∼ π ϕ (cid:2) f
( s ϕ , a ) (cid:3) =
E s ∼ d π ,a ∼ π [ f ( ϕ ( s ) , a )] .
Proof
[Equation]
where (a) is due to Definition 4 and we can replace s ϕ with ϕ ( s ) when we know s ∈ s ϕ .
Theorem 1.
Var
[Equation]
Proof
Before comparing the variances
[Equation]
we note that due to Assumption 4 and Lemma 2:
[Equation]
Denote,
[Equation]
Var
[Equation]
and V
[Equation]
Var
[Equation]
Now consider the difference between the two variances
[Equation]

We can analyze this difference by looking at one abstract state and one action and all the states that belong to it.
That is, for a fixed abstract state, s ϕ
[Equation]
and fixed action
[Equation]
a , we have:
[Equation]
where (a) is due to Definition 4.
If we can show that D ′ ≥ 0 for all possible sizes of | s ϕ | , we will the have the original difference, D , is a sum of only non-negative terms, thus proving Theorem 1.
We will prove D ′ ≥ 0 by inductive proof on the size of | s ϕ | from 1 to some n ≤ |S| .
Let our statement to prove, P ( n ) be that D ′ ≥ 0 where n = | s ϕ | .
This is trivially true for P (1) where the ground state equals the abstract state.
Now consider the inductive hypothesis, P ( n ) is true for n ≥ 1 .
Now with the inductive step, we must show that P ( n + 1) is true given P ( n ) is true.
Starting with the inductive hypothesis:
[Equation]
We define
[Equation]
After making the substitutions
[Equation]
we have:
[Equation]
We have the above result holding true for when the | s ϕ | = n .
Now consider the inductive step in relation to the inductive hypothesis where a new state, s n +1 is added to the abstract state.
We have the following difference:
[Equation]
For ease in notation
[Equation]
let x = d π e
( s n +1 )
π e ( a | s n +1 ) and y =
d π D ( s n +1 )
π D ( a | s n +1 ) .
The above difference is then:
[Equation]
The above difference, D ′′ , is minimized most when C is as large as possible.
From the inductive hypothesis, we have C ≤ √ SC ′ .
The minimum difference can be written as:
[Equation]
So we have D ′′ ≥ 0 for | s ϕ | = n + 1 , which means D ′ ≥ 0 .
We have showed that P ( n ) is true for all n .
We now have the original difference, D , to be a sum of non-negative terms after performing this same grouping for all abstract states and actions, which results in:
[Equation]
Thus, we have:
[Equation]
Proposition 1.
If Assumption 1 holds, the average reward of ground policy π executed in ground
MDP M , ρ ( π ) , is equal to the average reward of abstract policy π ϕ executed in abstract MDP M ϕ constructed with w π , ρ ( π ϕ ) .
That is, ρ ( π ) =
ρ ( π ϕ ) .
Proof
Consider the definition of R ϕπ :
[Equation]
where (a) is due to Definition 4, (b) is due to Definition 4 and Assumption
Theorem 4.
If Assumption 1 holds, our estimator, ˆ ρ ( π ϕe ) as defined in Equation 3, is an unbiased estimator of ρ ( π e ) .
Proof
We first consider the expectation of a single sample,
[Equation]
[Equation]
where (c) is due to Definition 4 and Assumption 1, (d) is due to Assumption 4, and (e) is due to
Proposition 1.
We have the bias defined as:
[Equation]
Bias [
[Equation]
where (a) is due to linearity of expectation and (b) is due to expectation of a single sample.
Theorem 2.
Our estimator, ˆ ρ ( π ϕe ) , given in Equation 3 is an asymptotically consistent estimator of ρ ( π e ) in terms of MSE: lim N →∞ E
[(ˆ ρ ( π ϕe
[Equation]
Proof
We have the MSE of ˆ ρ ( π ϕe ) w.r.t ρ ( π e ) defined in terms of the bias and variance as follows:
[Equation]
 where (a) is because ˆ ρ is an unbiased estimator as shown in Theorem 4.
Due to Assumptions 4 and 5, (cid:32)(cid:80)
[Equation]
is a bounded value.
Thus, as mT → ∞ , Var [ˆ ρ ( π ϕe )] → 0 .
We then have lim mT →∞ E
[(ˆ ρ ( π ϕe ) − ρ ( π e )) 2 ] = 0 .
Thus, the estimator ˆ ρ ( π ϕe ) is consistent in MSE.
Corollary 1.
If Assumption 1 holds, then our estimator, ˆ ρ ( π ϕe ) as defined in Equation 3 is an asymptotically strongly consistent estimator of ρ ( π e ) .
Proof
Theorem 2 showed that ˆ ρ ( π ϕe ) is consistent in terms of MSE.
Then by applying Lemma 1, we have ˆ ρ ( π ϕe ) to be an asymptotically strongly consistent estimator of ρ ( π e ) .
That is, ˆ ρ ( π ϕe ) a.s.
→ ρ ( π e ) .
Note that this proof is very similar to the one given in Theorem 1, where the main the difference is that we are analyzing the variance of an OPE estimator on given batch of data.
Theorem 3.
If
Assumption and if for any fixed ≤ t < k ≤
T ,
[Equation]
Cov (cid:16)
[Equation]
) d
[Equation] [Equation]
a t
[Equation]
( s t
[Equation]
a t )
[Equation]
d πe ( s k
[Equation]
a k ) d D ( s k
[Equation]
a k ) r ( s k
[Equation]
a k ) (cid:17) hold then Var (ˆ ρ ( π ϕe ))
≤ Var (ˆ ρ ( π e )) .
Proof
We first consider the general form of the variance of the baseline estimator, ˆ ρ ( π e ) (and the similar form applies to ˆ ρ ( π ϕe ) ).
For simplicity in notation, we take the batch size, m =
[Equation]
but the analysis holds for general m .
[Equation]
 
Here we have V gt to be the variance of a single sample at time t .
Similarly, for our estimator, we have
V ϕt =
Var
[Equation]
We will show that V ϕt ≤ V gt .
We drop the subscript t for convenience since the analysis applies for all t .
Before comparing the variances, we note that due to Assumption 1 and 4, Lemma 2, and Proposition 1:
[Equation]
Consider the difference between the two variances
[Equation]
where (a) is due to Lemma 2.
We can analyze this difference by looking at one abstract state and one action and all the states that belong to it.
That is, for a fixed abstract state, s ϕ , and fixed action
[Equation]
a , we have:
[Equation]
where (a) is due to Definition 4 and Assumption 1.
If we can show that D ′ ≥ 0 for all possible sizes of | s ϕ | , we will the have the original difference,
D , a sum of only non-negative terms, thus proving Theorem 3.
We will prove D ′ ≥ 0 by proof by induction on the size of | s ϕ | from 1 to some n ≤ |S| .
First note that r ( s, a ) 2 > 0 , so we can ignore this term.
Let our statement to prove, P ( n ) be that D ′ ≥ 0 where n = | s ϕ | .
This is trivially true for P (1) where the ground state equals the abstract state.
Now consider the inductive hypothesis, P ( n ) is true for n ≥ 1 .
Now with the inductive step, we must show that P ( n + 1) is true given P ( n ) is true.
Starting with the inductive hypothesis:
[Equation]
After making the substitutions, we have:
[Equation]
[Equation]
We have the above result holding true for when the | s ϕ | = n .
Now consider the inductive step in relation to the inductive hypothesis where a new state, s n +1 is added to the abstract state.
We have the following difference:
[Equation]
For ease in notation, let x = d π e ( s n +1 )
π e ( a | s n +1 ) and y =
d π D ( s n +1 )
π D ( a | s n +1 ) .
The above difference is then:
[Equation]
The above difference, D ′′ , is minimized most when C is as large as possible.
From the inductive hypothesis, we have C ≤ √ SC ′ .
The minimum difference can be written as:
[Equation]
So we have D ′′ ≥ 0 for | s ϕ | = n + 1 , which means D ′ ≥ 0 .
We have showed that P ( n ) is true for all n .
We now have the original difference, D , to be a sum of non-negative terms after performing this same grouping for all abstract states and actions, which results in:
[Equation]
Thus
[Equation]
when the covariance terms interact favorably
[Equation]
that is
[Equation]
if for any fixed 1 ≤ t < k ≤ T ,
[Equation]
we have:
[Equation]
A.4 Additional AbstractBestDICE Derivation Details
We proceed assuming ϕ satisfies Assumptions 1,
[Equation]
and 3 from the main text.
We base the following derivation on that of DualDICE Nachum et al.
The only difference between DualDICE and
BestDICE (which we use in our experiments) is that the latter enforces a positivity and unit mean constraint on the ratios.
Technical observation Consider the same technical observation made in DualDICE: the solution to the scalar convex optimization problem
min x J ( x ) :
= 12 mx 2 − nx , where m ∈ R > 0 and n ∈ R ≥ 0 , is x ∗ = nm .
This observation can then be connected to a similar convex problem but in terms of functions where the ratio nm corresponds to the true abstract ratios, d π ϕe ( s ϕ , a )
/d D ϕ ( s ϕ , a ) .
Consider the following convex problem where x ∈ S ϕ ×
A → [0 , C ] ⊂ R :
[Equation]
The solution to this optimization problem is the true abstract state-action ratios: x ∗ ( s ϕ , a ) =
d π ϕe ( s ϕ , a ) /d
D ϕ
( s ϕ , a ) .
Change-of-variables As also done by Nachum et al.
(2019), we next consider a change-ofvariables trick to obtain an objective that can be approximated with samples from D ϕ .
However, as noted in Section 3, there are two challenges we must overcome before applying this procedure on D ϕ .
To overcome these challenges, we identified Assumptions 2 and 3 as conditions that ϕ must satisfy.
With a satisfactory ϕ , we can proceed as follows.
Consider an arbitrary function ν ∈ S ϕ ×
A → R that satisfies ν ( s ϕ , a ) =
x ( s ϕ , a ) + B π ϕe ν ( s ϕ , a ) where
B π ϕe ν
( s ϕ , a ) =
γ E s ′ ϕ ∼ P ϕ
( s ϕ ,a ) ,a ′ ∼ π ϕe ( s ′ ϕ )
[ ν ( s ′ ϕ , a ′ )] .
Since x ( s ϕ , a ) ∈ [0 , C ] and γ < 1 , ν is well-defined.
We can the replace x ( s ϕ , a ) with ( ν − B π ϕe ν )( s ϕ , a ) .
For the second expectation, we define β ϕt ( s ϕ ) :
= P
[Equation]
ϕ ,
[Equation]
∼ π ϕe ( s
[Equation]
) , s ϕk +1 ∼ P ϕ
( s ϕk , a k )) , 0 ≤ k < t as the abstract-state visitation probability at time-step t by the π ϕe in the abstract MDP, and then have:
[Equation]
where β ϕ
= d 0 ϕ is the initial abstract state distribution.
After the two changes, we then have:
[Equation]
where we have the optimal solution x
[Equation]
d D ϕ ( s ϕ , a ) .
A.4.1 Optimization techniques
The optimization problem in Equation (7) presents a couple of optimization challenges.
These challenges are the same as in DualDICE Nachum et al.
(2019) (page 5).
Namely,
• The quantity ( ν − ν B π ϕe )( s ϕ , a ) involves a conditional exepctation inside the square.
When the environment dynamics are stochastic and/or the action space is large or continuous, this quantity may not be readily optimized using stochastic techniques • Even once ν ∗ is determined, ( ν ∗ − ν ∗ B π ϕe )( s ϕ , a ) may not be easily computable due to the same reasons as above.
To overcome these challenges, Fenchal duality is invoked Rockafellar (1970) where for a convex function f , f ( x ) = max ζ x · ζ − f ∗ ( ζ ) , where f ∗ is the Fenchal conjugate of f .
When f ( x ) = 12 x 2 , f ∗ ( ζ ) = 12 ζ 2 .
Thus, the optimization given in Equation 7 becomes:
[Equation]
Then by the interchangability principle Rockafellar and Wets (1998)
; Dai et al.
(2016), we replace the inner maximization over scalars to a maximization over ζ : S ϕ ×
A → R , resulting in the optimization given by Equation (8):
[Equation]
On applying the KKT conditions to the inner optimization, which is convex and quadratic, for any ν the optimal ζ ( s ϕ , a ) =
( ν − B π ϕe ν )( s ϕ , a ) .
Therefore, for the optimal ν ∗ , we have ζ ∗ ( s ϕ , a ) =
( ν ∗ − B π ϕe ν ∗ )( s ϕ , a ) =
d π ϕe ( s ϕ , a ) /d
D ϕ
( s ϕ , a ) .
This derivation with DualDICE naturally extends to
BestDICE since BestDICE solves the same optimization with the added constraints that E [ ζ ] =and ζ ≥ 0 , which are properties we know the true ratios would satisfy.
A.5 Additional Tabular Experiments and Details
In this section, we include the remaining tabular experiments and information.
• The true value of ρ ( π e ) was determined by executing π e for 200 episodes and averaging the results.
• In all experiments, we use γ = 0 .•
In the ˆ d π ϕe vs. d π ϕe plots, we use batch size of 300 .
In practice, Assumptions 2 and 3 may be hard to validate.
To study the impact of violation of these assumptions, we consider the following variations (Figure 5) to the original TwoPath MDP presented in Figure 2.
In all these MDPS, we only aggregate s 1 and s 2 .
The policies in the last two MDPs are the same.
Figure 6 illustrates the accuracy of AbstractBestDICE’s true ratio estimation.
In general, we do see AbstractBestDICE fails to compute the true abstract ratios when these assumptions are violated.
A.6 Additional Function Approximation Experiments and Details A.6.1 Oracle ρ ( π e ) Values
On each domain, we executed π e for 200 episodes and averaged the results.
A.6.2 Policies
For each of the domains, we used the following policies:
• Reacher: We trained a policy using PPO Schulman et al.
π e was the trained policy after 100 k time-steps with a standard deviation of 0 .
1 on the action dimensions while π b used 0 .
5 as the standard deviation.
• Walker2D: We trained a policy using PPO Schulman et al.
π e was the trained policy after 100 k time-steps with a standard deviation of 0 .
1 on the action dimensions while π b used 0 .
5 as the standard deviation.
• Pusher: We trained a policy using PPO Schulman et al.
π e was the trained policy after 100 k time-steps with a standard deviation of 0 .
1 on the action dimensions while π b used 0 .
5 as the standard deviation.
• AntUMaze:
We used the policies made available Fu et al.
π e was the final 10th snapshot saved and π b was the 5th snapshot.
Each also had 0 .
1 standard deviation on the action dimensions.
[Figure]
A.6.3 Trajectory Length
For each of the domains, the trajectory length is: 200 for Reacher, 500 for Walker2D, 300 for Pusher, and 500 for AntUMaze.
A.6.4 Hyperparameters
For BestDICE and AbstractBestDICE, we fixed the following hyperparameters:
• γ = 0 .
995 in all experiments.
• Neural net architecture: All neural networks are 2 layers with 64 hidden units using tanh activation.
• Unit mean constraint learning rate Zhang et al.
(2020b)
; Yang et al.
(2020): λ = 1 e − 3 .
• Optimizer: Adam optimizer with default parameters in Pytorch.
• Positivity constraint: squaring function on the last layer of the neural network.
We conducted a search for the learning rate of ν ( α ν ) and learning rate of ζ ( α ζ ), .
The learning rate search for ( α ν , α ζ ) was over { (5 e − 5 , 5 e − 5 ) , (1 e − 4 , 1 e − 4 ) , (3 e − 4 , 3 e − 4 ) , (7 e − 4 , 7 e − 4 ) , (1 e − 3 , 1 e − 3 ) } .
The optimal hyperparameters ( α ν = α ζ ) for each environment and batch size were:
[Table] [Table]
A.6.5 Empirical Estimator
In practice we use a weighted importance sampling Yang et al.
(2020) approach for the function approximation cases to estimate ρ ( π e ) (same for BestDICE):
[Equation]
[Table]
A.6.6 Misc Abstraction Details
• For Walker2D, we modified the default reward function from incremental distance covered at each time-step to distance from start location at each time-step to ensure Assumption 1 is satisfied.
• For AntUMaze, the reward function is originally r ( s ′ )
i.e. it is based on the next state that the ant moves to.
To ensure Assumption 1 is satisfied, we changed this reward function to be of the current state, r ( s ) .
A.6.7 Additional Results
Additional Hyperparameter Robustness Results In general, we can see AbstractBestDICE can be much more robust than BestDICE to hyperparameter tuning.
Training Stability In Figure 8 we show that ϕ can improve training stability.
Abstract Quality and Data-Efficiency .
We find that not all abstractions that satisfy Assumptionlead to better performance.
For example, the following are valid abstractions on the Reacher task: 1) the Euclidean distance between the arm and goal, s ϕ ∈ R and the 3D vector between the arm and goal, s
ϕ ∈ R 3 (Figure 9).
However, in practice we found that these were unreliable.
One possible reason for this unreliability is that these abstractions are incredibly extreme and the algorithm may be unable to differentiate between abstract state, resulting in outputting similar ζ ϕ ( s ϕ , a ) ∀ s ϕ .
A.7 Hardware For Experiments
• Distributed cluster on HTCondor framework • Intel(R)
Xeon(R) CPU E5-2470 0 @ 2.30GHz • RAM: 5GB • Disk space: 4GB
.d
π φe ( s φ ,a ) 0 .0 .0 .0 .0 .0 .
[Figure] [Figure]
a
[Figure]
Assumption 2 (transition dynamics) violated 0 .0 .0 .
.d
π φe ( s φ ,a ) 0 .0 .0 .0 .0 .0 .0 .
[Figure]
( s
[Figure] [Figure]
[Figure]
) Assumption 3 ( π e action equality) violated
[Figure] [Figure]
d
[Figure] [Figure]
[Figure] [Figure]
t i v
[Figure] [Figure]
( π
[Figure] [Figure]
(a) Reacher (Batch size: 5)
[Figure]
[Figure] [Figure]
(Batch
[Figure]
(d) Walker2D (Batch size: 5)
[Figure]
(e)
[Figure] [Figure]
[Figure] [Figure]