SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhang-Wei Hong $^{1,*}$ , Aviral Kumar $^{2}$ , Sathwik Karnik $^{1}$ , Abhishek Bhandwaldar $^{3}$ , Akash Srivastava $^{3}$ , Joni Pajarinen $^{4}$ , Romain Laroche $^{6}$ , Abhishek Gupta $^{5}$ , Pulkit Agrawal $^{1}$
Offline policy learning is aimed at learning decision-making policies using existing datasets of trajectories without collecting additional data. The primary motivation for using reinforcement learning (RL) instead of supervised learning techniques such as behavior cloning is to find a policy that achieves a higher average return than the trajectories constituting the dataset. However, we empirically find that when a dataset is dominated by suboptimal trajectories, state-of-the-art offline RL algorithms do not substantially improve over the average return of trajectories in the dataset. We argue this is due to an assumption made by current offline RL algorithms of staying close to the trajectories in the dataset. If the dataset primarily consists of sub-optimal trajectories, this assumption forces the policy to mimic the suboptimal actions. We overcome this issue by proposing a sampling strategy that enables the policy to only be constrained to "good data" rather than all actions in the dataset (i.e., uniform sampling). We present a realization of the sampling strategy and an algorithm that can be used as a plug-and-play module in standard offline RL algorithms. Our evaluation demonstrates significant performance gains in 72 imbalanced datasets, D4RL dataset, and across three different offline RL algorithms. Code is available at https://github.com/Improbable-AI/dw-offline-rl.
Offline reinforcement learning (RL) [23, 27] aims to learn a decision-making policy that maximizes the expected return (i.e., the sum of rewards over time) using a pre-collected dataset of trajectories, making it appealing for applications where data collection is infeasible or expensive (e.g., recommendation systems [28]). Without loss of generality, it can be assumed that the dataset is generated from an unknown policy  $\pi_{\mathcal{D}}(a|s)$ , also known as the behavior policy [24]. The goal in offline RL is to learn a policy,  $\pi_{\theta}(a|s)$  with parameters  $\theta$ , that exceeds the performance of the behavior policy. In offline RL, a widely recognized issue is the overestimation of  $Q$ -values for out-of-distribution state-action pairs, leading to suboptimal policies [8, 20, 22]. This stems from incomplete coverage of the state-action space in the dataset, causing the learning algorithm to consider absent states and actions during optimization.
Most state-of-the-art offline RL algorithms [7, 9, 20, 22, 25] mitigate the issue of OOD Q-values by constraining the distribution of actions of the learned policy  $\pi_{\theta}(a|s)$ , to be close to the distribution of actions in the dataset. This results in a generic objective with the following form:
$$
\max _ {\pi_ {\theta}} J (\pi_ {\theta}) - \alpha \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ \mathcal {C} (s, a) \right],
$$
where  $J(\pi_{\theta})$  denotes the expected return of the policy  $\pi_{\theta}$ ,  $\mathcal{D}$  denotes the dataset,  $\mathcal{C}$  is a regularization term that penalizes the policy  $\pi_{\theta}$  for deviating from the state-action pairs in the dataset, and  $\alpha$  is the
*Correspondence: zwhong@mit.edu, ImprobableAI Lab, Massachusetts Institute of Technology<sup>1</sup>, RAIL Lab, UC Berkeley<sup>2</sup>, MIT-IBM Lab<sup>3</sup>, Aalto University<sup>4</sup>, University of Washington<sup>5</sup>, and independent researcher<sup>6</sup>.
Figure 1: The dots represent actions in the dataset, where imbalanced datasets have more low-return actions. (a) Regularized offline RL algorithms [22, 7, 18] equally regularize the policy  $\pi_{\theta}$  on each action, leading to imitation of low-return actions and a low-performing  $\pi_{\theta}$ . The color under the curves shows the policy's performance  $J(\pi_{\theta})$ , with red indicating higher performance and blue indicating lower performance. (b) Re-weighting the dataset based on actions' returns allows the algorithm to only regularize on actions with high returns, enabling the policy  $\pi_{\theta}$  to imitate high-return actions while ignoring low-return actions.
hyper-parameter balancing the conflicting objectives of maximizing returns while also staying close to the data distribution. This prevents offline RL algorithms from learning behaviors that produce action distributions that diverge significantly from the behavior policy.
An easy-to-understand example of choice for  $\mathcal{C}$  is the squared distance between the policy and the data [7],  $\mathcal{C}(s,a)\coloneqq \| \pi_{\theta}(s) - a\| _2^2$ , where  $\pi_{\theta}(s)$  denotes the mean of the action distribution  $\pi_{\theta}(.|s)$  and  $a$  is an action sampled from the dataset. When the collected dataset is good, i.e., mostly comprising high-return trajectories, staying close to the data distribution is aligned with the objective of maximizing return, and existing offline RL algorithms work well. However, in scenarios where the dataset is skewed or imbalanced, i.e., contains only a few high-return trajectories and many low-return trajectories, staying close to the data distribution amounts to primarily imitating low-performing actions and is, therefore, detrimental. Offline RL algorithms struggle to learn high-return policies in such scenarios [12]. We present a method that overcomes this fundamental limitation of offline RL. Our method is plug-and-play in the sense that it is agnostic to the choice of the offline RL algorithm.
Our key insight stems from the observation that current methods are unnecessarily conservative by forcing the policy  $\pi_{\theta}$  to stay close to all the data. Instead, we would ideally want the policy  $\pi_{\theta}$  to be close to the best parts of the offline dataset. This suggests that we should constrain  $\pi_{\theta}$  to only be close to state-action pairs that would be generated from a policy that achieves high returns, for instance, the (nearly) optimal policy  $\pi^{*}$ . In offline scenarios, where collecting additional data is prohibited, to mirror the data distribution of  $\pi^{*}$  as much as possible, we can re-weight existing data (i.e., importance sampling [17]). We instantiate this insight in the following way: Represent the distribution induced by a better policy  $\pi_{\mathcal{D}_w}$  (initially unknown) by re-weighting data points in the dataset with importance weights  $w(s,a)$  and denote this distribution as  $\mathcal{D}_w(s,a)$ . Under this weighting, the offline RL algorithm's training objective can be written as:
$$
\max  _ {\pi_ {\theta}} J (\pi_ {\theta}) - \alpha \mathbb {E} _ {(s, a) \sim \mathcal {D}} [ w (s, a) \mathcal {C} (s, a) ],
$$
Solving for  $\pi_{\theta}$  using the re-weighted objective constrains the policy to be close to the better policy  $\pi_{\mathcal{D}}$ , and, therefore, allows learning of performant policies. The key challenge is determining the weights  $w$  since the state-action distribution of the better policy  $\pi_{\mathcal{D}_w}$  is initially unknown. To address this, we employ off-policy evaluation techniques [32, 26, 44] to connect the data distribution  $D_w$  with the expected return of the policy  $\pi_{\mathcal{D}_w}$  that would generate it. This allows one to optimize the importance weights with respect to its expected return as follows:
$$
\max _ {w} J (\pi_ {\mathcal {D} _ {w}}) = \mathbb {E} _ {(s, a) \sim \mathcal {D} _ {w}} [ r (s, a) ] \approx \mathbb {E} _ {(s, a) \sim \mathcal {D}} [ w (s, a) r (s, a) ].
$$
Here  $r(s,a)$  denotes reward of state-action pair  $(s,a)$ . By exploiting this connection, we optimize the importance weights  $w$  to maximize the expected return  $J(\pi_{\mathcal{D}_w})$  of its corresponding policy  $\pi_{\mathcal{D}_w}$  subject to necessary constraints (i.e., Bellman flow conservation constraint [35, 32]). This enables us to obtain a better importance weights  $w$ .
We evaluate our method with state-of-the-art offline RL algorithms [22, 18] and demonstrate performance gain on 72 imbalanced datasets [12, 6]. Our method significantly outperforms prior work [12]
in challenging datasets with more diverse initial states and fewer trajectories ( $20 \times$  smaller than existing datasets). These datasets pose greater challenges, yet they are crucial for practical applications since real-world datasets are often small and exhibit diverse initial states (e.g., robots with different starting positions, where data comes from human teleoperation).
Typical (online) RL. Reinforcement learning is a formalism that enables us to optimize an agent's policy in a Markov decision process (MDP [35]). The agent (i.e., decision-maker) starts from an initial state  $s_0$  sampled from an initial state distribution  $\rho_0(\cdot)$ . At each timestep  $t$ , the agent perceives the state  $s_t$ , takes an action  $a_t \sim \pi(.|s_t)$  with its policy  $\pi$ , receives a reward  $r_t = r(s_t, a_t)$  from the environment, and transitions to a next state  $s_{t+1}$  sampled from the environment dynamics  $\mathcal{T}(.|s_t, a_t)$  until reaching terminal states. The goal in RL is to learn a policy  $\pi$  to maximize the  $\gamma$ -discounted expected, infinite-horizon return  $J^{\gamma}(\pi) = \mathbb{E}_{s_0 \sim \rho_0, a_t \sim \pi(.|s_t), s_{t+1} \sim \mathcal{T}(.|s_t, a_t)}\left[\sum_{t=0}^{\infty} \gamma^t r(s_t, a_t)\right]$ . Typical (online) RL algorithms estimate the policy  $\pi$ 's expected return (policy evaluation) from trajectories  $\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \dots)$  generated by rolling out  $\pi$ , and update the policy  $\pi$  toward increasing  $J^{\gamma}(\pi)$  (policy improvement), and repeat the processing by performing rollouts with the updated policy.
Offline RL. With no interaction with the environment allowed during the course of learning, offline RL algorithms aim to learn a policy  $\pi$  that maximizes return, entirely using a fixed dataset  $\mathcal{D}$  that was collected by an arbitrary and unknown "behavior policy"  $\pi_{\mathcal{D}}$  (e.g., humans or pre-programmed controllers). These methods typically aim to estimate the return of a policy  $\pi$  via techniques such as Q-learning or actor-critic, only using batches of state-action pairs  $(s_t, a_t)$  uniformly drawn from  $\mathcal{D}$ . We will denote this estimate value of the return as  $\widehat{J}_{\mathcal{D}}^{\gamma}(\pi)$ . The dataset  $\mathcal{D}$  consists of  $N$  trajectories rolled out by  $\pi_{\mathcal{D}}$ :
$$
\mathcal {D} := \left\{\tau_ {i} = \left(s _ {0}, a _ {0}, r _ {0}, s _ {1}, a _ {1}, r _ {1} \dots s _ {T _ {i}}\right) _ {i} \right\} _ {i = 1} ^ {N}, \tag {1}
$$
where  $T_{i}$  denotes the length of  $\tau_{i}$ . In practice, a limit on trajectory length is required since we can only collect finite-length trajectories [34]. When the states that the policy  $\pi$  would encounter and the actions that  $\pi$  would take are not representative in the dataset  $\mathcal{D}$ , the estimated return  $\widehat{J}_{\mathcal{D}}^{\gamma}(\pi)$  is typically inaccurate [8, 20, 22]. Thus, most offline RL algorithms learn the policy  $\pi$  with pessimistic or conservative regularization that penalizes shift of  $\pi$  from the behavior policy  $\pi_{\mathcal{D}}$  that collected the dataset. Typically, implicitly or explicitly, the policy  $\pi$  learned by most offline RL algorithms can be thought of as optimizing the following regularized objective:
$$
\max  _ {\pi} \widehat {J} _ {\mathcal {D}} ^ {\gamma} (\pi) - \alpha \mathbb {E} _ {(s _ {t}, a _ {t}) \sim \mathcal {D}} \left[ \mathcal {C} (s _ {t}, a _ {t}) \right], \tag {2}
$$
where  $\mathcal{C}$  measures some kind of divergence (e.g., Kullback-Leibler divergence [7]) between  $\pi$  and  $\pi_{\mathcal{D}}$ , and  $\alpha \in \mathbb{R}^{+}$  denotes the strength of regularization.
3 Problem Statement: Unnecessary Conservativeness in Imbalanced Datasets
In this section, we describe the issue of offline RL in imbalanced datasets. While algorithms derived from the regularized offline RL objective (Equation 2) attain good performance on several standard benchmarks [6], recent work [12] showed that it leads to "unnecessary conservativeness" on imbalanced datasets [12] due to the use of constant regularization weight on each state-action pairs  $(s,a)$  in Equation 2. To illustrate why, we start by defining imbalance of a dataset  $\mathcal{D}$  using the positive-sided variance of the returns of the dataset (RPSV [12] defined in Definition 3.1). In essence, RPSV measures the dispersion of trajectory returns in the dataset. It indicates the room for improvement of the dataset. Figure 2 illustrates the distribution of trajectory returns in imbalanced datasets with high and low RPSV. Datasets with a low RPSV exhibit a pronounced concentration of returns around the mean value, whereas datasets with a high RPSV display a return distribution that extends away from the mean, towards higher returns. Intuitively, a dataset with high RPSV has trajectories with far higher returns than the average return of the dataset, indicating high chances of finding better data distribution through reweighting. Throughout this paper, we will use the term imbalanced datasets to denote datasets with high RPSV.
Definition 3.1 (Dataset imbalance). RPSV of a dataset,  $\mathbb{V}_{+}[G(\tau_{i})]$ , corresponds to the second-order moment of the positive component of the difference between trajectory return:  $G(\tau_i)\coloneqq \sum_{t = 0}^{T_i - 1}\gamma^t r(s_t^i,a_t^i)$  and its expectation, where  $\tau_{i}$  denote trajectory in the dataset:
$$
\mathbb {V} _ {+} [ G (\tau_ {i}) ] \doteq \mathbb {E} _ {\tau_ {i} \sim \mathcal {D}} \left[ (G (\tau_ {i}) - \mathbb {E} _ {\tau_ {i} \sim \mathcal {D}} [ G (\tau_ {i}) ]) _ {+} ^ {2} \right] \quad \text {w i t h} \quad x _ {+} = \max  \{x, 0 \}, \tag {3}
$$
Imbalanced datasets are common in real-world scenarios, as collecting high-return trajectories is often more costly than collecting low-return ones. An example of an imbalanced offline dataset is autonomous driving, where most trajectories are from average drivers, with limited data from very good drivers. Due to the dominance of low-return trajectories, state-action pairs  $(s,a)$  from these trajectories are oversampled in Equation 2. Consequently, optimizing the regularized objective (Equation 2) would result in a policy that closely imitates the actions from low-performing trajectories that constitute the majority of dataset  $\mathcal{D}$ , but ideally, we want the policy to imitate actions only on state-action pairs from high-performing trajectories. How-
Figure 2: Return distribution of datasets with high and low RPSV. Low RPSV datasets have returns centered at the mean, while high RPSV datasets have a wider distribution extending towards higher returns. See Appendix A.4 for details.
ever, current offline RL algorithms [21, 18, 7] use constant regularization weight  $\alpha$  (Equation 2). As a result, each state-action pairs is weighted equally, which leads the algorithm to be unnecessarily conservative on all data (i.e., imitating all actions of state-action pairs in the dataset). Further analysis of imbalanced datasets can be found in Appendix A.1.
4 Mitigating Unnecessary Conservativeness By Weighting Samples
In this section, we seek to develop an approach to address the unnecessary conservativeness issue (Section 3) of regularized offline RL algorithms in imbalanced datasets. Adding more experiences from high-performing policies to the dataset would regularize the policy to keep close to high-performing policies and hence easily mitigate the unnecessary conservativeness issue. Though collecting additional experiences (i.e., state-action pairs  $(s,a)$ ) from the environment is prohibited in offline RL, importance sampling [32] can emulate sampling from another dataset  $\mathcal{D}_w$  since the weighting can be regarded as the density ratio shown below:
$$
w (s, a) = \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)}, \tag {4}
$$
where  $\mathcal{D}_w(s,a)$  and  $\mathcal{D}(s,a)$  denote the probability density of state-action pairs  $(s,a)$  in dataset  $\mathcal{D}_w$  and  $\mathcal{D}$ . Note that  $\mathcal{D}_w$  is unknown but implicitly defined through a given weighting function  $w$ . This allows us to adjust the sampling distribution that we train the policy  $\pi$  to, as suggested in the following equivalence:
$$
\max  _ {\pi} \widehat {J} _ {\mathcal {D}} ^ {\gamma} (\pi) - \alpha \mathbb {E} _ {(s _ {t}, a _ {t}) \sim \mathcal {D}} [ w (s, a) \mathcal {C} (s _ {t}, a _ {t}) ] \Longleftrightarrow \max  _ {\pi} \widehat {J} _ {\mathcal {D}} ^ {\gamma} (\pi) - \alpha \mathbb {E} _ {(s _ {t}, a _ {t}) \sim \mathcal {D} _ {w}} [ \mathcal {C} (s _ {t}, a _ {t}) ]. \tag {5}
$$
The remaining question is: how can we determine the weighting  $w(s, a)$  so that we emulate sampling from a better dataset  $\mathcal{D}_w$  collected by a policy that achieves higher return than the behavior policy  $\pi_{\mathcal{D}}$  that collected the original dataset  $\mathcal{D}$ ?
4.1 Optimizing the Weightings: Emulating Sampling from High-Performing Policies
Our goal is to discover a weighting function  $w$  that can emulate drawing state-action samples from a better dataset  $\mathcal{D}_w$  that is collected by an alternative behavior policy  $\pi_{\mathcal{D}_w}$  with higher return than the behavior policy  $\pi_{\mathcal{D}}$  that collected the original dataset  $\mathcal{D}$  (i.e.,  $J^{\gamma}(\pi_{\mathcal{D}_w}) \geq J^{\gamma}(\pi_{\mathcal{D}})$ ). We make use of density-ratio-based off-policy evaluation methods [32, 29, 44] to determine if a weighting function corresponds to a high-return policy. Note that we do not propose a new off-policy evaluation approach but rather apply the existing off-policy evaluation technique in our problem. By using these
techniques, we can relate the weighting  $w$  to the expected return of the alternative behavior policy  $J(\pi_{\mathcal{D}_w})$  via importance sampling formulation as follows:
$$
J ^ {\gamma} \left(\pi_ {\mathcal {D} _ {w}}\right) \approx \mathbb {E} _ {(s, a) \sim \mathcal {D} _ {w}} [ r (s, a) ] = \mathbb {E} _ {(s, a) \sim \mathcal {D}} [ w (s, a) r (s, a) ]. \tag {6}
$$
In Equation 6,  $J^{\gamma}(\pi_{\mathcal{D}_w})$  evaluates the quality of a given weighting function  $w$ . It also provides a feasible objective to optimize  $w$ , since it only requires obtaining samples from the original dataset  $\mathcal{D}$ . However, it is important to note that Equation 6 measures  $\gamma$ -discounted return only when the the dataset  $\mathcal{D}_w$  represents a stationary state-action distribution that satisfies Bellman flow conservation constraint [35, 32, 44] in the MDP, as shown in the following equation:
$$
\mathcal {D} _ {w} \left(s ^ {\prime}\right) = (1 - \gamma) \rho_ {0} \left(s ^ {\prime}\right) + \gamma \sum_ {s, a} \mathcal {T} \left(s ^ {\prime} \mid s, a\right) \mathcal {D} _ {w} (s, a) \forall s ^ {\prime} \in \mathcal {S}, \quad \mathcal {D} _ {w} \left(s ^ {\prime}\right) := \sum_ {a ^ {\prime} \in \mathcal {A}} \mathcal {D} _ {w} \left(s ^ {\prime}, a ^ {\prime}\right) \tag {7}
$$
where  $\mathcal{S}$  and  $\mathcal{A}$  denote the state and action spaces, respectively, and the discount factor  $\gamma$  determines what discount factor corresponds to in  $J^{\gamma}(\pi_{\mathcal{D}_w})$  in Equation 6. We slightly abuse the notation, denoting state marginal as  $\mathcal{D}_w(s')$ .
To estimate  $J^{\gamma}(\pi_{\mathcal{D}_w})$  from the weighting function  $w$ , it is required to impose Bellman flow conservation constraint (Equaton 7) on  $w$ . However, it is difficult to impose this constraint due to the dependence of initial state distribution  $\rho_0$  in Equation 7. Estimating  $\rho_0$  from the first state of each trajectory in the dataset is an option, but it is infeasible when the trajectories do not consistently start from initial states sampled from the distribution  $\rho_0$ . While we could make the assumption that all trajectories begin from initial states sampled from  $\rho_0$ , it would limit the applicability of our method to datasets where trajectories start from arbitrary states. We thus choose not to make this assumption since current offline RL algorithms do not require it.
Instead, since the Bellman flow conservation constraint (Equation 7) only depends on the initial state distribution  $\rho_0$  when  $\gamma \neq 1$ , it is possible to bypass this dependence, if we maximize the undiscounted return  $J(\pi_{\mathcal{D}_w}) = J^{\gamma = 1}(\pi_{\mathcal{D}_w})$  (i.e., setting  $\gamma = 1$  in Equation 6) of the alternative behavior policy  $\pi_{\mathcal{D}_w}$ . While it deviates from the RL objective presented in Equation 2, undiscounted return is often more aligned with the true objective in various RL applications, as suggested in [13]. Many RL algorithms resort to employing discounted return as an approximation of undiscounted return instead due to the risk of divergence when estimating the undiscounted return using Q-learning [39]. Thus, we constrain the weighting function  $w$  to satisfy the Bellman flow conservation constraint with  $\gamma = 1$  as shown below:
$$
\mathcal {D} _ {w} \left(s ^ {\prime}\right) = \sum_ {s, a} \mathcal {T} \left(s ^ {\prime} \mid s, a\right) \mathcal {D} _ {w} (s, a) \forall s ^ {\prime} \in \mathcal {S}. \tag {8}
$$
To connect the constraint in Equation 8 to the objective in Equation 6, we rewrite Equation 8 in terms of weightings  $w$  according to [29, 32], as shown below:
$$
\mathcal {D} \left(s ^ {\prime}\right) w \left(s ^ {\prime}\right) = \sum_ {s, a} \mathcal {T} \left(s ^ {\prime} \mid s, a\right) w (s, a) \forall s ^ {\prime} \in \mathcal {S}, \quad w (s) := \sum_ {a \in \mathcal {A}} \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} \tag {9}
$$
where  $w(s)$  denotes state marginal weighting. Putting the objective (Equation 6) and the constraint (Equation 9) together, we optimize  $w$  to maximize the undiscounted expected return of the corresponding alternative behavior policy  $\pi_{\mathcal{D}_w}$ , as shown in the following:
$$
\max  _ {w} J \left(\pi_ {\mathcal {D} _ {w}}\right) = \mathbb {E} _ {(s, a) \sim \mathcal {D}} [ w (s, a) r (s, a) ] \tag {10}
$$
subject to  $\mathbb{E}_{(s,a,s')\sim \mathcal{D}}[w(s') - w(s,a)\mid s'] = 0\forall s'\in S.$
As the weightings  $w$  can be viewed as the density ratio (Equation 4), we call our method as Density-ratio Weighting (DW). We then re-weight offline RL algorithm, as shown in Equation 5. Note that while these weights correspond to  $(s,a)$  in the dataset, this is sufficient to reweight policy optimization for offline RL.
4.2 Practical Implementation
Optimizing weightings. We begin by addressing the parameterization of the weighting function  $w(s, a)$  and its state marginal  $w(s')$  in Equation 10. Though state marginal  $w(s')$  can derived from
See Equation 24 in [29]
summing  $w(s, a)$  over action space  $\mathcal{A}$ , as defined in Equation 9, it can difficult to take summation over a continuous or infinite action space. Thus we opt to parameterize the weightings  $w(s, a)$  and its state marginal  $w(s')$  separately. By using the identities [32]  $\mathcal{D}_w(s, a) = \mathcal{D}_w(s) \pi_{\mathcal{D}_w}(a|s)$  and  $\mathcal{D}(s, a) = \mathcal{D}(s) \pi_{\mathcal{D}}(a|s)$ , we can represent  $w(s, a)$  as the product of two ratios:
$$
w (s, a) \doteq \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} = \frac {\mathcal {D} _ {w} (s) \pi_ {\mathcal {D} _ {w}} (a | s)}{\mathcal {D} (s) \pi_ {\mathcal {D}} (a | s)} = \frac {\mathcal {D} _ {w} (s)}{\mathcal {D} (s)} \times \frac {\pi_ {\mathcal {D} _ {w}} (a | s)}{\pi_ {\mathcal {D}} (a | s)}. \tag {11}
$$
Michel et al. [31] showed that ratios can be parameterized by neural networks with exponential output. Thus, we represent state-action weighting  $w(s,a)$  as  $w_{\phi,\psi}(s,a)$  and its state marginal as  $w_{\phi}(s)$ , as shown below:
$$
w _ {\phi , \psi} (s, a) = \exp \phi (s) \exp \psi (s, a), \quad w _ {\phi} (s) = \exp \phi (s) \tag {12}
$$