SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiong-Hui Chen $^{1,2,4,*}$ , Yang Yu $^{1,2,4,*,\dagger}$ , Zheng-Mao Zhu $^{1,2}$ , Zhihua Yu $^{1,2,3}$ , Zhenjun Chen $^{1,2,3}$ , Chenghe Wang $^{1,2}$ , Yinan Wu $^{3}$ , Hongqiu Wu $^{1,2}$ , Rong-Jun Qin $^{1,2,4}$ , Ruijin Ding $^{5}$ , Fangsheng Huang $^{3}$
$^{1}$  National Key Laboratory for Novel Software Technology, Nanjing University  $^{2}$  School of Artificial Intelligence, Nanjing University,  $^{3}$  Meituan,  $^{4}$  Polixir.ai,  $^{5}$  Tsinghua University  $\{\mathrm{chenxh}, \mathrm{yuy}, \mathrm{zhuzm}\} @$  lamda.nju.edu.cn,  $\{\mathrm{yuzh}, \mathrm{chenzj}\} @$  smail.nju.edu.cn wangch@lamda.nju.edu.cn, wuyinan02@meituan.com,  $\{\mathrm{wuhq}, \mathrm{qinrj}\} @$  lamda.nju.edu.cn drj17@mails.tsinghua.edu.cn, huangfangsheng@meituan.com
An accurate environment dynamics model is crucial for various downstream tasks in sequential decision-making, such as counterfactual prediction, off-policy evaluation, and offline reinforcement learning. Currently, these models were learned through empirical risk minimization (ERM) by step-wise fitting of historical transition data. This way was previously believed unreliable over long-horizon rollouts because of the compounding errors, which can lead to uncontrollable inaccuracies in predictions. In this paper, we find that the challenge extends beyond just long-term prediction errors: we reveal that even when planning with one step, learned dynamics models can also perform poorly due to the selection bias of behavior policies during data collection. This issue will significantly mislead the policy optimization process even in identifying single-step optimal actions, further leading to a greater risk in sequential decision-making scenarios. To tackle this problem, we introduce a novel model-learning objective called adversarial weighted empirical risk minimization (AWRM). AWRM incorporates an adversarial policy that exploits the model to generate a data distribution that weakens the model's prediction accuracy, and subsequently, the model is learned under this adversarial data distribution. We implement a practical algorithm, GALILEO, for AWRM and evaluate it on two synthetic tasks, three continuous-control tasks, and a real-world application. The experiments demonstrate that GALILEO can accurately predict counterfactual actions and improve various downstream tasks, including offline policy evaluation and improvement, as well as online decision-making.
A good environment dynamics model for action-effect prediction is essential for many downstream tasks. For example, humans or agents can leverage this model to conduct simulations to understand future outcomes, evaluate other policies' performance, and discover better policies. With environment models, costly real-world trial-and-error processes can be avoided. These tasks are vital research problems in counterfactual predictions [52, 1], off-policy evaluation (OPE) [32, 35], and offline reinforcement learning (Offline RL) [31, 59, 12, 13, 10]. In these problems, the core role of the models is to answer queries on counterfactual data unbiasedly, that is, given states, correctly answer what might happen if we were to carry out actions unseen in the training data. However, addressing counterfactual queries differentiates environment model learning from standard supervised learning (SL), which directly fits the offline dataset for empirical risk minimization (ERM). In essence,
*The authors contribute equally
$\dagger$  Yang Yu is the corresponding author.
(a) a selection bias example
Figure 1: An example of selection bias and predictions under counterfactual queries. Suppose a ball locates in a 2D plane whose position is  $s_t = (x_t, y_t)$  at time  $t$ . The ball will move to  $s_{t+1} = (x_{t+1}, y_{t+1})$  according to  $x_{t+1} = x_t + 1$  and  $y_{t+1} \sim \mathcal{N}(y_t + a_t, 2)$ . Here,  $a_t$  is chosen by a control policy  $a_t \sim \pi_\phi(a|s_t) = \mathcal{N}((\phi - y_t)/15, 0.05)$  parameterized by  $\phi$ , which tries to keep the ball near the line  $y = \phi$ . In Fig. 1(a), the behavior policy  $\mu$  is  $\pi_{62.5}$ . Fig. 1(b) shows the collected training data and the learned models' prediction of the next position of  $y$ . Besides, the dataset superficially presents the relation that the corresponding next  $y$  will be smaller with a larger action. However, the truth is not because the larger  $a_t$  causes a smaller  $y_{t+1}$ , but the policy selects a small  $a_t$  when  $y_t$  is close to the target line. Mistakenly exploiting the "association" will lead to local optima with serious factual errors, e.g., believed that  $y_{t+1} \propto \pi_\phi^{-1}(y_t|a) + a_t \propto \phi - 14a_t$ , where  $\pi_\phi^{-1}$  is the inverse function of  $\pi_\phi$ . When we estimate the response curves by fixing  $y_t$  and reassigning action  $a_t$  with other actions  $a_t + \Delta a$ , where  $\Delta a \in [-1, 1]$  is a variation of action value, we found that the model of SL indeed exploits the association and give opposite responses, while in AWRM and its practical implementation GALILEO, the predictions are closer to the ground truths ( $y_{t+1} \propto y_t + a_t$ ). The result is in Fig. 1(c), where the darker a region is, the more samples are fallen in. AWRM injects data collected by adversarial policies for model learning to eliminate the unidentifiability between  $y_{t+1} \propto \pi_\phi^{-1}(y_t|a) + a_t$  and  $wy_{t+1} \propto y_t + a_t$  in offline data.
(b) responses to actions in (c) responses to actions in the training data the action space
the problem involves training the model on one dataset and testing it on another with a shifted distribution, specifically, the dataset generated by counterfactual queries. This challenge surpasses the SL's capability as it violates the independent and identically distributed (i.i.d.) assumption [31].
In this paper, we concentrate on faithful dynamics model learning in sequential decision-making settings like RL. Specifically, we first highlight a distinct situation of distribution shift that can easily lead to catastrophic failures in the model's predictions: In many real-world applications, offline data is often gathered using a single policy with exhibiting selection bias, meaning that, for each state, actions are chosen unfairly. As illustrated in Fig. 1(a), to maintain the ball's trajectory along a target line, a behavior policy applies a smaller force when the ball's location is nearer to the target line. When a dataset is collected with such selection bias, the association between the states (location) and actions (force) will make SL hard to identify the correct causal relationship of the states and actions to the next states respectively (see Fig. 1(c)). Then when we query the model with counterfactual data, the predictions might be catastrophic failures. Finally, offline policy optimization based on this SL model, even for just seeking one-step optimal actions, will select also a totally opposite direction of policy improvement, making the offline policy learning system fail. The selection bias can be regarded as an instance of the distributional-shift problem in offline model-based RL, which has also received great attention [31, 59, 14, 34]. However, previous methods employing naive super
Figure 2: An illustration of the prediction error in counterfactual datasets. The error of SL is small only in training data  $(\phi = 62.5)$  but becomes much larger in the dataset "far away from" the training data. AWRM-oracle selects the oracle worst counterfactual dataset for training for each iteration (pseudocode is in Alg. 2) which reaches small MSE in all datasets and gives correct response curves (Fig. 1(c)). GALILEO approximates the optimal adversarial counterfactual data distribution based on the training data and model. Although the MSE of GALILEO is a bit larger than SL in the training data, in the counterfactual datasets, the MSE is on the same scale as AWRM-oracle.
vised learning for environment model learning tend to overlook this issue during the learning process,
addressing it instead by limiting policy exploration and learning in high-risk regions. So far, how to learn a faithful environment model that can alleviate the problem directly has rarely been discussed.
In this work, we focus on faithful environment model learning techniques. The work is first inspired by weighted empirical risk minimization (WERM), which is a typical solution to solve the selection bias problem in causal inference for individual treatment effects (ITEs) estimation in many scenarios like patients' treatment selection [26, 1, 41]. ITEs measure the effects of treatments on individuals by administering treatments uniformly and evaluating the differences in outcomes. To estimate ITEs from offline datasets with selection bias, they estimate an inverse propensity score (IPS) to reweight the training data, approximating the data distribution under a uniform policy, and train the model under this reweighted distribution. Compared with ITEs estimation, the extra challenge of faithful model learning in sequential decision-making settings include: (1) the model needs to answer queries on numerous different policies, resulting in various and unknown target data distributions for reweighting, and (2) the IPS should account for the cumulative effects of behavior policies on state distribution rather than solely focusing on bias of actions. To address these issues, we propose an objective called adversarial weighted empirical risk minimization (AWRM). For each iteration, AWRM employs adversarial policies to construct an adversarial counterfactual dataset that maximizes the model's prediction error, and drive the model to reduce the prediction risks under the adversarial counterfactual data distribution. However, obtaining the adversarial counterfactual data distribution is infeasible in the offline setting. Therefore, we derive an approximation of the counterfactual data distribution queried by the optimal adversarial policy and provide a tractable solution to learn a model from the approximated data distribution. As a result, we propose a practical approach named Generative Adversarial offLIne counterfactual Environment mOdel learning (GALILEO) for AWRM. Fig. 2 illustrates the difference in the prediction errors learned by these algorithms.
Experiments are conducted in two synthetic tasks, three continuous-control tasks, and a real-world application. We first verify that GALILEO can make accurate predictions on counterfactual data queried by other policies compared with baselines. We then demonstrate that the model learned by GALILEO is helpful to several downstream tasks including offline policy evaluation and improvement, and online decision-making in a large-scale production environment.
We first introduce weighted empirical risk minimization (WERM) through inverse propensity scoring (IPS), which is commonly used in individualized treatment effects (ITEs) estimation [43]. It can be regarded as a scenario of single-step model learning. We define  $M^{*}(y|x,a)$  as the one-step environment, where  $x$  denotes the state vector containing pre-treatment covariates (such as age and weight),  $a$  denotes the treatment variable which is the action intervening with the state  $x$ , and  $y$  is the feedback of the environment. When the offline dataset is collected with a behavior policy  $\mu(a|x)$  which has selection bias, a classical solution to handle the above problem is WERM through IPS  $\omega$  [48, 3, 29]:
Definition 2.1. The learning objective of WERM through IPS is formulated as
$$
\min  _ {M \in \mathcal {M}} \mathbb {E} _ {x, a, y \sim p _ {M ^ {*}} ^ {\mu}} [ \omega (x, a) \ell (M (y | x, a), y) ], \tag {1}
$$
where  $\omega(x, a) \coloneqq \frac{\beta(a|x)}{\mu(a|x)}$ ,  $\beta$  and  $\mu$  denote the policies in testing and training domains, and the joint probability  $p_{M^*}^\mu(x, a, y) \coloneqq \rho_0(x)\mu(a|x)M^*(y|x, a)$  in which  $\rho_0(x)$  is the distribution of state.  $\mathcal{M}$  is the model space.  $\ell$  is a loss function.
The  $\omega$  is also known as importance sampling (IS) weight, which corrects the sampling bias by aligning the training data distribution with the testing data. By selecting different  $\hat{\omega}$  to approximate  $\omega$  to learn the model  $M$ , current environment model learning algorithms employing reweighting are fallen into the framework. For standard ITEs estimation,  $\omega = \frac{1}{\hat{\mu}}$  (i.e.,  $\beta$  is a uniform policy) for balancing treatments, where  $\hat{\mu}$  is an approximated behavior policy  $\mu$ . Note that it is a reasonable weight in ITEs estimation: ITEs are defined to evaluate the differences of effect on each state under a uniform policy.
In sequential decision-making setting, decision-making processes in a sequential environment are often formulated into Markov Decision Process (MDP) [51]. MDP depicts an agent interacting with the environment through actions. In the first step, states are sampled from an initial state distribution  $x_0 \sim \rho_0(x)$ . Then at each time-step  $t \in \{0, 1, 2, \ldots\}$ , the agent takes an action  $a_t \in \mathcal{A}$  through a policy  $\pi(a_t | x_t) \in \Pi$  based on the state  $x_t \in \mathcal{X}$ , then the agent receives a reward  $r_t$  from
a reward function  $r(x_{t},a_{t})\in \mathbb{R}$  and transits to the next state  $x_{t + 1}$  given by a transition function  $M^{*}(x_{t + 1}|x_{t},a_{t})$  built in the environment.  $\Pi ,\mathcal{X}$ , and  $\mathcal{A}$  denote the policy, state, and action spaces.
3 Related Work
We give related adversarial algorithms for model learning in the following and leave other related work in Appx. F. In particular, ITEs in Rubin causal model [48] and causal effect estimation in structural causal model [38] attracted widespread attention in recent years [56, 55, 58, 7]. GANTIE [58] uses a generator to fill counterfactual outcomes for each data pair and a discriminator to judge the source (treatment group or control group) of the filled data pair. The generator is trained to minimize the output of the discriminator. [7] propose SCIGAN to extend GANITE to continuous ITEs estimation via a hierarchical discriminator architecture. In real-world applications, environment model learning based on Generative Adversarial Imitation Learning (GAIL) has also been adopted for sequential decision-making problems [25]. GAIL is first proposed for policy imitation [25], which uses the imitated policy to generate trajectories by interacting with the environment. The policy is learned with the trajectories through RL which maximizes the cumulative rewards given by the discriminator. [47, 11] use GAIL for environment model learning by regarding the environment model as the generator and the behavior policy as the "environment" in standard GAIL. [16] further inject the technique into a unified objective for model-based RL, which joints model and policy optimization. Our study reveals the connection between adversarial model learning and the WERM through IPS, where previous adversarial model learning methods can be regarded as partial implementations of GALILEO, explaining the effectiveness of the former.
4 Adversarial Counterfactual Environment Model Learning
In this section, we first propose a new offline model-learning objective for sequential decision-making setting in Sec. 4.1; In Sec. 4.2, we give a surrogate objective to the proposed objective; Finally, we give a practical solution in Sec. 4.3.
4.1 Problem Formulation
In scenarios like offline policy evaluation and improvement, it is crucial for the environment model to have generalization ability in counterfactual queries, as we need to query accurate feedback from  $M$  for numerous different policies. Referring to the formulation of WERM through IPS in Def. 2.1, these requirements necessitate minimizing counterfactual-query risks for  $M$  under multiple unknown policies, rather than focusing on a specific target policy  $\beta$ . More specifically, the question is: If  $\beta$  is unknown and can be varied, how can we generally reduce the risks in counterfactual queries? In this article, we call the model learning problem in this setting "counterfactual environment model learning" and propose a new objective to address the issue. To be compatible with multi-step environment model learning, we first define a generalized WERM through IPS based on Def. 2.1:
Definition 4.1. In an MDP, given a transition function  $M^{*}$  that satisfies  $M^{*}(x^{\prime}|x,a) > 0,\forall x\in \mathcal{X},\forall a\in \mathcal{A},\forall x^{\prime}\in \mathcal{X}$  and  $\mu$  satisfies  $\mu (a|x) > 0,\forall a\in \mathcal{A},\forall x\in \mathcal{X}$ , the learning objective of generalized WERM through IPS is:
$$
\min  _ {M \in \mathcal {M}} \mathbb {E} _ {x, a, x ^ {\prime} \sim \rho_ {M ^ {*}} ^ {\mu}} [ \omega (x, a, x ^ {\prime}) \ell_ {M} (x, a, x ^ {\prime}) ], \tag {2}
$$
where  $\omega(x, a, x') = \frac{\rho_{M^*}^\beta(x, a, x')}{\rho_{M^*}^\mu(x, a, x')}$ ,  $\rho_{M^*}^\mu$  and  $\rho_{M^*}^\beta$  the training and testing data distributions collected by policy  $\mu$  and  $\beta$  respectively. We define  $\ell_M(x, a, x') := \ell(M(x'|x, a), x')$  for brevity.
In an MDP, for any given policy  $\pi$ , we have  $\rho_{M^*}^\pi(x, a, x') = \rho_{M^*}^\pi(x)\pi(a|x)M^*(x'|x, a)$  where  $\rho_{M^*}^\pi(x)$  denotes the occupancy measure of  $x$  for policy  $\pi$ . This measure can be defined as  $(1 - \gamma)\mathbb{E}_{x_0 \sim \rho_0}\left[\sum_{t=0}^{\infty} \gamma^t \operatorname{Pr}(x_t = x|x_0, M^*)\right]$  [51, 25] where  $\operatorname{Pr}^\pi[x_t = x|x_0, M^*]$  is the discounted state visitation probability that the policy  $\pi$  visits  $x$  at time-step  $t$  by executing in the environment  $M^*$  and starting at the state  $x_0$ . Here  $\gamma \in [0,1]$  is the discount factor. We also define  $\rho_{M^*}^\pi(x, a) := \rho_{M^*}^\pi(x)\pi(a|x)$  for simplicity.
With this definition,  $\omega$  can be rewritten:  $\omega = \frac{\rho_{M^*}^\beta(x)\beta(a|x)M^*(x'|x,a)}{\rho_{M^*}^\mu(x)\mu(a|x)M^*(x'|x,a)} = \frac{\rho_{M^*}^\beta(x,a)}{\rho_{M^*}^\mu(x,a)}$ . In single-step environments, for any policy  $\pi$ ,  $\rho_{M^*}^\pi(x) = \rho_0(x)$ . Consequently, we obtain  $\omega = \frac{\rho_0(x)\beta(a|x)}{\rho_0(x)\mu(a|x)} = \frac{\beta(a|x)}{\mu(a|x)}$ , and the objective degrade to Eq. (1). Therefore, Def. 2.1 is a special case of this generalized form.
Remark 4.2.  $\omega$  is referred to as density ratio and is commonly used to correct the weighting of rewards in off-policy datasets to estimate the value of a specific target policy in off-policy evaluation
[35, 60]. Recent studies in offline RL also provide similar evidence through upper bound analysis, suggesting that offline model learning should be corrected to specific target policies' distribution using  $\omega$  [57]. We derive the objective from the perspective of selection bias correlation, further demonstrate the necessity and effects of this term.
In contrast to previous studies, in this article, we would like to propose an objective for faithful model learning which can generally reduce the risks in counterfactual queries in scenarios where  $\beta$  is unknown and can be varied. To address the problem, we introduce adversarial policies that can iteratively induce the worst prediction performance of the current model and propose to optimize WERM under the adversarial policies. In particular, we propose Adversarial Weighted empirical Risk Minimization (AWRM) based on Def. 4.1 to handle this problem.
Definition 4.3. Given the MDP transition function  $M^{*}$ , the learning objective of adversarial weighted empirical risk minimization through IPS is formulated as
$$
\min  _ {M \in \mathcal {M}} \max  _ {\beta \in \Pi} L \left(\rho_ {M ^ {*}} ^ {\beta}, M\right) = \min  _ {M \in \mathcal {M}} \max  _ {\beta \in \Pi} \mathbb {E} _ {x, a, x ^ {\prime} \sim \rho_ {M ^ {*}} ^ {\mu}} [ \omega (x, a | \rho_ {M ^ {*}} ^ {\beta}) \ell_ {M} (x, a, x ^ {\prime}) ], \tag {3}
$$
where  $\omega(x, a|\rho_{M^*}^\beta) = \frac{\rho_{M^*}^\beta(x, a)}{\rho_{M^*}^\mu(x, a)}$ , and the re-weighting term  $\omega(x, a|\rho_{M^*}^\beta)$  is conditioned on the distribution  $\rho_{M^*}^\beta$  of the adversarial policy  $\beta$ . In the following, we will ignore  $\rho_{M^*}^\beta$  and use  $\omega(x, a)$  for brevity.
In a nutshell, Eq. (3) minimizes the maximum model loss under all counterfactual data distributions  $\rho_{M^*}^\beta, \beta \in \Pi$  to guarantee the generalization ability for counterfactual queried by policies in  $\Pi$ .
4.2 Surrogate AWRM through Optimal Adversarial Data Distribution Approximation
The main challenge of solving AWRM is constructing the data distribution  $\rho_{M^*}^{\beta^*}$  of the best-response policy  $\beta^*$  in  $M^*$  since obtaining additional data from  $M^*$  can be expensive in real-world applications. In this paper, instead of deriving the optimal  $\beta^*$ , our solution is to offline estimate the optimal adversarial distribution  $\rho_{M^*}^{\beta^*}(x,a,x')$  with respect to  $M$ , enabling the construction of a surrogate objective to optimize  $M$  without directly querying the real environment  $M^*$ .
In the following, we select  $\ell_{M}$  as the negative log-likelihood loss for our full derivation, instantiating  $L(\rho_{M^{*}}^{\beta}, M)$  in Eq. (3) as:  $\mathbb{E}_{x, a \sim \rho_{M^{*}}^{\mu}}[\omega(x, a | \rho_{M^{*}}^{\beta}) \mathbb{E}_{M^{*}}(-\log M(x' | x, a))]$ , where  $\mathbb{E}_{M^{*}}[\cdot]$  denotes  $\mathbb{E}_{x' \sim M^{*}(x' | x, a)}[\cdot]$ . Ideally, for any given  $M$ , it is obvious that the optimal  $\beta$  is the one that makes  $\rho_{M^{*}}^{\beta}(x, a)$  assign all densities to the point with the largest negative log-likelihood. However, this maximization process is impractical, particularly in continuous spaces. To provide a tractable yet relaxed solution, we introduce an  $L_{2}$  regularizer to the original objective in Eq. (3).
$$
\min  _ {M \in \mathcal {M}} \max  _ {\beta \in \Pi} \bar {L} \left(\rho_ {M ^ {*}} ^ {\beta}, M\right) = \min  _ {M \in \mathcal {M}} \max  _ {\beta \in \Pi} L \left(\rho_ {M ^ {*}} ^ {\beta}, M\right) - \frac {\alpha}{2} \| \rho_ {M ^ {*}} ^ {\beta} (\cdot , \cdot) \| _ {2} ^ {2}, \tag {4}
$$
where  $\alpha$  denotes the regularization coefficient of  $\rho_{M^*}^\beta$  and  $\| \rho_{M^*}^\beta (\cdot ,\cdot)\| _2^2 = \int_{\mathcal{X},\mathcal{A}}(\rho_{M^*}^\beta (x,a))^2\mathrm{d}a\mathrm{d}x$ . Now we present the final results and the intuitions behind them, while providing a full derivation in Appx.A. Supposing we have  $\bar{\rho}_{M^*}^{\bar{\beta}^*}$  representing the approximated data distribution of the approximated best-response policy  $\bar{\beta}^*$  under model  $M_{\theta}$  parameterized by  $\theta$ , we can find the optimal  $\theta^{*}$  of  $\min_{\theta}\max_{\beta \in \Pi}\bar{L} (\rho_{M^*}^\beta ,M_\theta)$  (Eq. (4)) through iterative optimization of the objective  $\theta_{t + 1} = \min_{\theta}\bar{L} (\bar{\rho}_{M^*}^{\bar{\beta}^*},M_\theta)$ . To this end, we approximate  $\bar{\rho}_{M^*}^{\bar{\beta}^*}$  via the last-iteration model  $M_{\theta_t}$  and derive an upper bound objective for  $\min_{\theta}\bar{L} (\bar{\rho}_{M^*}^{\bar{\beta}^*},M_\theta)$ :
$$
\theta_ {t + 1} = \min  _ {\theta} \mathbb {E} _ {\rho_ {M ^ {*}} ^ {\mu}} \left[ \frac {- 1}{\alpha_ {0} (x , a)} \log M _ {\theta} \left(x ^ {\prime} | x, a\right) \underbrace {\left(\underbrace {f \left(\frac {\rho_ {M _ {\theta_ {t}}} ^ {\mu} \left(x , a , x ^ {\prime}\right)}{\rho_ {M ^ {*}} ^ {\mu} \left(x , a , x ^ {\prime}\right)}\right)} _ {\text {d i s c r e p a n c y}} - \underbrace {f \left(\frac {\rho_ {M _ {\theta_ {t}}} ^ {\mu} \left(x , a\right)}{\rho_ {M ^ {*}} ^ {\mu} \left(x , a\right)}\right)} _ {\text {d e n s i t y - r a t i o b a s e l i n e}} + \underbrace {H _ {M ^ {*}} \left(x , a\right)} _ {\text {s t o c h a s t i c i t y}}\right)} \right],
$$