SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\mathbb{E}_{\rho_{M^*}^\mu}[\cdot ]$  denotes  $\mathbb{E}_{x,a,x'\sim \rho_{M^*}^\mu}[\cdot ]$ ,  $f$  is a convex and lower semi-continuous (l.s.c.) function satisfying  $f^{\prime}(x)\leq 0,\forall x\in \mathcal{X}$ , which is also called  $f$  function in  $f$ -divergence [2],  $\alpha_0(x,a) = \alpha_{M_{\theta_t}}\rho_{M^*}^\mu (x,a)$ , and  $H_{M^{*}}(x,a)$  denotes the entropy of  $M^{*}(\cdot |x,a)$ .
Remark 4.4 (Intuition of  $W$ ). After derivation, we found that the optimal adversarial data distribution can be approximated by  $\bar{\rho}_{M^*}^{\bar{\beta}^*}(x,a) = \int_{\mathcal{X}}\rho_{M^*}^\mu (x,a,x')W(x,a,x')\mathrm{d}x'$  (see Appx. A), leading to the upper-bound objective Eq. (5), which is a WERM dynamically weighting by  $W$ . Intuitively,  $W$  assigns more learning propensities for data points with (1) larger discrepancy between  $\rho_{M_{\theta_t}}^\mu$  (generated by model) and  $\rho_{M^*}^\mu$  (real-data distribution), or (2) larger stochasticity of the real model  $M^*$ . The latter is contributed by the entropy  $H_{M^*}$ , while the former is contributed by the first two terms combined. In particular, through the definition of  $f$ -divergence, we known that the discrepancy of two distribution  $P$  and  $Q$  can be measured by  $\int_{\mathcal{X}}Q(x)f(P(x) / Q(x))\mathrm{d}x$ , thus the terms  $f(\rho_{M_{\theta_t}}^\mu (x,a,x') / \rho_{M^*}^\mu (x,a,x'))$  can be interpreted as the discrepancy measure unit between  $\rho_{M_{\theta_t}}^\mu (x,a,x')$  and  $\rho_{M^*}^\mu (x,a,x')$ , while  $f(\rho_{M_{\theta_t}}^\mu (x,a) / \rho_{M^*}^\mu (x,a))$  serves as a baseline on  $x$  and  $a$  measured by  $f$  to balance the discrepancy contributed by  $x$  and  $a$ , making  $M$  focus on errors on  $x'$ .
In summary, by adjusting the weights  $W$ , the learning process will iteratively exploit subtle errors of the current model in any data point, regardless of how many proportions it contributes in the original data distribution, to eliminate potential unidentifiability on counterfactual data caused by selection bias.
4.3 Tractable Solution
In Eq. (5), the terms  $f(\rho_{M_{\theta_t}}^\mu (x,a,x') / \rho_{M^*}^\mu (x,a,x')) - f(\rho_{M_{\theta_t}}^\mu (x,a) / \rho_{M^*}^\mu (x,a))$  are still intractable. Thanks to previous successful practices in GAN [19] and GAIL [25], we achieve the objective via a generator-discriminator-paradigm objective through similar derivation. We show the results as follows and leave the complete derivation in Appx. A.4. In particular, by introducing two discriminators  $D_{\varphi_0^*}(x,a,x')$  and  $D_{\varphi_0^*}(x,a)$ , we can optimize the surrogate objective Eq. (5) via:
$$
\theta_ {t + 1} = \max _ {\hat {\theta}} \left(\mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\hat {\mu}}} \left[ A _ {\varphi_ {0} ^ {*}, \varphi_ {1} ^ {*}} (x, a, x ^ {\prime}) \log M _ {\theta} (x ^ {\prime} | x, a) \right] + \mathbb {E} _ {\rho_ {M ^ {*}} ^ {\mu}} \left[ \left(H _ {M ^ {*}} (x, a) - A _ {\varphi_ {0} ^ {*}, \varphi_ {1} ^ {*}} (x, a, x ^ {\prime})\right) \log M _ {\theta} (x ^ {\prime} | x, a) \right]\right)
$$
$$
s. t. \quad \varphi_ {0} ^ {*} = \arg \max  _ {\varphi_ {0}} \left(\mathbb {E} _ {\rho_ {M ^ {*}} ^ {\mu}} \left[ \log D _ {\varphi_ {0}} (x, a, x ^ {\prime}) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\hat {\mu}}} \left[ \log \left(1 - D _ {\varphi_ {0}} (x, a, x ^ {\prime})\right) \right]\right)
$$
$$
\varphi_ {1} ^ {*} = \arg \max  _ {\varphi_ {1}} \left(\mathbb {E} _ {\rho_ {M ^ {*}} ^ {\mu}} [ \log D _ {\varphi_ {1}} (x, a) ] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\hat {\mu}}} [ \log (1 - D _ {\varphi_ {1}} (x, a)) ]\right), \tag {6}
$$
where  $\mathbb{E}_{\rho_M^\mu}[\cdot ]$  is a simplification of  $\mathbb{E}_{x,a,x^{\prime}\sim \rho_{M}^{\mu}}[\cdot ]$ ,  $A_{\varphi_0^*,\varphi_1^*}(x,a,x') = \log D_{\varphi_0^*}(x,a,x') - \log D_{\varphi_1^*}(x,a)$ , and  $\varphi_0$  and  $\varphi_{1}$  are the parameters of  $D_{\varphi_0}$  and  $D_{\varphi_1}$  respectively. We learn a policy  $\hat{\mu}\approx \mu$  via imitation learning based on the offline dataset  $\mathcal{D}_{\mathrm{real}}$  [40, 25]. Note that in the process, we ignore the term  $\alpha_0(x,a)$  for simplifying the objective. The discussion on the impacts of removing  $\alpha_0(x,a)$  is left in App. B.
The overall optimization pipeline is illustration in Fig. 3. In Eq. (6), the reweighting term  $W$  from Eq. (5) is split into two terms in the RHS of the equation: the first term is a GAIL-style objective [25], treating  $M_{\theta}$  as the policy generator,  $\hat{\mu}$  as the environment, and  $A$  as the advantage function, while the second term is WERM through  $H_{M^{*}} - A_{\varphi_{0}^{*},\varphi_{1}^{*}}$ . The first term resembles the previous adversarial model learning objectives [46, 47, 11]. These two terms have intuitive explanations: the first term assigns learning weights on data generated by the model  $M_{\theta_t}$ . If the predictions of the model appear realistic, mainly assessed by  $D_{\varphi_0^*}$ , the propensity weights would be increased, encouraging the model to generate more
Figure 3: Illustration of the GALILEO workflow.
such kind of data; Conversely, the second term assigns weights on real data generated by  $M^{*}$ . If the model's predictions seem unrealistic (mainly assessed by  $-D_{\varphi_0^*}$ ) or stochastic (evaluated by  $H_{M^*}$ ), the propensity weights will be increased, encouraging the model to pay more attention to these real data points when improving the likelihoods.
(a) GNFC
(b) TCGA
Figure 4: Illustration of the performance in GNFC and TCGA. The grey bar denotes the standard error ( $\times 0.3$  for brevity) of 3 random seeds.
Based on the above implementations, we propose Generative Adversarial offLIne counterfactualEnvironment mOdel learning (GALILEO) for environment model learning. We list a brief of GALILEO in Alg. 1 and the details in Appx. E.
Algorithm 1 Pseudocode for GALILEO
$\mathcal{D}_{\mathrm{real}}$ : offline dataset sampled from  $\rho_{M^*}^\mu$  where  $\mu$  is the behavior policy;  $N$ : total iterations;
Process:
1: Approximate a behavior policy  $\hat{\mu}$  via behavior cloning through offline dataset  $\mathcal{D}_{\mathrm{real}}$
2: Initialize an environment model  $M_{\theta_1}$
3: for  $t = 1:N$  do
4: Use  $\hat{\mu}$  to generate a dataset  $\mathcal{D}_{\mathrm{gen}}$  with the model  $M_{\theta_t}$
5: Update the discriminators  $D_{\varphi_0}$  and  $D_{\varphi_1}$  through the second and third equations in Eq. (6) where  $\rho_{M_{\theta_+}}^{\hat{\mu}}$  is estimated by  $\mathcal{D}_{\mathrm{gen}}$  and  $\rho_{M^{*}}^{\mu}$  is estimated by  $\mathcal{D}_{\mathrm{real}}$
6: Generative adversarial training for  $M_{\theta_t}$  by regarding  $A_{\varphi_0^*,\varphi_1^*}$  as the advantage function and computing the gradient to  $M_{\theta_t}$ , named  $g_{\mathrm{pg}}$ , with a standard policy gradient method like TRPO [44] or PPO [45] based on  $\mathcal{D}_{\mathrm{gen}}$ .
7: Regard  $H_{M^*} - A_{\varphi_0^*, \varphi_1^*}$  as the reweighting term for WERM and compute the gradient to  $M_{\theta_t}$  based on  $\mathcal{D}_{\mathrm{real}}$ . Record it as  $g_{\mathrm{sl}}$ .
8: Update the model  $\theta_{t + 1}\gets \theta_t + g_{\mathrm{pg}} + g_{\mathrm{sl}}$
9: end for
In this section, we first conduct experiments in two synthetic environments to quantify the performance of GALILEO on counterfactual queries  $^{3}$ . Then we deploy GALILEO in two complex environments: MuJoCo in Gym [53] and a real-world food-delivery platform to test the performance of GALILEO in difficult tasks. The results are in Sec. 5.2. Finally, to further verify the ability GALILEO, in Sec. 5.3, we apply models learned by GALILEO to several downstream tasks including off-policy evaluation, offline policy improvement, and online decision-making in production environment. The algorithms compared are: (1) SL: using standard empirical risk minimization for model learning; (2) IPW [50]: a standard implementation of WERM based IPS; (3) SCIGAN [7]: an adversarial algorithms for model learning used for causal effect estimation, which can be roughly regarded as a partial implementation of GALILEO (Refer to Appx. E.2). We give a detailed description in Appx. G.2.
5.1 Environment Settings
Synthetic Environments Previous experiments on counterfactual environment model learning are based on single-step semi-synthetic data simulation [7]. As GALILEO is compatible with single-step environment model learning, we first benchmark GALILEO in the same task named TCGA as previous studies do [7]. Based on the three synthetic response functions, we construct 9 tasks by choosing different parameters of selection bias on  $\mu$  which is constructed with beta distribution, and
3code https://github.com/xionghuichen/galileo.
design a coefficient  $c$  to control the selection bias. We name the tasks with the format of "t?_bias?" For example, t1.bias2 is the task with the first response functions and  $c = 2$ . The details of TCGA is in Appx. G.1.2. Besides, for sequential environment model learning under selection bias, we construct a new synthetic environment, general negative feedback control (GNFC), which can represent a classic type of task with policies having selection bias, where Fig. 1(a) is also an instance of GNFC. We construct 9 tasks on GNFC by adding behavior policies  $\mu$  with different scales of uniform noise  $U(-e, e)$  with probabilities  $p$ . Similarly, we name them with the format "e?_p?"
Continuous-control Environments We select 3 MuJoCo environments from D4RL [17] to construct our model learning tasks. We compare it with a standard transition model learning algorithm used in the previous offline model-based RL algorithms [59, 30], which is a standard supervised learning. We name the method OFF-SL. Besides, we also implement IPW and SCIGAN as the baselines. In D4RL benchmark, only the "medium" tasks is collected with a fixed policy, i.e., the behavior policy is with 1/3 performance to the expert policy), which is most matching to our proposed problem. So we train models in datasets HalfCheetah-medium, Walker2d-medium, and Hopper-medium.
A Real-world Large-scale Food-delivery Platform We finally deploy GALILEO in a real-world large-scale food-delivery platform. We focus on a Budget Allocation task to the Time period (BAT) in the platform (see Appx. G.1.3 for details). The goal of the BAT task is to handle the imbalance problem between the demanded orders from customers and the supply of delivery clerks in different time periods by allocating reasonable allowances to those time periods. The challenge of the environment model learning in BAT tasks is similar to the challenge in Fig. 1: the behavior policy is a human-expert policy, which tends to increase the budget of allowance in the time periods with a lower supply of delivery clerks, otherwise tends to decrease the budget (We gives a real-data instance in Appx. G.1.3).
5.2 Prediction Accuracy on Shifted Data Distributions
Test in Synthetic Environments For all of the tasks, we select mean-integrated-square error  $\mathrm{MISE} = \mathbb{E}\left[\int_{\mathcal{A}}(M^{*}(x^{\prime}|x,a) - M(x^{\prime}|x,a))^{2}\mathrm{d}a\right]$  as the metric, which is a metric to measure the accuracy in counterfactual queries by considering the prediction errors in the whole action space. The results are summarized in Fig. 4 and the detailed results can be found in Appx. H. The results show that the property of the behavior policy (i.e.,  $e$  and  $p$ ) dominates the generalization ability of the baseline algorithms. When  $e = 0.05$ , almost all of the baselines fail and give a completely opposite response curve, while GALILEO gives the correct response. (see Fig. 5). IPW still performs well when  $0.2 \leq e \leq 1.0$  but fails when  $e = 0.05, p <= 0.2$ . We also found that SCIGAN can reach a better performance than other baselines when  $e = 0.05, p <= 0.2$ , but the results in other tasks are unstable. GALILEO is the only algorithm that is robust to the selection bias and outputs correct response curves in all of the tasks. Based on the experiment, we also indicate that the commonly used overlap assumption is unreasonable to a certain extent especially in real-world applications since it is impractical to inject noises into the whole action space.
The problem of overlap assumption being violated, e.g.,  $e < 1$  in our setting, should be taken into consideration otherwise the algorithm will be hard to use in practice if it is sensitive to the noise range. On the other hand, we found the phenomenon in TCGA experiment is similar to the one in GNFC, which demonstrates the compatibility of GALILEO to single-step environments.
We also found that the results of IPW are unstable in TCGA experiment. It might be because the behavior policy is modeled with beta distribution while the propensity score  $\hat{\mu}$  is modeled with Gaussian distribution. Since IPW directly reweight loss with  $\frac{1}{\hat{\mu}}$ , the results are sensitive to the error of  $\hat{\mu}$ .
Figure 5: Illustration of the averaged response curves in task e0.05_p0.2.
Finally, we plot the averaged response curves which are constructed by equidistantly sampling action from the action space and averaging the feedback of the states in the dataset as the averaged response. One result is in Fig. 5 (all curves can be seen in Appx. H). For those tasks where base
lines fail in reconstructing response curves, GALILEO not only reaches a better MISE score but reconstructs almost exact responses, while the baselines might give completely opposite responses.
Test in MuJoCo Benchmarks We test the prediction error of the learned model in corresponding unseen "expert" and "medium-replay" datasets. Fig. 6 illustrates the results in halfcheetah. We can see that all algorithms perform well in the training datasets. OFF-SL can even reach a bit lower error. However, when we verify the models through "expert" and "medium-replay" datasets, which
(a) medium (train)
(b) medium-replay (test)
Figure 6: Illustration of learning curves of the halfcheetah tasks (full results are in Appx. H.5). The figures with titles ending in “(train)” means the dataset is used for training while the titles ending in “(test)” means the dataset is just used for testing. The X-axis records the steps of the environment model update, and the Y-axis is the prediction errors in the corresponding steps evaluated by the datasets. The solid curves are the mean reward and the shadow is the standard error of three seeds.
(c) expert (test)
are collected by other policies, the performance of GALILEO is more stable and better than all other algorithms. As the training continues, the baseline algorithms even gets worse and worse. The phenomenon are similar among three datasets, and we leave the full results in Appx. H.5.
(a) response curves in City A
Figure 7: Parts of the results in BAT tasks. Fig. 7(a) demonstrate the averaged response curves of the SL and GALILEO model in City A. It is expected to be monotonically increasing through our prior knowledge. In Fig. 7(b) show the AUCC curves, where the model with larger areas above the "random" line makes better predictions in randomized-controlled-trials data [61].
(b) AUUC curves
(c) A/B test in the City A
Test in a Real-world Dataset We first learn a model to predict the supply of delivery clerks (measured by fulfilled order amount) on given allowances. Although the SL model can efficiently fit the offline data, the tendency of the response curve is easily to be incorrect. As can be seen in Fig. 7(a), with a larger budget of allowance, the prediction of the supply is decreased in SL, which obviously goes against our prior knowledge. This is because, in the offline dataset, the corresponding supply will be smaller when the allowance is larger. It is conceivable that if we learn a policy through the model of SL, the optimal solution is canceling all of the allowances, which is obviously incorrect in practice. On the other hand, the tendency of GALILEO's response is correct. In Appx. H.7, we plot all the results in 6 cities. We further collect some randomized controlled trials data, and the Area Under the Uplift Curve (AUUC) [6] curve in Fig. 7(b) verify that GALILEO gives a reasonable sort order on the supply prediction while the standard SL technique fails to achieve this task.
5.3 Apply GALILEO to Downstream Tasks
Off-policy Evaluation (OPE) We first verify the ability of the models in MuJoCo environments by adopting them into off-policy evaluation tasks. We use 10 unseen policies constructed by DOPE benchmark [18] to conduct our experiments. We select three common-used metrics: value gap, regret@1, and rank correlation and averaged the results among three tasks in Tab. 1. The baselines and the corre
Table 1: Results of OPE on DOPE benchmark. We list the averaged performances on three tasks. The detailed results are in Appx. H.6.  $\pm$  is the standard deviation among the tasks. We bold the best scores for each metric.
Algorithm	Norm. value gap	Rank corr.	Regret@1
GALILEO	0.37 ± 0.24	0.44 ± 0.10	0.09 ± 0.02
Best DICE	0.48 ± 0.19	0.15 ± 0.04	0.42 ± 0.28
VPM	0.71 ± 0.04	0.29 ± 0.15	0.17 ± 0.11
FQE (L2)	0.54 ± 0.09	-0.19 ± 0.10	0.34 ± 0.03
IS	0.67 ± 0.01	-0.40 ± 0.15	0.36 ± 0.27
Doubly Rubost	0.57 ± 0.07	-0.14 ± 0.17	0.33 ± 0.06
Table 2: Results of policy performance directly optimized through SAC [20] using the learned dynamics models and deployed in MuJoCo environments. MAX-RETURN is the policy performance of SAC in the MuJoCo environments, and "avg. norm." is the averaged normalized return of the policies in the 9 tasks, where the returns are normalized to lie between 0 and 100, where a score of 0 corresponds to the worst policy, and 100 corresponds to MAX-RETURN.
Task	Hopper	Walker2d	HalfCheetah	avg. norm.
Horizon	H=10	H=20	H=40	H=10	H=20	H=40	H=10	H=20	H=40	/
GALILEO	13.0 ± 0.1	33.2 ± 0.1	53.5 ± 1.2	11.7 ± 0.2	29.9 ± 0.3	61.2 ± 3.4	0.7 ± 0.2	-1.1 ± 0.2	-14.2 ± 1.4	51.1
OFF-SL	4.8 ± 0.5	3.0 ± 0.2	4.6 ± 0.2	10.7 ± 0.2	20.1 ± 0.8	37.5 ± 6.7	0.4 ± 0.5	-1.1 ± 0.6	-13.2 ± 0.3	21.1
IPW	5.9 ± 0.7	4.1 ± 0.6	5.9 ± 0.2	4.7 ± 1.1	2.8 ± 3.9	14.5 ± 1.4	1.6 ± 0.2	0.5 ± 0.8	-11.3 ± 0.9	19.7
SCIGAN	12.7 ± 0.1	29.2 ± 0.6	46.2 ± 5.2	8.4 ± 0.5	9.1 ± 1.7	1.0 ± 5.8	1.2 ± 0.3	-0.3 ± 1.0	-11.4 ± 0.3	41.8
MAX-RETURN	13.2 ± 0.0	33.3 ± 0.2	71.0 ± 0.5	14.9 ± 1.3	60.7 ± 11.1	221.1 ± 8.9	2.6 ± 0.1	13.3 ± 1.1	49.1 ± 2.3	100.0
sponding results we used are the same as the one proposed by [18]. As seen in Tab. 1, compared with all the baselines, OPE by GALILEO always reach the better performance with a large margin (at least  $23\%$ ,  $193\%$  and  $47\%$  respectively), which verifies that GALILEO can eliminate the effect of selection bias and give correct evaluations on unseen policies.
Offline Policy Improvement We then verify the generalization ability of the models in MuJoCo environments by adopting them into offline model-based RL. To strictly verify the ability of the models, we abandon all tricks to suppress policy exploration and learning in risky regions as current offline model-based RL algorithms [59] do, and we just use the standard SAC algorithm [20] to fully exploit the models to search an optimal policy. Unfortunately, we found that the compounding error will still be inevitably large in the 1,000-step rollout, which is the standard horizon in MuJoCo tasks, leading all models to fail to derive a reasonable policy. To better verify the effects of models on policy improvement, we learn and evaluate the policies with three smaller horizons:  $H \in \{10,20,40\}$ . The results are listed in Tab. 2. We first averaged the normalized return (refer to "avg. norm.") under each task, and we can see that the policy obtained by GALILEO is significantly higher than other models (the improvements are  $24\%$  to  $161\%$ ). But in HalfCheetah, IPW works slightly better. However, compared with MAX-RETURN, it can be found that all methods fail to derive reasonable policies because their policies' performances are far away from the optimal policy. By further checking the trajectories, we found that all the learned policies just keep the cheetah standing in the same place or even going backward. This phenomenon is also similar to the results in MOPO [59]. In MOPO's experiment in the medium datasets, the truncated-rollout horizon used in Walker and Hopper for policy training is set to 5, while HalfCheetah has to be set to the minimal value: 1. These phenomena indicate that HalfCheetah may still have unknown problems, resulting in the generalization bottleneck of the models.
Online Decision-making in a Production Environment Finally, we search for the optimal policy via model-predict control (MPC) using cross-entropy planner [21] based on the learned model and deploy the policy in a real-world platform. The results of A/B test in City A is shown in Fig. 7(c). It can be seen that after the day of the A/B test, the treatment group (deploying our policy) significant improve the five-minute order-taken rate than the baseline policy (the same as the behavior policy). In summary, the policy improves the supply from 0.14 to 1.63 percentage points to the behavior policies in the 6 cities. The details of these results are in Appx. H.7.
6 Discussion and Future Work
In this work, we propose AWRM which handles the generalization challenges of the counterfactual environment model learning. By theoretical modeling, we give a tractable solution to handle AWRM and propose GALILEO. GALILEO is verified in synthetic environments, complex robot control tasks, and a real-world platform, and shows great generalization ability on counterfactual queries.
Giving correct answers to counterfactual queries is important for policy learning. We hope the work can inspire researchers to develop more powerful tools for counterfactual environment model learning. The current limitation lies in: There are several simplifications in the theoretical modeling process (further discussion is in Appx. B), which can be modeled more elaborately. Besides, experiments on MuJoCo indicate that these tasks are still challenging to give correct predictions on counterfactual data. These should also be further investigated in future work.
This work is supported by the National Key Research and Development Program of China (2020AAA0107200) and the National Science Foundation of China (61921006).