SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\arg \max  _ {\theta} \mathbb {E} _ {\rho_ {M ^ {*}} ^ {\kappa}} [ \log M _ {\theta} \left(x ^ {\prime} \mid x, a\right) ] \tag {21}
$$
Algorithm 2 AWRM with Oracle Counterfactual Datasets
$\Phi$  : policy space;  $N$  : total iterations
1: Generate counterfactual datasets  $\{\mathcal{D}_{\pi_{\phi}}\}$  for all adversarial policies  $\pi_{\phi},\phi \in \Phi$
2: Initialize an environment model  $M_{\theta}$
3: for i = 1:N do
4: Select  $\mathcal{D}_{\pi_{\phi}}$  with worst prediction errors through  $M_{\theta}$  from  $\{\mathcal{D}_{\pi_{\phi}}\}$
5: Optimize  $M_{\theta}$  with standard supervised learning based on  $\mathcal{D}_{\pi_{\phi}}$
6: end for
to replace the second term of the above objective when the output probability of  $D$  is far away from 0.5 ( $f'(1) = \log 0.5$ ).
In the offline model-learning setting, we only have a real-world dataset  $\mathcal{D}$  collected by the behavior policy  $\mu$ . We learn a policy  $\hat{\mu} \approx \mu$  via behavior cloning with  $\mathcal{D}$  [40, 25] and let  $\hat{\mu}$  be the policy  $\kappa$ . We regard  $\mathcal{D}$  as the empirical data distribution of  $\rho_{M^*}^\kappa$  and the trajectories collected by  $\hat{\mu}$  in the model  $M_{\theta_t}$  as the empirical data distribution of  $\rho_{M_{\theta_t}}^\kappa$ . But the assumption  $\forall x \in \mathcal{X}, \forall a \in \mathcal{A}, \mu(a|x) > 0$  might not be satisfied. In behavior cloning, we model  $\hat{\mu}$  with a Gaussian distribution and constrain the lower bound of the variance with a small value  $\epsilon_\mu > 0$  to keep the assumption holding. Besides, we add small Gaussian noises  $\mathbf{u} \sim \mathcal{N}(0, \epsilon_D)$  to the inputs of  $D_\varphi$  to handle the mismatch between  $\rho_{M^*}^\mu$  and  $\rho_{M^*}^{\hat{\mu}}$  due to  $\epsilon_\mu$ . In particular, for  $\varphi_0$  and  $\varphi_1$  learning, we have:
$$
\varphi_ {0} ^ {*} = \arg \max _ {\varphi_ {0}} \mathbb {E} _ {\rho_ {M ^ {*}, \mathbf {u}}} \left[ \log D _ {\varphi_ {0}} (x + u _ {x}, a + u _ {a}, x ^ {\prime} + u _ {x ^ {\prime}}) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}, \mathbf {u}}} \left[ \log (1 - D _ {\varphi_ {0}} (x + u _ {x}, a + u _ {a}, x ^ {\prime} + u _ {x ^ {\prime}})) \right]
$$
$$
\varphi_ {1} ^ {*} = \arg \max _ {\varphi_ {1}} \mathbb {E} _ {\rho_ {M ^ {*}, \mathbf {u}}} \left[ \log D _ {\varphi_ {1}} (x + u _ {x}, a + u _ {a}) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}, \mathbf {u}}} \left[ \log (1 - D _ {\varphi_ {1}} (x + u _ {x}, a + u _ {a})) \right],
$$
where  $\mathbb{E}_{\rho_{M_{\theta_t}}^{\kappa},\mathbf{u}}[\cdot ]$  is a simplification of  $\mathbb{E}_{x,a,x^{\prime}\sim \rho_{M_{\theta_t}}^{\kappa},\mathbf{u}\sim \mathcal{N}(0,\epsilon_D)}[\cdot ]$  and  $\mathbf{u} = [u_x,u_a,u_{x'}]$
On the other hand, we notice that the first term in Eq. (20) is similar to the objective of GAIL [25] by regarding  $M_{\theta}$  as the policy to learn and  $\kappa$  as the environment to generate data. For better capability in sequential environment model learning, here we introduce some practical tricks inspired by GAIL for model learning [47, 46]: we introduce an MDP for  $\kappa$  and  $M_{\theta}$ , where the reward is defined by the discriminator  $D$ , i.e.,  $r(x,a,x') = \log D(x,a,x')$ .  $M_{\theta}$  is learned to maximize the cumulative rewards. With advanced policy gradient methods [44, 45], the objective is converted to  $\max_{\theta}\left[A_{\varphi_0^*,\varphi_1^*}(x,a,x')\log M_\theta (x,a,x')\right]$ , where  $A = Q_{M_{\theta_t}}^\kappa -V_{M_{\theta_t}}^\kappa$ ,  $Q_{M_{\bar{\theta}}}^{\kappa}(x,a,x') = \mathbb{E}\left[\sum_{t = 0}^{\infty}\gamma^{t}r(x_{t},a_{t},x_{t + 1})\mid (x_{t},a_{t},x_{t + 1}) = (x,a,x'),\kappa ,M_{\theta_t}\right]$ , and  $V_{M_{\bar{\theta}}}^{\kappa}(x,a) = \mathbb{E}_{M_{\bar{\theta}}}\left[Q_{M_{\bar{\theta}}}^{\kappa}(x,a,x')\right]$ .  $A$  in Eq. (20) can also be constructed similarly. Although it looks unnecessary in theory since the one-step optimal model  $M_{\theta}$  is the global optimal model in this setting, the technique is helpful in practice as it makes  $A$  more sensitive to the compounding effect of one-step prediction errors: we would consider the cumulative effects of prediction errors induced by multi-step transitions in environments. In particular, to consider the cumulative effects of prediction errors induced by multi-step of transitions in environments, we overwrite function  $A_{\varphi_0^*,\varphi_1^*}$  as  $A_{\varphi_0^*,\varphi_1^*} = Q_{M_{\theta_t}}^\kappa -V_{M_{\theta_t}}^\kappa$  where  $Q_{M_{\theta_t}}^\kappa (x,a,x') = \mathbb{E}\left[\sum_t^\infty \gamma^t\log D_{\varphi_0^*}(x_t,a_t,x_{t + 1})|(x_t,a_t,x_{t + 1}) = (x,a,x'),\kappa ,M_{\theta_t}\right]$  and  $V_{M_{\theta_t}}^\kappa (x,a) = \mathbb{E}\left[\sum_t^\infty \gamma^t\log D_{\varphi_1^*}(x_t,a_t)|(x_t,a_t) = (x,a),\kappa ,M_{\theta_t}\right]$ . To give an algorithm for single-step environment model learning, we can just set  $\gamma$  in  $Q$  and  $V$  to 0.
Algorithm 3 GALILEO pseudocode
$\mathcal{D}_{\mathrm{real}}$ : offline dataset sampled from  $\rho_{M^*}^\mu$  where  $\mu$  is the behavior policy;
$N$ : total iterations;
1: Approximate a behavior policy  $\hat{\mu}$  via behavior cloning
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}$  via Eq. (25) and Eq. (26) respectively, where  $\rho_{M_{\theta_t}}^{\hat{\mu}}$  is estimated by  $\mathcal{D}_{\mathrm{gen}}$  and  $\rho_{M^{*}}^{\mu}$  is estimated by  $\mathcal{D}_{\mathrm{real}}$
6: Update  $Q$  and  $V$  via Eq. (23) and Eq. (24) through  $\mathcal{D}_{\mathrm{gen}}$ ,  $D_{\varphi_0}$ , and  $D_{\varphi_1}$
7: Update the model  $M_{\theta_t}$  via the first term of Eq. (22), which is implemented with a standard policy gradient method like TRPO [44] or PPO [45]. Record the policy gradient  $g_{\mathrm{pg}}$
8: if  $p_0 < \mathbb{E}_{\mathcal{D}_{\mathrm{gen}}}\left[D_{\varphi_0}(x_t, a_t, x_{t+1})\right] < p_1$  then
9: Compute the gradient of  $M_{\theta_t}$  via the second term of Eq. (22) and record it as  $g_{\mathrm{sl}}$
0: else
1: Compute the gradient of  $M_{\theta_t}$  via Eq. (21) and record it as  $g_{\mathrm{sl}}$
2: end if
3: Rescale  $g_{\mathrm{sl}}$  via Eq. (27)
4: Update the model  $M_{\theta_t}$  via the gradient  $g_{\mathrm{sl}}$  and obtain  $M_{\theta_{t + 1}}$
5: end for
Figure 9: Illustration of the workflow of the GALILEO algorithm.
By adopting the above implementation techniques, we convert the objective into the following formulation
$$
\theta_ {t + 1} = \arg \max  _ {\theta} \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}} \left[ A _ {\varphi_ {0} ^ {*}, \varphi_ {1} ^ {*}} \left(x, a, x ^ {\prime}\right) \log M _ {\theta} \left(x ^ {\prime} \mid x, a\right) \right] + \mathbb {E} _ {\rho_ {M ^ {*}} ^ {\kappa}} \left[ \left(H _ {M ^ {*}} (x, a) - A _ {\varphi_ {0} ^ {*}, \varphi_ {1} ^ {*}} \left(x, a, x ^ {\prime}\right)\right) \log M _ {\theta} \left(x ^ {\prime} \mid x, a\right) \right] \tag {22}
$$
s.t.  $Q_{M_{\theta_t}}^\kappa (x,a,x') = \mathbb{E}\left[\sum_t^\infty \gamma^t\log D_{\varphi_0^*}(x_t,a_t,x_{t + 1})|(x_t,a_t,x_{t + 1}) = (x,a,x'),\kappa ,M_{\theta_t}\right]$  (23)
$$
V _ {M _ {\theta_ {t}}} ^ {\kappa} (x, a) = \mathbb {E} \left[ \sum_ {t} ^ {\infty} \gamma^ {t} \log D _ {\varphi_ {1} ^ {*}} \left(x _ {t}, a _ {t}\right) \mid \left(x _ {t}, a _ {t}\right) = (x, a), \kappa , M _ {\theta_ {t}} \right] \tag {24}
$$
$$
\varphi_ {0} ^ {*} = \arg \max  _ {\varphi_ {0}} \mathbb {E} _ {\rho_ {M ^ {*}}, \mathbf {u}} \left[ \log D _ {\varphi_ {0}} \left(x + u _ {x}, a + u _ {a}, x ^ {\prime} + u _ {x ^ {\prime}}\right) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}, \mathbf {u}} \left[ \log \left(1 - D _ {\varphi_ {0}} \left(x + u _ {x}, a + u _ {a}, x ^ {\prime} + u _ {x ^ {\prime}}\right)\right) \right] \tag {25}
$$
$$
\varphi_ {1} ^ {*} = \arg \max  _ {\varphi_ {1}} \mathbb {E} _ {\rho_ {M ^ {*}}, \mathbf {u}} \left[ \log D _ {\varphi_ {1}} \left(x + u _ {x}, a + u _ {a}\right) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}, \mathbf {u}} \left[ \log \left(1 - D _ {\varphi_ {1}} \left(x + u _ {x}, a + u _ {a}\right)\right) \right], \tag {26}
$$
where  $A_{\varphi_0^*, \varphi_1^*}(x, a, x') = Q_{M_\theta}^\kappa(x, a, x') - V_{M_\theta}^\kappa(x, a)$ . In practice, GALILEO optimizes the first term of Eq. (22) with conservative policy gradient algorithms (e.g., PPO [45] or TRPO [44]) to avoid unreliable gradients for model improvements. Eq. (25) and Eq. (26) are optimized with supervised learning. The second term of Eq. (22) is optimized with supervised learning with a re-weighting term  $-A_{\varphi_0^*, \varphi_1^*} + H_{M^*}$ . Since  $H_{M^*}$  is unknown, we use  $H_{M_\theta}$  to estimate it. When the mean output probability of a batch of data is larger than 0.6 or small than 0.4, we replace the second term of Eq. (22) with a standard supervised learning in Eq. (21). Besides, unreliable gradients also exist in the process of optimizing the second term of Eq. (22). In our implementation, we use the scale of policy gradients to constrain the gradients of the second term of Eq. (22). In particular, we first compute the  $l_2$ -norm of the gradient of the first term of Eq. (22) via conservative policy gradient algorithms, named  $||g_{\mathrm{pg}}||_2$ . Then we compute the  $l_2$ -norm of the gradient of the second term of Eq. (22), name  $||g_{\mathrm{sl}}||_2$ . Finally, we rescale the gradients of the second term  $g_{\mathrm{sl}}$  by
$$
g _ {\mathrm {s l}} \leftarrow g _ {\mathrm {s l}} \frac {\left\| g _ {\mathrm {p g}} \right\| _ {2}}{\max  \left\{\left\| g _ {\mathrm {p g}} \right\| _ {2} , \left\| g _ {\mathrm {s l}} \right\| _ {2} \right\}}. \tag {27}
$$
For each iteration, Eq. (22), Eq. (25), and Eq. (26) are trained with certain steps (See Tab. 6) following the same framework as GAIL. Based on the above techniques, we summarize the pseudocode of GALILEO in Alg. 3, where  $p_0$  and  $p_1$  are set to 0.4 and 0.6 in all of our experiments. The overall architecture is shown in Fig. 9.
E.2 Connection with Previous Adversarial Algorithms
Standard GAN [19] can be regarded as a partial implementation including the first term of Eq. (22) and Eq. (25) by degrading them into the single-step scenario. In the context of GALILEO, the objective of GAN is
$$
\begin{array}{l} \theta_ {t + 1} = \arg \max  _ {\theta} \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}} \left[ A _ {\varphi^ {*}} \left(x, a, x ^ {\prime}\right) \log M _ {\theta} \left(x ^ {\prime} \mid x, a\right) \right] \\ s. t. \varphi * = \arg \max _ {\varphi} \mathbb {E} _ {\rho_ {M ^ {*}} ^ {\kappa}} \left[ \log D _ {\varphi} (x, a, x ^ {\prime}) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}} \left[ \log (1 - D _ {\varphi} (x, a, x ^ {\prime})) \right], \\ \end{array}
$$
where  $A_{\varphi^*}(x,a,x') = \log D_{\varphi^*}(x,a,x')$ . In the single-step scenario,  $\rho_{M_{\theta_t}}^\kappa(x,a,x') = \rho_0(x)\kappa(a|x)M_{\theta_t}(x'|a,x)$ . The term  $\mathbb{E}_{\rho_{M_{\theta_t}}^\kappa}[A_{\varphi^*}(x,a,x')\log M_\theta(x'|x,a)]$  can convert to  $\mathbb{E}_{\rho_{M_\theta}^\kappa}[\log D_{\varphi^*}(x,a,x')]$  by replacing the gradient of  $M_{\theta_t}(x'|x,a)\nabla_\theta\log M_\theta(x'|x,a)$  with  $\nabla_\theta M_\theta(x'|x,a)$  [51]. Previous algorithms like GANITE [58] and SCIGAN [7] can be regarded as variants of the above training framework.
The first term of Eq. (22) and Eq. (25) are similar to the objective of GAIL by regarding  $M_{\theta}$  as the "policy" to imitate and  $\hat{\mu}$  as the "environment" to collect data. In the context of GALILEO, the objective of GAIL is:
$$
\begin{array}{l} \theta_ {t + 1} = \arg \max  _ {\theta} \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}} \left[ A _ {\varphi^ {*}} \left(x, a, x ^ {\prime}\right) \log M _ {\theta} \left(x ^ {\prime} \mid x, a\right) \right] \\ s. t. \quad Q _ {M _ {\theta_ {t}}} ^ {\kappa} (x, a, x ^ {\prime}) = \mathbb {E} \left[ \sum_ {t} ^ {\infty} \gamma^ {t} \log D _ {\varphi^ {*}} (x _ {t}, a _ {t}, x _ {t + 1}) | (x _ {t}, a _ {t}, x _ {t + 1}) = (x, a, x ^ {\prime}), \kappa , M _ {\theta_ {t}} \right] \\ \varphi^ {*} = \underset {\varphi} {\arg \max} \mathbb {E} _ {\rho_ {M ^ {*}} ^ {\kappa}} \left[ \log D _ {\varphi} (x, a, x ^ {\prime}) \right] + \mathbb {E} _ {\rho_ {M _ {\theta_ {t}}} ^ {\kappa}} \left[ \log (1 - D _ {\varphi} (x, a, x ^ {\prime})) \right], \\ \end{array}
$$
where  $A_{\varphi^*}(x,a,x') = Q_{M_\theta}^\kappa (x,a,x') - V_{M_\theta}^\kappa (x,a)$  and  $V_{M_{\theta_t}}^\kappa (x,a) = \mathbb{E}_{M_{\theta_t}(x,a)}[Q^\kappa (x,a,x')]$ .
F Additional Related Work
Our primitive objective is inspired by weighted empirical risk minimization (WERM) based on inverse propensity score (IPS). WERM is originally proposed to solve the generalization problem of domain adaptation in machine learning literature. For instance, we would like to train a predictor  $M(y|x)$  in a domain with distribution  $P_{\mathrm{train}}(x)$  to minimize the prediction risks in the domain with distribution  $P_{\mathrm{test}}(x)$ , where  $P_{\mathrm{test}} \neq P_{\mathrm{test}}$ . To solve the problem, we can train a weighted objective with  $\max_M \mathbb{E}_{x \sim P_{\mathrm{train}}}[\frac{P_{\mathrm{test}}(x)}{P_{\mathrm{train}}(x)} \log M(y|x)]$ , which is called weighted empirical risk minimization methods [5, 4, 15, 9, 42]. These results have been extended and applied to causal inference, where the predictor is required to be generalized from the data distribution in observational studies (source domain) to the data distribution in randomized controlled trials (target domain) [48, 3, 22, 29, 28].
In this case, the input features include a state  $x$  (a.k.a. covariates) and an action  $a$  (a.k.a. treatment variable) which is sampled from a policy. We often assume the distribution of  $x$ ,  $P(x)$  is consistent between the source domain and the test domain, then we have  $\frac{P_{\mathrm{test}}(x)}{P_{\mathrm{train}}(x)} = \frac{P(x)\beta(a|x)}{P(x)\mu(a|x)} = \frac{\beta(a|x)}{\mu(a|x)}$ , where  $\mu$  and  $\beta$  are the policies in source and target domains respectively. In [48, 3, 22], the policy in randomized controlled trials is modeled as a uniform policy, then  $\frac{P_{\mathrm{test}}(x)}{P_{\mathrm{train}}(x)} = \frac{P(x)\beta(a|x)}{P(x)\mu(a|x)} = \frac{\beta(a|x)}{\mu(a|x)} \propto \frac{1}{\mu(a|x)} \cdot \frac{1}{\mu(a|x)}$  is also known as inverse propensity score (IPS). In [28], it assumes that the policy in the target domain is predefined as  $\beta(a|x)$  before environment model learning, then it uses  $\frac{\beta}{\mu}$  as the IPS. The differences between AWRM and previous works are fallen in two aspects: (1) We consider the distribution-shift problem in the sequential decision-making scenario. In this scenario, we not only consider the action distribution mismatching between the behavior policy  $\mu$  and the policy to evaluation  $\beta$ , but also the follow-up effects of policies to the state distribution; (2) For faithful offline policy optimization, we require the environment model to have generalization ability in numerous different policies. The objective of AWRM is proposed to guarantee the generalization ability of  $M$  in numerous different policies instead of a specific policy.
On a different thread, there are also studies that bring counterfactual inference techniques of causal inference into model-based RL [8, 39, 49]. These works consider that the transition function is relevant to some hidden noise variables and use Pearl-style structural causal models (SCMs), which is a directed acyclic graphs to define the causality of nodes in an environment, to handle the problem. SCMs can help RL in different ways: [8] approximate the posterior of the noise variables based on the observation of data, and environment models are learned based on the inferred noises. The generalization ability is improved if we can infer the correct value of the noise variables. [39] discover several local causal structural models of a global environment model, then data augmentation strategies by leveraging these local structures to generate counterfactual experiences. [49] proposes a representation learning technique for causal factors, which is an instance of the hidden noise variables, in partially observable Markov decision processes (POMDPs). With the learned representation of causal factors, the performance of policy learning and transfer in downstream tasks will be improved.
Instead of considering the hidden noise variables in the environments, our study considers the environment model learning problem in the fully observed setting and focuses on unbiased causal effect estimation in the offline dataset under behavior policies collected with selection bias.
In offline model-based RL, the problem is called distribution shift [59, 31, 14] which has received great attentions. However, previous algorithms do not handle the model learning challenge directly but propose techniques to suppress policy sampling and learning in risky regions [59, 30]. Although these algorithms have made great progress in offline policy optimization in many tasks, so far, how to learn a better environment model in this scenario has rarely been discussed.
We are not the first article to use the concept of density ratio for weighting. In off-policy estimation, [32, 35, 60] use density ratio to evaluate the value of a given target policy  $\beta$ . These methods attempt to solve an accurate approximation of  $\omega(s, a|\rho^{\beta})$ . The objective of our work, AWRM, is for the model to provide faithful feedback for different policies, formalized as minimizing the model error of the density function for any  $\beta$  in the policy space  $\Pi$ . The core of this problem is how to obtain an approximation of the density function of the best-response  $\beta^{*}$  corresponding to the current model, and then approximate the corresponding  $\omega(s, a|\rho^{\beta^{*}})$ . It should be emphasized that since  $\beta^{*}$  is unknown in advance and will change as the induction bias of  $M$  is changed, the solutions proposed in [32, 35, 60] cannot be applied to AWRM; Recently, [24, 57] use density ratio weighting to learn a model. The purpose of weighting is to make the model adapt to the current policy learning and adjust the model learning according to the current policy, which is the same as the WERM objective proposed in Def. 4.1. [54] also utilizes density ratio weighting to learn a model. Instead of estimating them based on the offline dataset and target policy as previous works do, they propose to design an adversarial model learning objective by constructing two function classes  $\mathcal{V}$  and  $\mathcal{W}$ , satisfying the target policy's value  $V_{\beta}$  and density ratio  $\omega_{\beta}$  are covered, i.e.,  $V_{\beta} \in \mathcal{V}$  and  $\omega_{\beta} \in \mathcal{W}$ . Different from previous articles, our approach uses adversarial weighting to learn a universal model that provides good feedback for any target policy  $\beta \in \Pi$ , i.e., AWRM, instead of learning a model suitable to a specific target policy.
G Experiment Details
G.1 Settings
G.1.1 General Negative Feedback Control (GNFC)
The design of GNFC is inspired by a classic type of scenario that behavior policies  $\mu$  have selection bias and easily lead to counterfactual risks: For some internet platforms, we would like to allocate budgets to a set of targets (e.g., customers or cities) to increase the engagement of the targets in the platforms. Our task is to train a model to predict targets' feedback on engagement given targets' features and allocated budgets.
In these tasks, for better benefits, the online working policy (i.e., the behavior policy) will tend to cut down the budgets if targets have better engagement, otherwise, the budgets might be increased. The risk of counterfactual environment model learning in the task is that: the object with better historical engagement will be sent to smaller budgets because of the selection bias of the behavior policies, then the model might exploit this correlation for learning and get a conclusion that: increasing budgets will reduce the targets' engagement, which violates the real causality. We construct an environment and a behavior policy to mimic the above process. In particular, the behavior policy  $\mu_{GNFC}$  is
$$
\mu_ {G N F C} (x) = \frac {(6 2 . 5 - \mathrm {m e a n} (x))}{1 5} + \epsilon ,
$$
where  $\epsilon$  is a sample noise, which will be discussed later. The environment includes two parts:
(1) response function  $M_{1}(y|x,a)$ :
$$
M _ {1} (y | x, a) = \mathcal {N} (\operatorname {m e a n} (x) + a, 2)
$$
(2) mapping function  $M_2(x'|x,y)$ :
$$
M _ {2} \left(x ^ {\prime} | x, a, y\right) = y - \operatorname {m e a n} (x) + x
$$
The transition function  $M^{*}$  is a composite of  $M^{*}(x^{\prime}|x,a) = M_{2}(x^{\prime}|x,a,M_{1}(y|x,a))$ . The behavior policies have selection bias: the actions taken are negatively correlated with the states, as illustrated in Fig. 10(a) and Fig. 10(b). We control the difficulty of distinguishing the correct causality of  $x$ ,  $a$ , and  $y$  by designing different strategies of noise sampling on  $\epsilon$ . In principle, with a larger number or more pronounced disturbances, there are more samples violating the correlation between  $x$  and  $a$ , then more samples can be used to find the correct causality. Therefore, we can control the difficulty of counterfactual environment model learning by controlling the strength of disturbance. In particular, we sample  $\epsilon$  from a uniform distribution  $U(-e,e)$  with probability  $p$ . That is,  $\epsilon = 0$  with probability  $1 - p$  and  $\epsilon \sim U(-e,e)$  with probability  $p$ . Then with larger  $p$ , there are more samples in the dataset violating the negative correlation (i.e.,  $\mu_{GNFC}$ ), and with larger  $e$ , the difference of the feedback will be more obvious. By selecting different  $e$  and  $p$ , we can construct different tasks to verify the effectiveness and ability of the counterfactual environment model learning algorithm.
(a)  $y$  in collected trajectories.
(b)  $a$  in collected trajectories.
Figure 10: Illustration of information about the collected dataset in GNFC. Each color of the line denotes one of the collected trajectories. The X-axis denotes the timestep of a trajectory.
G.1.2 The Cancer Genomic Atlas (TCGA)
The Cancer Genomic Atlas (TCGA) is a project that has profiled and analyzed large numbers of human tumors to discover molecular aberrations at the DNA, RNA, protein, and epigenetic levels. The resulting rich data provide a significant opportunity to accelerate our understanding of the molecular basis of cancer. We obtain features,  $\mathbf{x}$ , from the TCGA dataset and consider three continuous treatments as done in SCIGAN [7]. Each treatment,  $a$ , is associated with a set of parameters,  $\mathbf{v}_1$ ,  $\mathbf{v}_2$ ,  $\mathbf{v}_3$ , that are sampled randomly by sampling a vector from a standard normal distribution and scaling it with its norm. We assign interventions by sampling a treatment,  $a$ , from a beta distribution,  $a \mid \mathbf{x} \sim \mathrm{Beta}(\alpha, \beta)$ .  $\alpha \geq 1$  controls the sampling bias and  $\beta = \frac{\alpha - 1}{a^*} + 2 - \alpha$ , where  $a^*$  is the optimal treatment. This setting of  $\beta$  ensures that the mode of Beta  $(\alpha, \beta)$  is  $a^*$ .
The calculation of treatment response and optimal treatment are shown in Table 4.
Table 4: Treatment response used to generate semi-synthetic outcomes for patient features x. In the experiments, we set  $C = 10$ .
Treatment	Treatment Response	Optimal treatment
1	f1(x, a1) = C((v1)T x + 12 (v2)T x a1 - 12 (v3)T x a12)	a1* = (v12)T x / 2(v31)T x
2	f2(x, a2) = C((v12)T x + sin(π(v22T x / v32T x) a2))	a2* = (v32)T x / 2(v22)T x
3	f3(x, a3) = C((v13)T x + 12a3(a3-b)2, where b = 0.75(v23)T x / (v33)T x)	3/b if b ≥ 0.75, 1 if b < 0.75