Title: Adversarial Counterfactual Environment Model Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: environment model learning

offline reinforcement learning

off-policy evaluation

individual treatment effects estimation

causal inference

adversarial learning
Abstract: 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  \textit{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.

Adversarial Counterfactual Environment Model Learning

Abstract

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 longterm 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.

1 Introduction

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,

x

𝑦

𝑠

!

=(1, 10)

𝑠

"

=(2, 45)

𝑎

#

= 0.4

𝑠

$

=(4, 61)

𝑠

%

=(6, 62.5)

𝑎

&

= 0.02

𝑎

!

= 3.5

𝑎

"

= 1.17

𝑎

$

= 0.13

𝑎

%

= 0

𝑠

#

=(3, 57)

𝑠

&

=(5, 62.3)

selection bias

trajectory

target line

state

action

0

2

4

action

0

10

20

30

40

50

60

70

80

ball's y-axis location

real

GALILEO

SL

1.0

0.5

0.0

0.5

1.0

variation of action value

8

6

4

2

0

2

4

6

8

variation of next y prediction

real

GALILEO

SL

(a) a selection bias example

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-

0.0

12.5

25.0

37.5

50.0

62.5

75.0

87.5 100.0 112.5 125.0

adversarial policy's parameter

of the counterfactual dataset

0

25

50

75

100

125

150

175

200

mean square error (MSE) in the dataset

GALILEO

SL

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, AWRMemploys 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 G enerative A dversarial off LI ne counterfactua L E nvironment m O del 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.

2 Preliminaries

We first introduce weighted empirical risk minimization (WERM) through inverse propensity scoring (IPS), which is commenly 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 µ ( a | x ) which has selection bias, a classical solution to handle the above problem is WERM through IPS ω [48, 3, 29]:

Definition 2.1. The learning objective of WERM through IPS is formulated as

where ω ( x, a ) := β ( a | x ) µ ( a | x ) , β and µ denote the policies in testing and training domains, and the joint probability p µ M ∗ ( x, a, y ) := ρ 0 ( x ) µ ( a | x ) M ∗ ( y | x, a ) in which ρ 0 ( x ) is the distribution of state. M is the model space. ℓ is a loss function.

The ω 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 ˆ ω to approximate ω to learn the model M , current environment model learning algorithms employing reweighting are fallen into the framework. For standard ITEs estimation, ω = 1 ˆ µ (i.e., β is a uniform policy) for balancing treatments, where ˆ µ is an approximated behavior policy µ . 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 ∼ ρ 0 ( x ) . Then at each time-step t ∈ { 0 , 1 , 2 , ... } , the agent takes an action a t ∈ A through a policy π ( a t | x t ) ∈ Π based on the state x t ∈ X , then the agent receives a reward r t from

a reward function r ( x t , a t ) ∈ 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. Π , X , and A denote the policy, state, and action spaces.

3 Related Work

Wegive 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 β . More specifically, the question is: If β 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 ′ | x, a ) > 0 , ∀ x ∈ X , ∀ a ∈ A , ∀ x ′ ∈ X and µ satisfies µ ( a | x ) > 0 , ∀ a ∈ A , ∀ x ∈ X , the learning objective of generalized WERM through IPS is:

where ω ( x, a, x ′ ) = ρ β M ∗ ( x,a,x ′ ) ρ µ M ∗ ( x,a,x ′ ) , ρ µ M ∗ and ρ β M ∗ the training and testing data distributions collected by policy µ and β respectively. We define ℓ M ( x, a, x ′ ) := ℓ ( M ( x ′ | x, a ) , x ′ ) for brevity.

In an MDP, for any given policy π , we have ρ π M ∗ ( x, a, x ′ ) = ρ π M ∗ ( x ) π ( a | x ) M ∗ ( x ′ | x, a ) where ρ π M ∗ ( x ) denotes the occupancy measure of x for policy π . This measure can be defined as (1 -γ ) E x 0 ∼ ρ 0 [ ∑ ∞ t =0 γ t Pr( x t = x | x 0 , M ∗ )] [51, 25] where Pr π [ x t = x | x 0 , M ∗ ] is the discounted state visitation probability that the policy π visits x at time-step t by executing in the environment M ∗ and starting at the state x 0 . Here γ ∈ [0 , 1] is the discount factor. We also define ρ π M ∗ ( x, a ) := ρ π M ∗ ( x ) π ( a | x ) for simplicity.

With this definition, ω can be rewritten: ω = ρ β M ∗ ( x ) β ( a | x ) M ∗ ( x ′ | x,a ) ρ µ M ∗ ( x ) µ ( a | x ) M ∗ ( x ′ | x,a ) = ρ β M ∗ ( x,a ) ρ µ M ∗ ( x,a ) . In single-step environments, for any policy π , ρ π M ∗ ( x ) = ρ 0 ( x ) . Consequently, we obtain ω = ρ 0 ( x ) β ( a | x ) ρ 0 ( x ) µ ( a | x ) = β ( a | x ) µ ( a | x ) , and the objective degrade to Eq. (1). Therefore, Def. 2.1 is a special case of this generalized form.

Remark 4.2. ω 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 ω [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 β 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 A dversarial W eighted empirical R isk M inimization (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

where ω ( x, a | ρ β M ∗ ) = ρ β M ∗ ( x,a ) ρ µ M ∗ ( x,a ) , and the re-weighting term ω ( x, a | ρ β M ∗ ) is conditioned on the distribution ρ β M ∗ of the adversarial policy β . In the following, we will ignore ρ β M ∗ and use ω ( x, a ) for brevity.

In a nutshell, Eq. (3) minimizes the maximum model loss under all counterfactual data distributions ρ β M ∗ , β ∈ Π to guarantee the generalization ability for counterfactual queried by policies in Π .

4.2 Surrogate AWRM through Optimal Adversarial Data Distribution Approximation

The main challenge of solving AWRM is constructing the data distribution ρ β ∗ M ∗ of the best-response policy β ∗ in M ∗ since obtaining additional data from M ∗ can be expensive in real-world applications. In this paper, instead of deriving the optimal β ∗ , our solution is to offline estimate the optimal adversarial distribution ρ β ∗ M ∗ ( 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 ℓ M as the negative log-likelihood loss for our full derivation, instantiating L ( ρ β M ∗ , M ) in Eq. (3) as: E x,a ∼ ρ µ M ∗ [ ω ( x, a | ρ β M ∗ ) E M ∗ ( -log M ( x ′ | x, a ))] , where E M ∗ [ · ] denotes E x ′ ∼ M ∗ ( x ′ | x,a ) [ · ] . Ideally, for any given M , it is obvious that the optimal β is the one that makes ρ β M ∗ ( 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).

where α denotes the regularization coefficient of ρ β M ∗ and ∥ ρ β M ∗ ( · , · ) ∥ 2 2 = ∫ X , A ( ρ β M ∗ ( x, a )) 2 d a d x .

Now we present the final results and the intuitions behind them, while providing a full derivation in Appx.A. Suppossing we have ¯ ρ ¯ β ∗ M ∗ representing the approximated data distribution of the approximated best-response policy ¯ β ∗ under model M θ parameterized by θ , we can find the optimal θ ∗ of min θ max β ∈ Π ¯ L ( ρ β M ∗ , M θ ) (Eq. (4)) through iterative optimization of the objective θ t +1 = min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) . To this end, we approximate ¯ ρ ¯ β ∗ M ∗ via the last-iteration model M θ t and derive an upper bound objective for min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) :

where E ρ µ M ∗ [ · ] denotes E x,a,x ′ ∼ ρ µ M ∗ [ · ] , f is a convex and lower semi-continuous (l.s.c.) function satisfying f ′ ( x ) ≤ 0 , ∀ x ∈ X , which is also called f function in f -divergence [2], α 0 ( x, a ) = α M θ t ρ µ M ∗ ( x, a ) , and H M ∗ ( x, a ) denotes the entropy of M ∗ ( ·| x, a ) .

Remark 4.4 (Intuition of W ) . After derivation, we found that the optimal adversarial data distribution can be approximated by ¯ ρ ¯ β ∗ M ∗ ( x, a ) = ∫ X ρ µ M ∗ ( x, a, x ′ ) W ( x, a, x ′ ) 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 ρ µ M θ t (generated by model) and ρ µ M ∗ (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 ∫ X Q ( x ) f ( P ( x ) /Q ( x )) d x , thus the terms f ( ρ µ M θ t ( x, a, x ′ ) /ρ µ M ∗ ( x, a, x ′ )) can be interpreted as the discrepancy measure unit between ρ µ M θ t ( x, a, x ′ ) and ρ µ M ∗ ( x, a, x ′ ) , while f ( ρ µ M θ t ( x, a ) /ρ µ M ∗ ( 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 ( ρ µ M θ t ( x, a, x ′ ) /ρ µ M ∗ ( x, a, x ′ )) -f ( ρ µ M θ t ( x, a ) /ρ µ M ∗ ( 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 φ ∗ 0 ( x, a, x ′ ) and D φ ∗ 0 ( x, a ) , we can optimize the surrogate objective Eq. (5) via:

where E ρ µ M [ · ] is a simplification of E x,a,x ′ ∼ ρ µ M [ · ] , A φ ∗ 0 ,φ ∗ 1 ( x, a, x ′ ) = log D φ ∗ 0 ( x, a, x ′ ) -log D φ ∗ 1 ( x, a ) , and φ 0 and φ 1 are the parameters of D φ 0 and D φ 1 respectively. We learn a policy ˆ µ ≈ µ via imitation learning based on the offline dataset D real [40, 25]. Note that in the process, we ignore the term α 0 ( x, a ) for simplifying the objective. The discussion on the impacts of removing α 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 θ as the policy generator, ˆ µ as the environment, and A as the advantage function, while the second term is WERM through H M ∗ -A φ ∗ 0 ,φ ∗ 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 θ t . If the predictions of the model appear realistic, mainly assessed by D φ ∗ 0 , the propensity weights would be increased, encouraging the model to generate more

Generated

data

Offline

data

(2) update discriminators

Discriminator

Dynamics model

(1) data generation

rollout with the behavior policy

adversarial

reweighting

(3) update the dynamics model

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 φ ∗ 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.

e1_p1

e1_p0.2

e1_p0.05

e0.2_p1

e0.2_p0.2

e0.2_p0.05

e0.05_p1

e0.05_p0.2

5

10

15

20

25

30

35

root MISE

GALILEO

SL

IPW

SCIGAN

(a) GNFC

t0_bias2

t0_bias20

t1_bias2

t1_bias6

t1_bias8

t2_bias2

t2_bias6

t2_bias8

0

1

2

3

4

5

6

7

root MISE

GALILEO

SL

IPW

SCIGAN

(b) TCGA

Based on the above implementations, we propose G enerative A dversarial off LI ne counterfactua L E nvironment m O del 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

Input :

Process :

D real : offline dataset sampled from ρ µ M ∗ where µ is the behavior policy; N : total iterations;

1: Approximate a behavior policy ˆ µ via behavior cloning through offline dataset D real

2: Initialize an environment model M θ 1

3: for t = 1 : N do

4: Use ˆ µ to generate a dataset D gen with the model M θ t

5: Update the discriminators D φ 0 and D φ 1 through the second and third equations in Eq. (6) where ρ ˆ µ M θ t is estimated by D gen and ρ µ M ∗ is estimated by D real

6: Generative adversarial training for M θ t by regarding A φ ∗ 0 ,φ ∗ 1 as the advantage function and computing the gradient to M θ t , named g pg , with a standard policy gradient method like TRPO [44] or PPO [45] based on D gen .

7: Regard H M ∗ -A φ ∗ 0 ,φ ∗ 1 as the reweighting term for WERM and compute the gradient to M θ t based on D real . Record it as g sl .

8: Update the model θ t +1 ← θ t + g pg + g sl .

9: end for

5 Experiments

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 abiliy 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 µ which is constructed with beta distribution, and

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 µ 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 MISE = E [ ∫ A ( M ∗ ( x ′ | x, a ) -M ( x ′ | x, a )) 2 d a ] 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 ≤ e ≤ 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 ˆ µ is modeled with Gaussian distribution. Since IPW directly reweight loss with 1 ˆ µ , the results are sensitive to the error of ˆ µ .

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-

0.0

0.2

0.4

0.6

0.8

1.0

normalized action

30

35

40

45

50

55

60

65

averaged response

real

GALILEO

SL

IPW

SCIGAN

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

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

0.4

0.6

0.8

1.0

1.2

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(a) medium (train)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

1.4

1.6

1.8

2.0

2.2

2.4

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(b) medium-replay (test)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

2.25

2.50

2.75

3.00

3.25

3.50

3.75

4.00

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(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.

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.90

0.92

0.94

0.96

0.98

1.00

rescaled averaged response

GALILEO

SL

(a) response curves in City A

0.0

0.2

0.4

0.6

0.8

1.0

cumulative samples (rescaled to 1)

1.0

0.5

0.0

0.5

1.0

rescaled cumulative treatment effects

random

GALILEO

SL

(b) AUUC curves

0

5

10

15

20

day number

0.84

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(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-

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 ∈ { 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.

References

A. M. Alaa and M. van der Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design. In Proceedings of the 35th International Conference on Machine Learning , volume 80, pages 129-138, Stockholm, Sweden, 2018.

S. M. Ali and S. D. Silvey. A general class of coefficients of divergence of one distribution from another. Journal of the Royal Statistical Society: Series B (Methodological) , 28(1):131-142, 1966.

S. Assaad, S. Zeng, C. Tao, S. Datta, N. Mehta, R. Henao, F. Li, and L. Carin. Counterfactual representation learning with balancing weights. In The 24th International Conference on Artificial Intelligence and Statistics , volume 130 of Proceedings of Machine Learning Research , pages 1972-1980, Virtual Event, 2021.

S. Ben-David, J. Blitzer, K. Crammer, A. Kulesza, F. Pereira, and J. W. Vaughan. A theory of learning from different domains. Mach. Learn. , 79(1-2):151-175, 2010.

S. Ben-David, J. Blitzer, K. Crammer, and F. Pereira. Analysis of representations for domain adaptation. In Advances in Neural Information Processing Systems 19 , pages 137-144, British Columbia, Canada, 2006. MIT Press.

A. Betlei, E. Diemert, and M. Amini. Treatment targeting by AUUC maximization with generalization guarantees. CoRR , abs/2012.09897, 2020.

I. Bica, J. Jordon, and M. van der Schaar. Estimating the effects of continuous-valued interventions using generative adversarial networks. In Advances in Neural Information Processing Systems 33 , virtual event, 2020.

L. Buesing, T. Weber, Y. Zwols, N. Heess, S. Racanière, A. Guez, and J. Lespiau. Woulda, coulda, shoulda: Counterfactually-guided policy search. In 7th International Conference on Learning Representations , New Orleans, LA, 2019.

J. Byrd and Z. C. Lipton. What is the effect of importance weighting in deep learning? In Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 872-881, Long Beach, CA, 2019.

X. Chen, B. He, Y. Yu, Q. Li, Z. T. Qin, W. Shang, J. Ye, and C. Ma. Sim2rec: A simulator-based decision-making approach to optimize real-world long-term user engagement in sequential recommender systems. In 39th IEEE International Conference on Data Engineering , pages 3389-3402, Anaheim, CA, 2023. IEEE.

X. Chen, S. Li, H. Li, S. Jiang, Y. Qi, and L. Song. Generative adversarial user model for reinforcement learning based recommendation system. In Proceedings of the 36th International Conference on Machine Learning , pages 1052-1061, Long Beach, CA, 2019.

X. Chen, Y. Yu, Q. Li, F. Luo, Z. T. Qin, W. Shang, and J. Ye. Offline model-based adaptable policy learning. In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34 , pages 8432-8443, virtual, 2021.

X.-H. Chen, F.-M. Luo, Y. Yu, Q. Li, Z. Qin, W. Shang, and J. Ye. Offline model-based adaptable policy learning for decision-making in out-of-support regions. IEEE Transactions on Pattern Analysis Machine Intelligence , (01):1-16, 5555.

X.-H. Chen, Y. Yu, Q. Li, F.-M. Luo, Z. T. Qin, S. Wenjie, and J. Ye. Offline model-based adaptable policy learning. In Advances in Neural Information Processing Systems 34 , Virtual Conference, 2021.

C. Cortes, Y. Mansour, and M. Mohri. Learning bounds for importance weighting. In Advances in Neural Information Processing Systems 23 , pages 442-450, British Columbia, Canada, 2010. Curran Associates, Inc.

B. Eysenbach, A. Khazatsky, S. Levine, and R. Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based RL. CoRR , abs/2110.02758, 2021.

J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR , abs/2004.07219, 2020.

J. Fu, M. Norouzi, O. Nachum, G. Tucker, Z. Wang, A. Novikov, M. Yang, M. R. Zhang, Y. Chen, A. Kumar, C. Paduraru, S. Levine, and T. Paine. Benchmarks for deep off-policy evaluation. In 9th International Conference on Learning Representations , Virtual Event, Austria, 2021.

I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. C. Courville, and Y. Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems 27 , pages 2672-2680, Montréal, Canada, 2014.

T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft Actor-Critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning , pages 1856-1865, Stockholmsmässan, Sweden, 2018.

D. Hafner, T. P. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 2555-2565, Long Beach, CA, 2019.

N. Hassanpour and R. Greiner. Counterfactual regression with importance sampling weights. In Proceedings of the 28th International Joint Conference on Artificial Intelligence , pages 5880-5887, Macao, China, 2019.

J. Hiriart-Urruty and C. Lemaréchal. Fundamentals of Convex Analysis . 2001.

T. Hishinuma and K. Senda. Weighted model estimation for offline model-based reinforcement learning. In Advances in Neural Information Processing Systems 34 , pages 17789-17800, virtual, 2021.

J. Ho and S. Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems 29 , pages 4565-4573, Barcelona, Spain, 2016.

G. Imbens. The role of the propensity score in estimating dose-response functions. Econometrics eJournal , 1999.

E. L. Ionides. Truncated importance sampling. Journal of Computational and Graphical Statistics , 17(2):295-311, 2008.

F. D. Johansson, N. Kallus, U. Shalit, and D. Sontag. Learning weighted representations for generalization across designs. arXiv preprint arXiv:1802.08598 , 2018.

Y. Jung, J. Tian, and E. Bareinboim. Learning causal effects via weighted empirical risk minimization. In Advances in Neural Information Processing Systems 33 , Virtual Event, 2020.

R. Kidambi, A. Rajeswaran, P. Netrapalli, and T. Joachims. MOReL : Model-based offline reinforcement learning. CoRR , abs/2005.05951, 2020.

S. Levine, A. Kumar, G. Tucker, and J. Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR , abs/2005.01643, 2020.

Q. Liu, L. Li, Z. Tang, and D. Zhou. Breaking the curse of horizon: Infinite-horizon off-policy estimation. In Advances in Neural Information Processing Systems 31 , pages 5361-5371, Montréal, Canada, 2018.

Y. Liu, A. Swaminathan, A. Agarwal, and E. Brunskill. Provably good batch off-policy reinforcement learning without great exploration. In Advances in Neural Information Processing Systems 33 , virtual event, 2020.

C. Lu, P. J. Ball, J. Parker-Holder, M. A. Osborne, and S. J. Roberts. Revisiting design choices in offline model based reinforcement learning. In The 10th International Conference on Learning Representations , Virtual Event, 2022.

O. Nachum, Y. Chow, B. Dai, and L. Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. In Advances in Neural Information Processing Systems 32 , pages 2315-2325, BC, Canada, 2019.

X. Nguyen, M. J. Wainwright, and M. I. Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Trans. Inf. Theory , 56(11):5847-5861, 2010.

S. Nowozin, B. Cseke, and R. Tomioka. f-GAN: Training generative neural samplers using variational divergence minimization. In Advances in Neural Information Processing Systems 29 , pages 271-279, Barcelona, Spain, 2016.

J. Pearl. Causality . Cambridge University Press, 2009.

S. Pitis, E. Creager, and A. Garg. Counterfactual data augmentation using locally factored dynamics. In Advances in Neural Information Processing Systems 33 , virtual, 2020.

D. Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation , 3(1):88-97, 1991.

T. Qin, T. Wang, and Z. Zhou. Budgeted heterogeneous treatment effect estimation. In Proceedings of the 38th International Conference on Machine Learning, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 8693-8702, 2021.

J. Quinonero-Candela, M. Sugiyama, A. Schwaighofer, and N. D. Lawrence. Dataset shift in machine learning . Mit Press, 2008.

P. R. Rosenbaum and D. B. Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika , 70:41-55, 1983.

J. Schulman, S. Levine, P. Abbeel, M. I. Jordan, and P. Moritz. Trust region policy optimization. In Proceedings of the 32nd International Conference on Machine Learning , volume 37, pages 1889-1897, Lille, France, 2015.

J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. CoRR , abs/1707.06347, 2017.

W. Shang, Y. Yu, Q. Li, Z. T. Qin, Y. Meng, and J. Ye. Environment reconstruction with hidden confounders for reinforcement learning based recommendation. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 566-576, Anchorage, AK, 2019.

J. Shi, Y . Yu, Q. Da, S. Chen, and A. Zeng. Virtual-taobao: Virtualizing real-world online retail environment for reinforcement learning. In The 33th AAAI Conference on Artificial Intelligence , pages 4902-4909, Honolulu, Hawaii, 2019.

H. Shimodaira. Improving predictive inference under covariate shift by weighting the loglikelihood function. Journal of statistical planning and inference , 90(2):227-244, 2000.

S. A. Sontakke, A. Mehrjou, L. Itti, and B. Schölkopf. Causal curiosity: RL agents discovering self-supervised experiments for causal representation learning. In Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 9848-9858, Virtual Event, 2021.

P. Spirtes. Introduction to causal inference. J. Mach. Learn. Res. , 11:1643-1662, aug 2010.

R. S. Sutton and A. G. Barto. Reinforcement learning - an introduction . Adaptive computation and machine learning. MIT Press, 1998.

A. Swaminathan and T. Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. J. Mach. Learn. Res. , 16:1731-1755, 2015.

E. Todorov, T. Erez, and Y. Tassa. MuJoCo: A physics engine for model-based control. In Proceedings of the 24th IEEE/RSJ International Conference on Intelligent Robots and Systems , pages 5026-5033, Vilamoura, Portugal, 2012.

C. Voloshin, N. Jiang, and Y. Yue. Minimax model learning. In A. Banerjee and K. Fukumizu, editors, The 24th International Conference on Artificial Intelligence and Statistics , volume 130 of Proceedings of Machine Learning Research , pages 1612-1620, Virtual Event, 2021.

T. Wang, T. Qin, and Z. Zhou. Estimating possible causal effects with latent variables via adjustment. In International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA , volume 202, pages 36308-36335, 2023.

T.-Z. Wang, X.-Z. Wu, S.-J. Huang, and Z.-H. Zhou. Cost-effectively identifying causal effects when only response variable is observable. In Proceedings of the 37th International Conference on Machine Learning , pages 10060-10069, 2020.

S. Yang, S. Zhang, Y. Feng, and M. Zhou. A unified framework for alternating offline model training and policy learning. CoRR , abs/2210.05922, 2022.

J. Yoon, J. Jordon, and M. van der Schaar. GANITE: estimation of individualized treatment effects using generative adversarial nets. In 6th International Conference on Learning Representations , 2018.

T. Yu, G. Thomas, L. Yu, S. Ermon, J. Y. Zou, S. Levine, C. Finn, and T. Ma. MOPO: modelbased offline policy optimization. In Advances in Neural Information Processing Systems 33 , virtual event, 2020.

R. Zhang, B. Dai, L. Li, and D. Schuurmans. Gendice: Generalized offline estimation of stationary values. In 8th International Conference on Learning Representations , Addis Ababa, Ethiopia, 2020.

W. Y. Zou, S. Du, J. Lee, and J. O. Pedersen. Heterogeneous causal learning for effectiveness optimization in user marketing. CoRR , abs/2004.09702, 2020.

Appendix

Table of Contents

A Proof of Theoretical Results

In the proof section, we replace the notation of E with an integral for brevity. Now we rewrite the original objective ¯ L ( ρ β M ∗ , M ) as:

where ω ( x, a ) = ρ β M ∗ ( x,a ) ρ µ M ∗ ( x,a ) and ∥ ρ β M ∗ ( · , · ) ∥ 2 2 = ∫ X , A ρ β M ∗ ( x, a ) 2 d a d x , which is the squared l 2 -norm. In an MDP, given any policy π , ρ π M ∗ ( x, a, x ′ ) = ρ π M ∗ ( x ) π ( a | x ) M ∗ ( x ′ | x, a ) where ρ π M ∗ ( x ) denotes the occupancy measure of x for policy π , which can be defined [51, 25] as ρ π M ∗ ( x ) := (1 -γ ) E x 0 ∼ ρ 0 [ ∑ ∞ t =0 γ t Pr( x t = x | x 0 , M ∗ )] where Pr π [ x t = x | x 0 , M ∗ ] is the state visitation probability that π starts at state x 0 in model M ∗ and receive x at timestep t and γ ∈ [0 , 1] is the discount factor.

min

M

∈M

β

∈

Π

L

(

ρ

β

M

∗

, M

)

estimate the optimal adversarial

distribution

ρ

¯

β

∗

M

∗

given

M

(Sec. A.1)

derive

¯

ρ

¯

β

∗

M

∗

as a generalized rep-

resentation of

ρ

¯

β

∗

M

∗

(Sec. A.2)

let

ρ

¯

β

∗

M

∗

as the distribution of the best-response

policy

β

∈

Π

L

(

ρ

β

M

∗

, M

)

, we have the sur-

rogate objective

min

M

∈M

L

(¯

ρ

¯

β

∗

M

∗

, M

)

(Sec. A.3)

tractable solution (Sec. A.4)

intermediary policy

κ

generator function

f

(model an easy-to-

estimate distribution of

the best-response policy

¯

β

∗

)

approximate

¯

ρ

¯

β

∗

M

∗

with

variational representation

The overall pipeline to model the tractable solution to AWRM is given in Fig. 8. In the following, we will summarize the modelling process based on Fig. 8. We first approximate the optimal distribution ρ ¯ β ∗ M ∗ via Lemma. A.1.

Lemma A.1. Given any M in ¯ L ( ρ β M ∗ , M ) , the distribution ρ ¯ β ∗ M ∗ ( x, a ) of the ideal best-response policy ¯ β ∗ satisfies:

where D KL ( M ∗ ( ·| x, a ) , M ( ·| x, a )) is the Kullback-Leibler (KL) divergence between M ∗ ( ·| x, a ) and M ( ·| x, a ) , H M ∗ ( x, a ) denotes the entropy of M ∗ ( ·| x, a ) , and α M is the regularization coefficient α in Eq. (4) and also as a normalizer of Eq. (8) .

Note that the ideal best-response policy ¯ β ∗ is not the real best-response policy β ∗ . The distribution ρ ¯ β ∗ M ∗ is an approximation of the real optimal adversarial distribution. We give a discussion of the rationality of the ideal best-response policy ¯ β ∗ as a replacement of the real best-response policy β ∗ in Remark A.4. Intuitively, ρ ¯ β ∗ M ∗ has larger densities on the data where the divergence between the approximation model and the real model (i.e., D KL ( M ∗ ( ·| x, a ) , M ( ·| x, a )) ) is larger or the stochasticity of the real model (i.e., H M ∗ ) is larger.

However, the integral process of D KL in Eq. (8) is intractable in the offline setting as it explicitly requires the conditional probability function of M ∗ . Our solution to solve the problem is utilizing the offline dataset D real as the empirical joint distribution ρ µ M ∗ ( x, a, x ′ ) and adopting practical

techniques for distance estimation on two joint distributions, like GAN [19, 37], to approximate Eq. (8). To adopt that solution, we should first transform Eq. (8) into a form under joint distributions. Without loss of generality, we introduce an intermediary policy κ , of which µ can be regarded as a specific instance. Then we have M ( x ′ | x, a ) = ρ κ M ( x, a, x ′ ) /ρ κ M ( x, a ) for any M if ρ κ M ( x, a ) > 0 . Assuming ∀ x ∈ X , ∀ a ∈ A , ρ κ M ∗ ( x, a ) > 0 if ρ ¯ β ∗ M ∗ ( x, a ) > 0 , which will hold when κ overlaps with µ , then Eq. (8) can transform to:

where α 0 ( x, a ) = α M ρ κ M ∗ ( x, a ) . We notice that the form ρ κ M ∗ log ρ κ M ∗ ρ κ M is the integrated function in reverse KL divergence, which is an instance of f function in f -divergence [2]. Replacing that form with f function, we obtain a generalized representation of ρ ¯ β ∗ M ∗ :

where f : R + → R is a convex and lower semi-continuous (l.s.c.) function. ¯ ρ ¯ β ∗ M ∗ gives a generalized representation of the optimal adversarial distribution to maximize the error of the model. Based on Eq. (9), we have a surrogate objective of AWRM which can avoid querying M ∗ to construct ρ β ∗ M ∗ :

Theorem A.2. Let ¯ ρ ¯ β ∗ M ∗ as the data distribution of the best-response policy ¯ β ∗ in Eq. (4) under model M θ parameterized by θ , then we can find the optimal θ ∗ of min θ max β ∈ Π ¯ L ( ρ β M ∗ , M θ ) (Eq. (4) ) via iteratively optimizing the objective θ t +1 = min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) , where ¯ ρ ¯ β ∗ M ∗ is approximated via the last-iteration model M θ t . Based on Corollary A.9, we derive an upper bound objective for min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) :

where E ρ κ M ∗ [ · ] denotes E x,a,x ′ ∼ ρ κ M ∗ [ · ] , f is a l.s.c function satisfying f ′ ( x ) ≤ 0 , ∀ x ∈ X , and α 0 ( x, a ) = α M θ t ρ κ M ∗ ( x, a ) .

Thm. A.2 approximately achieve AWRM by using κ and a pseudo-reweighting module W . W assigns learning propensities for data points with larger differences between distributions ρ κ M θ t and ρ κ M ∗ . By adjusting the weights, the learning process will exploit subtle errors in any data point, whatever how many proportions it contributes, to correct potential generalization errors on counterfactual data.

Remark A.3. In practice, we need to use real-world data to construct the distribution ρ κ M ∗ . In the offline model-learning setting, we only have a real-world dataset D real collected by the behavior policy µ , which is the empirical distribution of ρ µ M ∗ . Let κ = µ , we have

which is Eq. (5) in the main body.

In Thm. A.2, the terms f ( ρ κ M θ t ( x, a, x ′ ) /ρ κ M ∗ ( x, a, x ′ )) -f ( ρ κ M θ t ( x, a ) /ρ κ M ∗ ( 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 φ ∗ 0 ( x, a, x ′ ) and D φ ∗ 0 ( x, a ) , letting κ = µ , we can optimize the surrogate objective Eq. (5) via:

where E ρ µ M [ · ] is a simplification of E x,a,x ′ ∼ ρ µ M [ · ] , A φ ∗ 0 ,φ ∗ 1 ( x, a, x ′ ) = log D φ ∗ 0 ( x, a, x ′ ) -log D φ ∗ 1 ( x, a ) , and φ 0 and φ 1 are the parameters of D φ 0 and D φ 1 respectively. We learn a policy ˆ µ ≈ µ via imitation learning based on the offline dataset D real [40, 25]. Note that in the process, we ignore the term α 0 ( x, a ) for simplifying the objective. The discussion on the impacts of removing α 0 ( x, a ) is left in App. B.

A.1 Proof of Lemma A.1

Proof. Given a transition function M of an MDP, the distribution of the best-response policy β ∗ satisfies:

We know that the occupancy measure ρ β M ∗ is a density function with a constraint ∫ X ∫ A ρ β M ∗ ( x, a )d a d x = 1 . Assuming the occupancy measure ρ β M ∗ has an upper bound c , that is 0 ≤ ρ β M ∗ ( x, a ) ≤ c, ∀ a ∈ A , ∀ x ∈ X , constructing a regularization coefficient α M = ∫ X ∫ A ( D KL ( M ∗ ( ·| x, a ) , M ( ·| x, a )) + H M ∗ ( x, a ))d x d a as a constant value given any M , then we have

which is the optimal density function of Eq. (7) with α = α M .

Note that in some particular M ∗ , we still cannot construct a β that can generate an occupancy specified by g ( x, a ) /α M for any M . We can only claim the distribution of the ideal best-response policy ¯ β ∗ satisfies:

where α M is a normalizer that α M = ∫ X ∫ A ( D KL ( M ∗ ( ·| x, a ) , M ( ·| x, a )) + H M ∗ ( x, a ))d x d a . We give a discussion of the rationality of the ideal best-response policy ¯ β ∗ as a replacement of the real best-response policy β ∗ in Remark A.4.

Remark A.4. The optimal solution Eq. (10) relies on g ( x, a ) . In some particular M ∗ , it is intractable to derive a β that can generate an occupancy specified by g ( x, a ) /α M . Consider the following case: a state x 1 in M ∗ might be harder to reach than another state x 2 , e.g., M ∗ ( x 1 | x, a ) < M ∗ ( x 2 | x, a ) , ∀ x ∈ X , ∀ a ∈ A , then it is impossible to find a β that the occupancy satisfies ρ β M ∗ ( x 1 , a ) > ρ β M ∗ ( x 2 , a ) . In this case, Eq. (10) can be a sub-optimal solution. Since this work focuses on task-agnostic solution derivation while the solution to the above problem should rely on the specific description of M ∗ , we leave it as future work. However, we point out that Eq. (10) is a reasonable re-weighting term even as a sub-optimum: ρ ¯ β ∗ M ∗ gives larger densities on the data where the distribution distance between the approximation model and the real model (i.e., D KL ( M ∗ , M ) ) is larger or the stochasticity of the real model (i.e., H M ∗ ) is larger.

A.2 Proof of Eq. (9)

The integral process of D KL in Eq. (8) is intractable in the offline setting as it explicitly requires the conditional probability function of M ∗ . Our motivation for the tractable solution is utilizing the offline dataset D real as the empirical joint distribution ρ µ M ∗ ( x, a, x ′ ) and adopting practical techniques for distance estimation on two joint distributions, like GAN [19, 37], to approximate Eq. (8). To adopt that solution, we should first transform Eq. (8) into a form under joint distributions. Without loss of generality, we introduce an intermediary policy κ , of which µ can be regarded as a specific instance. Then we have M ( x ′ | x, a ) = ρ κ M ( x, a, x ′ ) /ρ κ M ( x, a ) for any M if ρ κ M ( x, a ) > 0 . Assuming ∀ x ∈ X , ∀ a ∈ A , ρ κ M ∗ ( x, a ) > 0 if ρ ¯ β ∗ M ∗ ( x, a ) > 0 , which will hold when κ overlaps with µ , then Eq. (8) can transform to:

where α 0 ( x, a ) = α M ρ κ M ∗ ( x, a ) .

Definition A.5 ( f -divergence) . Given two distributions P and Q , two absolutely continuous density functions p and q with respect to a base measure d x defined on the domain X , we define the f -divergence [37],

where the generator function f : R + → R is a convex, lower-semicontinuous function.

We notice that the terms ρ κ M ∗ ( x, a, x ′ ) log ρ κ M ∗ ( x,a,x ′ ) ρ κ M ( x,a,x ′ ) and ρ κ M ∗ ( x, a ) log ρ κ M ∗ ( x,a ) ρ κ M ( x,a ) are the integrated functions in reverse KL divergence, which is an instance of f function in f -divergence (See Reverse-

KL divergence of Tab.1 in [37] for more details). Replacing that form q log q p with qf ( p q ) , we obtain a generalized representation of ρ ¯ β ∗ M ∗ :

A.3 Proof of Thm. A.2

We first introduce several useful lemmas for the proof.

Lemma A.6. Rearrangement inequality The rearrangement inequality states that, for two sequences a 1 ≥ a 2 ≥ . . . ≥ a n and b 1 ≥ b 2 ≥ . . . ≥ b n , the inequalities a 1 b 1 + a 2 b 2 + · · · + a n b n ≥ a 1 b π (1) + a 2 b π (2) + · · · + a n b π ( n ) ≥ a 1 b n + a 2 b n -1 + · · · + a n b 1 hold, where π (1) , π (2) , . . . , π ( n ) is any permutation of 1 , 2 , . . . , n .

Lemma A.7. For two sequences a 1 ≥ a 2 ≥ . . . ≥ a n and b 1 ≥ b 2 ≥ . . . ≥ b n , the inequalities

hold.

Proof. By rearrangement inequality, we have

Then we have

Now we extend Lemma A.7 into the continuous integral scenario:

Lemma A.8. Given X ⊂ R , for two functions f : X → R and g : X → R that f ( x ) ≥ f ( y ) if and only if g ( x ) ≥ g ( y ) , ∀ x, y ∈ X , the inequality

holds, where p : X → R and p ( x ) > 0 , ∀ x ∈ X and ∫ X p ( x )d x = 1 .

Proof. Since ( f ( x ) -f ( y ))( g ( x ) -g ( y )) ≥ 0 , ∀ x, y ∈ X , we have

Corollary A.9. Let g ( p ( x ) q ( x ) ) = -log p ( x ) q ( x ) where p ( x ) > 0 , ∀ x ∈ X and q ( x ) > 0 , ∀ x ∈ X , for υ > 0 , the inequality

holds if f ′ ( x ) ≤ 0 , ∀ x ∈ X . It is not always satisfied for f functions of f -divergence. We list a comparison of f on that condition in Tab. 3.

Proof. g ′ ( x ) = -log x = -1 x < 0 , ∀ x ∈ X . Suppose f ′ ( x ) ≤ 0 , ∀ x ∈ X , we have f ( x ) ≥ f ( y ) if and only if g ( x ) ≥ g ( y ) , ∀ x, y ∈ X holds. Thus f ( υ p ( x ) q ( x ) ) ≥ f ( υ p ( y ) q ( y ) ) if and only if g ( p ( x ) q ( x ) ) ≥ g ( p ( y ) q ( y ) ) , ∀ x, y ∈ X holds for all υ > 0 . By defining F ( x ) = f ( υ p ( x ) q ( x ) )) and G ( x ) = g ( p ( x ) q ( x ) ) and using Lemma A.8, we have:

Then we know

holds.

Now, we prove Thm. A.2. For better readability, we first rewrite Thm. A.2 as follows:

Theorem A.10. Let ¯ ρ ¯ β ∗ M ∗ as the data distribution of the best-response policy ¯ β ∗ in Eq. (4) under model M θ parameterized by θ , then we can find the optimal θ ∗ of min θ max β ∈ Π ¯ L ( ρ β M ∗ , M θ ) (Eq. (4) ) via iteratively optimizing the objective θ t +1 = min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) , where ¯ ρ ¯ β ∗ M ∗ is approximated via the last-iteration model M θ t . Based on Corollary A.9, we have an upper bound objective for

min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) and derive the following objective

where α 0 ( x, a ) = α M θ t ρ κ M ∗ ( x, a ) , E ρ κ M ∗ [ · ] denotes E x,a,x ′ ∼ ρ κ M ∗ [ · ] , f is the generator function in f -divergence which satisfies f ′ ( x ) ≤ 0 , ∀ x ∈ X , and θ is the parameters of M . M θ t denotes a probability function with the same parameters as the learned model (i.e., ¯ θ = θ ) but the parameter is fixed and only used for sampling.

Proof. Let ¯ ρ ¯ β ∗ M ∗ as the data distribution of the best-response policy ¯ β ∗ in Eq. (4) under model M θ parameterized by θ , then we can find the optimal θ t +1 of min θ max β ∈ Π ¯ L ( ρ β M ∗ , M θ ) (Eq. (4)) via iteratively optimizing the objective θ t +1 = min θ ¯ L (¯ ρ ¯ β ∗ M ∗ , M θ ) , where ¯ ρ ¯ β ∗ M ∗ is approximated via the last-iteration model M θ t :

where M θ t is introduced to approximate the term ¯ ρ ¯ β ∗ M ∗ and fixed when optimizing θ . In Eq. (15), ∥ ρ β M ∗ ( · , · ) ∥ 2 2 for Eq. (7) is eliminated as it does not contribute to the gradient of θ . Assume f ′ ( x ) ≤ 0 , ∀ x ∈ X , let υ ( x, a ) := ρ κ Mθt ( x,a ) ρ κ M ∗ ( x,a ) > 0 , p ( x ′ | x, a ) = M θ ( x ′ | x, a ) , and q ( x ′ | x, a ) = M ∗ ( x ′ | x, a ) , the first inequality can be derived by adopting Corollary A.9 and eliminating the first H M ∗ since it does not contribute to the gradient of θ .

A.4 Proof of the Tractable Solution

Now we are ready to prove the tractable solution:

Proof. The core challenge is that the term f ( ρ κ Mθt ( x,a,x ′ ) ρ κ M ∗ ( x,a,x ′ ) ) -f ( ρ κ Mθt ( x,a ) ρ κ M ∗ ( x,a ) ) is still intractable. In the following, we give a tractable solution to Thm. A.2. First, we resort to the first-order approximation. Given some u ∈ (1 -ξ, 1 + ξ ) , ξ > 0 , we have

where f ′ is the first-order derivative of f . By Taylor's formula and the fact that f ′ ( u ) of the generator function f is bounded in (1 -ξ, 1 + ξ ) , the approximation error is no more than O ( ξ 2 ) . Substituting u with p ( x ) q ( x ) in Eq. (17), the pattern f ( p ( x ) q ( x ) ) in Eq. (16) can be converted to p ( x ) q ( x ) f ′ ( p ( x ) q ( x ) ) -f ′ ( p ( x ) q ( x ) ) + f (1) , then we have:

Note that the part ρ κ M ∗ ( x, a ) in ρ κ M ∗ ( x, a, x ′ ) can be canceled because of α 0 ( x, a ) = α M θ t ρ κ M ∗ ( x, a ) , but we choose to keep it and ignore α 0 ( x, a ) . The benefit is that we can estimate ρ κ M ∗ ( x, a, x ′ ) from an empirical data distribution through data collected by κ in M ∗ directly, rather than from a uniform distribution which is harder to be generated. Although keeping ρ κ M ∗ ( x, a ) incurs extra bias in theory, the results in our experiments show that it has not made significant negative effects in practice. We

leave this part of modeling in future work. In particular, by ignoring α 0 ( x, a ) , we have:

We can estimate f ′ ( ρ κ Mθt ( x,a ) ρ κ M ∗ ( x,a ) ) and f ′ ( ρ κ Mθt ( x,a,x ′ ) ρ κ M ∗ ( x,a,x ′ ) ) through Lemma A.11.

Lemma A.11 ( f ′ ( p q ) estimation [36]) . Given a function T φ : X → R parameterized by φ ∈ Φ , if f is convex and lower semi-continuous, by finding the maximum point of φ in the following objective:

we have f ′ ( p ( x ) q ( x ) ) = T φ ∗ ( x ) . f ∗ is Fenchel conjugate of f [23].

In particular,

then we have f ′ ( ρ κ Mθt ( x,a,x ′ ) ρ κ M ∗ ( x,a,x ′ ) ) ≈ T φ ∗ 0 ( x, a, x ′ ) and f ′ ( ρ κ Mθt ( x,a ) ρ κ M ∗ ( x,a ) ) ≈ T φ ∗ 1 ( x, a ) . Given φ ∗ 0 and φ ∗ 1 , let A φ ∗ 0 ,φ ∗ 1 ( x, a, x ′ ) = T φ ∗ 0 ( x, a, x ′ ) -T φ ∗ 1 ( x, a ) , then we can optimize θ via:

Based on the specific f -divergence, we can represent T and f ∗ ( T ) with a discriminator D φ . It can be verified that f ( u ) = u log u -( u +1)log( u +1) , T φ ( u ) = log D φ ( u ) , and f ∗ ( T φ ( u )) = -log(1 -D φ ( u )) proposed in [37] satisfies the condition f ′ ( x ) ≤ 0 , ∀ x ∈ X (see Tab. 3). We select the former in the implementation and convert the tractable solution to:

where A φ ∗ 0 ,φ ∗ 1 ( x, a, x ′ ) = log D φ ∗ 0 ( x, a, x ′ ) -log D φ ∗ 1 ( x, a ) , E ρ κ Mθt [ · ] is a simplification of E x,a,x ′ ∼ ρ κ Mθt [ · ] .

Remark A.12. In practice, we need to use the real-world data to construct the distribution ρ κ M ∗ and the generative data to construct ρ κ M θ t . In the offline model-learning setting, we only have a real-world dataset D real collected by the behavior policy µ . We can learn a policy ˆ µ ≈ µ via imitation learning based on D real [40, 25] and let ˆ µ be the policy κ . Then we can regard D real as the empirical

data distribution of ρ κ M ∗ and the trajectories collected by ˆ µ in the model M θ t as the empirical data distribution of ρ κ M θ t . Based on the above specializations, we have:

which is Eq. (6) in the main body.

B Discussion and Limitations of the Theoretical Results

We summarize the limitations of current theoretical results and future work as follows:

As discussed in Remark A.4, the solution Eq. (10) relies on ρ β M ∗ ( x, a ) ∈ [0 , c ] , ∀ a ∈ A , ∀ x ∈ X . In some particular M ∗ , it is intractable to derive a β that can generate an occupancy specified by g ( x, a ) /α M . If more knowledge of M ∗ or β ∗ is provided or some mild assumptions can be made on the properties of M ∗ or β ∗ , we may model ρ in a more sophisticated approach to alleviating the above problem.

In the tractable solution derivation, we ignore the term α 0 ( x, a ) = α M θ t ρ κ M ∗ ( x, a ) (See Eq. (19)). The benefit is that ρ κ M ∗ ( x, a, x ′ ) in the tractable solution can be estimated through offline datasets directly. Although the results in our experiments show that it does not produce significant negative effects in these tasks, ignoring ρ κ M ∗ ( x, a ) indeed incurs extra bias in theory. In future work, techniques for estimating ρ κ M ∗ ( x, a ) [33] can be incorporated to correct the bias. On the other hand, α M θ t is also ignored in the process. α M θ t can be regarded as a global rescaling term of the final objective Eq. (19). Intuitively, it constructs an adaptive learning rate for Eq. (19), which increases the step size when the model is better fitted and decreases the step size otherwise. It can be considered to further improve the learning process in future work, e.g., cooperating with empirical risk minimization by balancing the weights of the two objectives through α M θ t .

C Societal Impact

This work studies a method toward counterfactual environment model learning. Reconstructing an accurate environment of the real world will promote the wide adoption of decision-making policy optimization methods in real life, enhancing our daily experience. We are aware that decision-making policy in some domains like recommendation systems that interact with customers may have risks of causing price discrimination and misleading customers if inappropriately used. A promising way to reduce the risk is to introduce fairness into policy optimization and rules to constrain the actions (Also see our policy design in Sec. G.1.3). We are involved in and advocating research in such directions. We believe that business organizations would like to embrace fair systems that can ultimately bring long-term financial benefits by providing a better user experience.

D AWRM-oracle Pseudocode

We list the pseudocode of AWRM-oracle in Alg. 2.

E Implementation

E.1 Details of the GALILEO Implementation

The approximation of Eq. (17) holds only when p ( x ) /q ( x ) is close to 1 , which might not be satisfied. To handle the problem, we inject a standard supervised learning loss

Algorithm 2 AWRMwith Oracle Counterfactual Datasets

Input :

Φ : policy space; N : total iterations

Process :

1: Generate counterfactual datasets {D π ϕ } for all adversarial policies π ϕ , ϕ ∈ Φ

2: Initialize an environment model M θ

3: for i = 1:N do

4: Select D π ϕ with worst prediction errors through M θ from {D π ϕ }

5: Optimize M θ with standard supervised learning based on D π ϕ

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 D collected by the behavior policy µ . We learn a policy ˆ µ ≈ µ via behavior cloning with D [40, 25] and let ˆ µ be the policy κ . We regard D as the empirical data distribution of ρ κ M ∗ and the trajectories collected by ˆ µ in the model M θ t as the empirical data distribution of ρ κ M θ t . But the assumption ∀ x ∈ X , ∀ a ∈ A , µ ( a | x ) > 0 might not be satisfied. In behavior cloning, we model ˆ µ with a Gaussian distribution and constrain the lower bound of the variance with a small value ϵ µ > 0 to keep the assumption holding. Besides, we add small Gaussian noises u ∼ N (0 , ϵ D ) to the inputs of D φ to handle the mismatch between ρ µ M ∗ and ρ ˆ µ M ∗ due to ϵ µ . In particular, for φ 0 and φ 1 learning, we have:

On the other hand, we notice that the first term in Eq. (20) is similar to the objective of GAIL [25] by regarding M θ as the policy to learn and κ 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 κ and M θ , where the reward is defined by the discriminator D , i.e., r ( x, a, x ′ ) = log D ( x, a, x ′ ) . M θ is learned to maximize the cumulative rewards. With advanced policy gradient methods [44, 45], the objective is converted to max θ [ A φ ∗ 0 ,φ ∗ 1 ( x, a, x ′ ) log M θ ( x, a, x ′ ) ] , where A = Q κ M θ t -V κ M θ t Q κ M ¯ θ ( x, a, x ′ ) = E [ ∑ ∞ t =0 γ t r ( x t , a t , x t +1 ) | ( x t , a t , x t +1 ) = ( x, a, x ′ ) , κ, M θ t ] , and V κ M ¯ θ ( x, a ) =

, E M ¯ θ [ Q κ M ¯ θ ( x, a, x ′ ) ] . A in Eq. (20) can also be constructed similarly. Although it looks unnecessary in theory since the one-step optimal model M θ 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 φ ∗ 0 ,φ ∗ 1 as A φ ∗ 0 ,φ ∗ 1 = Q κ M θ t -V κ M θ t , where Q κ M θ t ( x, a, x ′ ) = E [∑ ∞ t γ t log D φ ∗ 0 ( x t , a t , x t +1 ) | ( x t , a t , x t +1 ) = ( x, a, x ′ ) , κ, M θ t ] and V κ M θ t ( x, a ) = E [∑ ∞ t γ t log D φ ∗ 1 ( x t , a t ) | ( x t , a t ) = ( x, a ) , κ, M θ t ] . To give an algorithm for singlestep environment model learning, we can just set γ in Q and V to 0 .

Algorithm 3 GALILEO pseudocode

Input :

N : total iterations;

D real : offline dataset sampled from ρ µ M ∗ where µ is the behavior policy;

Process :

1: Approximate a behavior policy ˆ µ via behavior cloning

2: Initialize an environment model M θ 1

3: for t = 1 : N do

4: Use ˆ µ to generate a dataset D gen with the model M θ t

5: Update the discriminators D φ 0 and D φ 1 via Eq. (25) and Eq. (26) respectively, where ρ ˆ µ M θ t is estimated by D gen and ρ µ ∗ is estimated by D real

M 6: Update Q and V via Eq. (23) and Eq. (24) through D gen , D φ 0 , and D φ 1

7: Update the model M θ 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 pg

8: if p 0 < E D gen [ D φ 0 ( x t , a t , x t +1 )] < p 1 then

9: Compute the gradient of M θ t via the second term of Eq. (22) and record it as g sl

10: else

11: Compute the gradient of M θ t via Eq. (21) and record it as g sl

12: end if

13: Rescale g sl via Eq. (27)

14: Update the model M θ t via the gradient g sl and obtain M θ t +1

15: end for

offline dataset

𝒟

!"#$

behavior policy

"

𝜇(𝑎|𝑠)

dynamics model

𝑀

%

!

pretrain: approximated the behavior

policy (line 1)

(1)

generate a dataset (line 4)

fake dataset

𝒟

&"'

discriminators and value functions

𝐷

(

"

(𝑠, 𝑎, 𝑠′)

,

𝐷

(

#

𝑠, 𝑎

,

𝑄

and

𝑉

(2)

update the discriminators

and the value functions

(line 5 to 6)

(3)

approximate the AWRM through Eq. 23

to update the dynamics model

(line 7 to 14)

neural networks

fake dataset

offline dataset

By adopting the above implementation techniques, we convert the objective into the following formulation

where A φ ∗ 0 ,φ ∗ 1 ( x, a, x ′ ) = Q κ M θ ( x, a, x ′ ) -V κ M θ ( 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 φ ∗ 0 ,φ ∗ 1 + H M ∗ . Since H M ∗ is unknown, we use H M θ 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 pg || 2 . Then we compute the l 2 -norm of the gradient of the second term of Eq. (22), name || g sl || 2 . Finally, we rescale the gradients of the second term g sl by

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

where A φ ∗ ( x, a, x ′ ) = log D φ ∗ ( x, a, x ′ ) . In the single-step scenario, ρ κ M θ t ( x, a, x ′ ) = ρ 0 ( x ) κ ( a | x ) M θ t ( x ′ | a, x ) . The term E ρ κ Mθt [ A φ ∗ ( x, a, x ′ ) log M θ ( x ′ | x, a )] can convert to E ρ κ Mθ [log D φ ∗ ( x, a, x ′ )] by replacing the gradient of M θ t ( x ′ | x, a ) ∇ θ log M θ ( x ′ | x, a ) with ∇ θ M θ ( 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 θ as the 'policy' to imitate and ˆ µ as the 'environment' to collect data. In the context of GALILEO, the objective of GAIL is:

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 train ( x ) to minimize the prediction risks in the domain with distribution P test ( x ) , where P test = P test . To solve the problem, we can train a weighted objective with max M E x ∼ P train [ P test ( x ) P 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 P test ( x ) P train ( x ) = P ( x ) β ( a | x ) P ( x ) µ ( a | x ) = β ( a | x ) µ ( a | x ) , where µ and β 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 P test ( x ) P train ( x ) = P ( x ) β ( a | x ) P ( x ) µ ( a | x ) = β ( a | x ) µ ( a | x ) ∝ 1 µ ( a | x ) . 1 µ ( a | x ) is also known as inverse propensity score (IPS). In [28], it assumes that the policy in the target domain is predefined as β ( a | x ) before environment model learning, then it uses β µ 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 µ and the policy to evaluation β , 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 β . These methods attempt to solve an accurate approximation of ω ( s, a | ρ β ) . 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 β in the policy space Π . The core of this problem is how to obtain an approximation of the density function of the best-response β ∗ corresponding to the current model, and then approximate the corresponding ω ( s, a | ρ β ∗ ) ). It should be emphasized that since β ∗ 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 V and W , satisfying the target policy's value V β and density ratio ω β are covered, i.e., V β ∈ V and ω β ∈ W . Different from previous articles, our approach uses adversarial weighting to learn a universal model that provides good feedback for any target policy β ∈ Π , 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 µ 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 µ GNFC is

where ϵ is a sample noise, which will be discussed later. The environment includes two parts:

(1) response function M 1 ( y | x, a ) :

(2) mapping function M 2 ( x ′ | x, y ) :

The transition function M ∗ is a composite of M ∗ ( x ′ | x, a ) = M 2 ( x ′ | 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 ϵ . 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 ϵ from a uniform distribution U ( -e, e ) with probability p . That is, ϵ = 0 with probability 1 -p and ϵ ∼ U ( -e, e ) with probability p . Then with larger p , there are more samples in the dataset violating the negative correlation (i.e., µ 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.

0

10

20

30

40

50

timestep

10

20

30

40

50

60

response

(a)

y

in collected trajectories.

0

10

20

30

40

50

timestep

0

1

2

3

4

action

(b)

a

in collected trajectories.

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, 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, v 1 , v 2 , 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 | x ∼ Beta( α, β ) . α ≥ 1 controls the sampling bias and β = α -1 a ∗ + 2 -α , where a ∗ is the optimal treatment. This setting of β ensures that the mode of Beta( α, β ) is a ∗ .

The calculation of treatment response and optimal treatment are shown in Table 4.

We conduct experiments on three different treatments separately and change the value of bias α to assess the robustness of different methods to treatment bias. When the bias of treatment is large, which means α is large, the training set contains data with a strong bias on treatment so it would be difficult for models to appropriately predict the treatment responses out of the distribution of training data.

G.1.3 Budget Allocation task to the Time period (BAT)

1. customers send

take-out food orders remotely

based on the platform

3. delivery clerks

take food

from the stores

2. stores make food

4. delivery clerks

send the food

to the customers

customers

stores

delivery

clerk

5. The platform will pay

the delivery clerks for

fulfilling the order

platform

We deploy GALILEO in a real-world large-scale food-delivery platform. The platform contains various food stores, and food delivery clerks. The overall workflow is as follows: the platform presents the nearby food stores to the customers and the customers make orders, i.e., purchase take-out foods from some stores on the platform. The food delivery clerks can select orders from the platform to fulfill. After an order is selected to fulfill, the delivery clerks will take the ordered take-out foods from the stores and then send the food to the customers. The platform will pay the delivery clerks (mainly in proportion to the distance between the store and the customers' location) once the orders are fulfilled. An illustration of the workflow can be found in Fig. 11.

0.0

0.2

0.4

0.6

0.8

1.0

normalized actions

0.2

0.4

0.6

0.8

1.0

normalized fulfilled order amount

However, there is an imbalance problem between the demanded orders from customers and the supply of delivery clerks to fulfill these orders. For example, at peak times like lunchtime, there will be many more demanded orders than at other times, and the existed delivery clerks might not be able to fulfill all of these orders timely. The goal of the Budget Allocation task to the Time period (BAT) is to handle the imbalance problem in time periods by sending reasonable allowances to different time periods. More precisely, the goal of BAT is to make all orders (i.e., the demand) sent in different time periods can be fulfilled (i.e., the supply) timely.

The core challenge of the environment model learning in BAT tasks is similar to the challenge in Fig. 1. Specifically, the behavior policy in BAT tasks is a human-expert policy, which will tend to increase the budget of allowance in the time periods with a lower supply of delivery clerks, otherwise will decrease the budget (Fig. 12 gives an instance of this phenomenon in the real data).

To handle the imbalance problem in different time periods, in the platform, the orders in different time periods t ∈ [0 , 1 , 2 ..., 23] will be allocated with different allowances c ∈ N + . For example, at 10 A.M. (i.e., t = 10 ), we add 0 . 5$ (i.e., c = 0 . 5 ) allowances to all of the demanded orders. From 10 A.M. to 11 A.M., the delivery clerks who take orders and send food to customers will receive extra allowances. Specifically, if the platform pays the delivery clerks 2$ for fulfilling the order, now he/she will receive 2 . 5$ . For each day, the budget of allowance C is fixed. We should find the best budget allocation policy π ∗ ( c | t ) of the limited budget C to make as many orders as possible can be taken timely.

To find the policy, we first learn a model to reconstruct the response of allowance for each delivery clerk ˆ M ( y t +1 | s t , p t , c t ) , where y t +1 is the taken orders of the delivery clerks in state s t , c t is the allowances, p t denotes static features of the time period t . In particular, the state s t includes historical order-taken information of the delivery clerks, current orders information, the feature of weather, city information, and so on. Then we use a rule-based mapping function f to fill the complete next time-period states, i.e., s t +1 = f ( s t , p t , c t , y t +1 ) . Here we define the composition of the above functions ˆ M and f as ˆ M f . Finally, we learn a budget allocation policy based on the learned model. For each day, the policy we would like to find is:

In our experiment, we evaluate the degree of balancing between demand and supply by computing the averaged five-minute order-taken rate, that is the percentage of orders picked up within five minutes. Note that the behavior policy is fixed for the long term in this application. So we directly use the data replay with a small scale of noise (See Tab. 6) to reconstruct the behavior policy for model learning in GALILEO.

Also note that although we model the response for each delivery clerk, for fairness, the budget allocation policy is just determining the allowance of each time period t and keeps the allowance to each delivery clerk s the same.

G.2 Baseline Algorithms

The algorithm we compared are: (1) Supervised Learning (SL): training a environment model to minimize the expectation of prediction error, without considering the counterfactual risks; (2) inverse propensity weighting (IPW) [50]: a practical way to balance the selection bias by re-weighting. It can be regarded as ω = 1 ˆ µ , where ˆ µ is another model learned to approximate the behavior policy; (3) SCIGAN: a recent proposed adversarial algorithm for model learning for continuous-valued interventions [7]. All of the baselines algorithms are implemented with the same capacity of neural networks (See Tab. 6).

G.2.1 Supervised Learning (SL)

As a baseline, we train a multilayer perceptron model to directly predict the response of different treatments, without considering the counterfactual risks. We use mean square error to estimate the performance of our model so that the loss function can be expressed as MSE = 1 n ∑ n i =1 ( y i -ˆ y i ) 2 , where n is the number of samples, y is the true value of response and ˆ y is the predicted response. In practice, we train our SL models using Adam optimizer and the initial learning rate 3 e -4 on both datasets TCGA and GNFC. The architecture of the neural networks is listed in Tab. 6.

G.2.2 Inverse Propensity Weighting (IPW)

Inverse propensity weighting [50] is an approach where the treatment outcome model uses sample weights to balance the selection bias by re-weighting. The weights are defined as the inverse propensity of actually getting the treatment, which can be expressed as 1 ˆ µ ( a | x ) , where x stands for the feature vectors in a dataset, a is the corresponding action and ˆ µ ( a | x ) indicates the action taken probability of a given the features x within the dataset. ˆ µ is learned with standard supervised learning. Standard IPW leads to large weights for the points with small sampling probabilities and finally makes the learning process unstable. We solve the problem by clipping the propensity score: ˆ µ ← min(ˆ µ, 0 . 05) , which is common used in existing studies [27]. The loss function can thus be expressed as 1 n ∑ n i =1 1 ˆ µ ( a i | x i ) ( y i -ˆ y i ) 2 . The architecture of the neural networks is listed in Tab. 6.

G.2.3 SCIGAN

SCIGAN [7] is a model that uses generative adversarial networks to learn the data distribution of the counterfactual outcomes and thus generate individualized response curves. SCIGAN does not place any restrictions on the form of the treatment-does response functions and is capable of estimating patient outcomes for multiple treatments, each with an associated parameter. SCIGAN first trains a generator to generate response curves for each sample within the training dataset. The learned generator can then be used to train an inference network using standard supervised methods. For fair comparison, we increase the number of parameters for the open-source version of SCIGAN so that the SCIGAN model can have same order of magnitude of network parameters as GALILEO. In addition, we also finetune the hyperparameters (Tab. 5) of the enlarged SCIGAN to realize its full strength. We set num_dosage_samples 9 and λ = 10 .

G.3 Hyper-parameters

We list the hyper-parameter of GALILEO in Tab. 6.

G.4 Computation Resources

We use one Tesla V100 PCIe 32GB GPU and a 32-core Intel(R) Xeon(R) Gold 5118 CPU @ 2.30GHz to train all of our models.

e1_p1

e1_p0.2

e1_p0.05

e0.2_p1

e0.2_p0.2

e0.2_p0.05

e0.05_p1

e0.05_p0.2

5

10

15

20

25

30

35

root MISE

GALILEO

SL

IPW

SCIGAN

(a) GNFC

t0_bias2

t0_bias20

t1_bias2

t1_bias6

t1_bias8

t2_bias2

t2_bias6

t2_bias8

0

1

2

3

4

5

6

7

root MISE

GALILEO

SL

IPW

SCIGAN

(b) TCGA

H Additional Results

H.1 Test in Single-step Environments

The results of GNFC tasks are summarized in Fig. 13(a) and the detailed results can be found in Tab. 11. 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. IPW still perform well when 0 . 2 ≤ e ≤ 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 should be taken into consideration otherwise the algorithm will be hard to use in practice if it is sensitive to the noise range.

The results of TCGA tasks are summarized in Fig. 13(b) and the detailed results can be found in Tab. 12. We found the phenomenon in this 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 this experiment. It might be because the behavior policy is modeled with beta distribution while the propensity score ˆ µ is modeled with Gaussian distribution. Since IPW directly reweight loss function with 1 ˆ µ , the results are sensitive to the error on ˆ µ . GALILEO also models ˆ µ with Gaussian distribution but the results are more stable since GALILEO does not re-weight through ˆ µ explicitly.

We give the averaged responses for all of the tasks and the algorithms in Fig. 21 to Fig. 28. We randomly select 20% of the states in the dataset and equidistantly sample actions from the action space for each sampled state, and plot the averaged predicted feedback of each action. The real response is slightly different among different figure as the randomly-selected states for testing is different. We sample 9 points in GNFC tasks and 33 points in TAGC tasks for plotting.

H.2 All of the Result Table

We give the result of CNFC in Tab. 11, TCGA in Tab. 12, BAT in Tab. 9, and MuJoCo in Tab. 10.

H.3 Ablation Studies

In Appx. E.1, we introduce several techniques to develop a practical GALILEO algorithm. Based on task e0.2_p0.05 in GNFC, we give the ablation studies to investigate the effects of these techniques. We first compare two variants that do not handle the assumptions violation problems: (1) NO_INJECT_NOISE : set ϵ µ and ϵ D to zero, which makes the overlap assumption not satisfied;; (2) SINGLE_SL : without replacing the second term in Eq. (6) with standard supervised learning even when the output probability of D is far away from 0 . 5 . Besides, we introduced several tricks inspired by GAIL and give a comparison of these tricks and GAIL: (3) ONE_STEP : use one-step reward instead of cumulative rewards (i.e., Q and V; see Eq. (23) and Eq. (24)) for re-weighting, which is implemented by set γ to 0 ; (4) SINGE_DIS : remove T φ ∗ 1 ( x, a ) and replace it with E M θ [ T φ ∗ 0 ( x, a, x ′ ) ] , which is inspired by GAIL that uses a value function as a baseline instead of using another discriminator; (5) PURE_GAIL : remove the

GALILEO

SINGLE_DIS

ONE_STEP

SINGLE_SL

PURE_GAIL

NO_INJECT_NOISE

0

2

4

6

8

10

sqrt MISE

second term in Eq. (6). It can be regarded as a naive adoption of GAIL and a partial implementation of GALILEO.

We summarize the results in Fig. 14. Based on the results of NO_INJECT_NOISE and SINGLE_SL , we can see that handling the assumption violation problems is important and will increase the ability on counterfactual queries. The results of PURE_GAIL tell us that the partial implementation of GALILEO is not enough to give stable predictions on counterfactual data; On the other hand, the result of ONE_STEP also demonstrates that embedding the cumulative error of one-step prediction is helpful for GALILEO training; Finally, we also found that SINGLE_DIS nearly has almost no effect on the results. It suggests that, empirically, we can use E M θ [ T φ ∗ 0 ( x, a, x ′ ) ] as a replacement for T φ ∗ 1 ( x, a ) , which can reduce the computation costs of the extra discriminator training.

H.4 Worst-Case Prediction Error

In theory, GALILEO increases the generalization ability by focusing on the worst-case samples' training to achieve AWRM. To demonstrate the property, we propose a new metric named Mean-Max Square Error (MMSE): E [ max a ∈A ( M ∗ ( x ′ | x, a ) -M ( x ′ | x, a )) 2 ] and give the results of MMSE for GNFC in Tab. 13 and for TCGA in Tab. 14.

H.5 Detailed Results in the MuJoCo Tasks

We select 3 environments from D4RL [17] to construct our model learning tasks. We compare it with a typical transition model learning algorithm used in the previous offline model-based RL algorithms [59, 30], which is a variant of standard supervised learning. We name the method OFFSL. Besides, we also implement IPW and SCIGAN as the baselines. We train models in datasets HalfCheetah-medium, Walker2d-medium, and Hopper-medium, which are collected by a behavior policy with 1/3 performance to the expert policy, then we test them in the corresponding expert dataset. For better training efficiency, the trajectories in the training and testing datasets are truncated, remaining the first 100 steps . We plot the converged results and learning curves of the three MuJoCo tasks in Tab. 10 and Fig. 15 respectively.

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

0.4

0.6

0.8

1.0

1.2

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(a) halfcheetah-medium (train)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

0.05

0.10

0.15

0.20

0.25

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(b) hopper-medium (train)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

0.4

0.5

0.6

0.7

0.8

0.9

1.0

1.1

1.2

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(c) walker2d-medium (train)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

1.4

1.6

1.8

2.0

2.2

2.4

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

0.40

0.45

0.50

0.55

0.60

0.65

0.70

0.75

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

1.0

1.1

1.2

1.3

1.4

1.5

1.6

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

2.25

2.50

2.75

3.00

3.25

3.50

3.75

4.00

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(g) halfcheetah-expert (test)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

0.3

0.4

0.5

0.6

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(h) hopper-expert (test)

0.0

0.2

0.4

0.6

0.8

1.0

update steps

×10 5

1.5

1.6

1.7

1.8

1.9

sqrt MSE

GALILEO

OFF-SL

SCIGAN

IPW

(i) walker2d-expert (test)

In Fig. 15, we can see that all algorithms perform well in the training datasets. OFF-SL and SCIGAN can even reach a bit lower error in halfcheetah and walker2d. However, when we verify the models through 'expert' and 'medium-replay' datasets, which are collected by other policies, the performance of GALILEO is significantly more stable and better than all other algorithms. As the training continues, the baseline algorithms even gets worse and worse. However, whether in GALILEO or other baselines, the performance for testing is at least 2x worse than in the training

dataset, and the error is large especially in halfcheetah. The phenomenon indicates that although GALILEO can make better performances for counterfactual queries, the risks of using the models are still large and still challenging to be further solved.

We then verify the generalization ability of the learned models above by adopting them into offline model-based RL. Instead of designing sophisticated tricks to suppress policy exploration and learning in risky regions as current offline model-based RL algorithms [59, 30] do, we just use the standard SAC algorithm [20] to exploit the models for policy learning to strictly verify the ability of the models. 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 optimization, we learn and evaluate the policies with three smaller horizons: H ∈ { 10 , 20 , 40 } .

The results have been listed in Tab. 7, where the learning curve in the dynamics models and the real environments is shown in Fig. 16 and Fig. 17. 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%). At the same time, we found that SCIGAN performed better in policy learning, while IPW performed similarly to SL. This is in line with our expectations, since IPW only considers the uniform policy as the target policy for debiasing, while policy optimization requires querying a wide variety of policies. Minimizing the prediction risks only under a uniform policy cannot yield a good environment model for policy optimization. On the other hand, SCIGAN, as a partial implementation of GALILEO (refer to Appx. E.2), also roughly achieves AWRM and considers the cumulative effects of policy on the state distribution, so its overall performance is better; In addition, we find that GALILEO achieves significant improvement in 6 of the 9 tasks. 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 and fall down 4 .

H.6 Off-policy Evaluation (OPE) in the MuJoCo Tasks

H.6.1 Training and Evaluation Settings

We select 3 environments from D4RL [17] to construct our model learning tasks as Appx. H.5. To match the experiment setting in DOPE [18], here we use the whole datasets to the train GALILEO model, instead of truncated dataset in Appx. H.5, for GALILEO model training.

OPE via a learned dynamics model is straightforward, which only needs to compute the return using simulated trajectories generated by the evaluated policy under the learned dynamics model. Due to the stochasticity in the model and the policy, we estimate the return for a policy with Monte-Carlo sampling. See Alg. 4 for pseudocode, where we use γ = 0 . 995 , N = 10 , H = 1000 for all of the tasks.

H.6.2 Metrics

The metrics we use in our paper are defined as follows:

0

1

2

3

4

time-step

×10 5

6

8

10

12

14

16

18

20

22

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(a) hopper

H

= 10

0

1

2

3

4

time-step

×10 5

10

20

30

40

50

60

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(b) hopper

H

= 20

0

1

2

3

4

time-step

×10 5

20

40

60

80

100

120

140

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(c) hopper

H

= 40

0

1

2

3

4

time-step

×10 5

0

5

10

15

20

25

30

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 10

0

1

2

3

4

time-step

×10 5

0

20

40

60

80

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 20

0

1

2

3

4

time-step

×10 5

0

50

100

150

200

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 40

0

1

2

3

4

time-step

×10 5

5

0

5

10

15

20

25

30

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(g) halfcheetah

H

= 10

0

1

2

3

4

time-step

×10 5

0

20

40

60

80

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 20

0

1

2

3

4

time-step

×10 5

0

50

100

150

200

returns in models

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 40

Absolute Error The absolute error is defined as the difference between the value and estimated value of a policy:

where V π is the true value of the policy and ˆ V π is the estimated value of the policy.

Rank correlation Rank correlation measures the correlation between the ordinal rankings of the value estimates and the true values, which can be written as:

0

1

2

3

4

time-step

×10 5

0

2

4

6

8

10

12

14

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(a) hopper

H

= 10

0

1

2

3

4

time-step

×10 5

0

5

10

15

20

25

30

35

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(b) hopper

H

= 20

0

1

2

3

4

time-step

×10 5

0

10

20

30

40

50

60

70

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(c) hopper

H

= 40

0

1

2

3

4

time-step

×10 5

10

5

0

5

10

15

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 10

0

1

2

3

4

time-step

×10 5

20

0

20

40

60

80

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 20

0

1

2

3

4

time-step

×10 5

0

50

100

150

200

250

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 40

0

1

2

3

4

time-step

×10 5

3

2

1

0

1

2

3

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

(g) halfcheetah

H

= 10

0

1

2

3

4

time-step

×10 5

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 20

0

1

2

3

4

time-step

×10 5

20

10

0

10

20

30

40

50

eval returns

GALILEO

GANITE

IPW

OFF-SL

ground-truth env

H

= 40

where topk( 1 : N ) denotes the indices of the top K policies as measured by estimated values ˆ V π .

Algorithm 4 Off-policy Evaluation with GALILEO model

Require: GALILEO model ( M θ ), evaluated policy π , number of rollouts N . set of initial states

H.6.3 Detailed Results

The results of OPE on three tasks are in Tab. 8. Firstly, we can see that for all recorded rankcorrelation scores, GALILEO is significantly better than all of the baseline methods, with at least 25 %improvements. As for regret@1, although GALILEO cannot reach the best performance among all of the tasks, it is always one of the top-2 methods among the three tasks, which demonstrates the stability of GALILEO. Finally, GALILEO has the smallest value gaps except for Halfcheetah. However, for all of the top-4 methods, the value gaps in halfcheetah are around 1,200. Thus we believe that their performances on value gaps are roughly at the same level.

H.7 Detailed Results in the BAT Task

The core challenge of the environment model learning in BAT tasks is similar to the challenge in Fig. 1. Specifically, the behavior policy in BAT tasks is a human-expert policy, which will tend to increase the budget of allowance in the time periods with a lower supply of delivery clerks, otherwise will decrease the budget (Fig. 12 gives an instance of this phenomenon in the real data).

Since there is no oracle environment model for querying, we have to describe the results with other metrics.

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.90

0.92

0.94

0.96

0.98

1.00

rescaled averaged response

GALILEO

SL

(a) City-A

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

rescaled averaged response

GALILEO

SL

(b) City-B

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.88

0.90

0.92

0.94

0.96

0.98

1.00

rescaled averaged response

GALILEO

SL

(c) City-C

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

rescaled averaged response

GALILEO

SL

(d) City-D

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.92

0.94

0.96

0.98

1.00

rescaled averaged response

GALILEO

SL

(e) City-E

0.0

0.2

0.4

0.6

0.8

1.0

normalized allocated allowances

0.93

0.94

0.95

0.96

0.97

0.98

0.99

1.00

rescaled averaged response

GALILEO

SL

(f) City-F

First, we review whether the tendency of the response curve is consistent. In this application, with a larger budget of allowance, the supply will not be decreased. As can be seen in Fig. 18, the tendency of GALILEO's response is valid in 6 cities but almost all of the models of SL give opposite directions to the response. 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.

Second, we conduct randomized controlled trials (RCT) in one of the testing cities. Using the RCT samples, we can evaluate the correctness of the sort order of the model predictions via Area Under the Uplift Curve (AUUC) [6]. To plot AUUC, we first sort the RCT samples based on the predicted treatment effects. Then the cumulative treatment effects are computed by scanning the sorted sample list. If the sort order of the model predictions is better, the sample with larger treatment effects will be computed early. Then

0.0

0.2

0.4

0.6

0.8

1.0

cumulative samples (rescaled to 1)

1.0

0.5

0.0

0.5

1.0

rescaled cumulative treatment effects

random

GALILEO

SL

the area of AUUC will be larger than the one via a random sorting strategy. The result of AUUC show GALILEO gives a reasonable sorting to the RCT samples (see Fig. 19).

Finally, we search for the optimal policy via the cross-entropy method planner [21] based on the learned model. We test the online supply improvement in 6 cities. The algorithm compared is a human-expert policy, which is also the behavior policy of the offline datasets. We conduct online

A/B tests for each of the cities. For each test, we randomly split a city into two partitions, one is for deploying the optimal policy learned from the GALILEO model, and the other is as a control group, which keeps the human-expert policy as before. Before the intervention, we collect 10 days' observation data and compute the averaged five-minute order-taken rates as the baselines of the treatment and control group, named b t and b c respectively. Then we start intervention and observe the five-minute order-taken rate in the following 14 days for the two groups. The results of the treatment and control groups are y t i and y c i respectively, where i denotes the i -th day of the deployment. The percentage points of the supply improvement are computed via difference-in-difference (DID):

where T is the total days of the intervention and T = 14 in our experiments.

The results are summarized in Tab. 9. The online experiment is conducted in 14 days and the results show that the policy learned with GALILEO can make better (the supply improvements are from 0.14 to 1.63 percentage points) budget allocation than the behavior policies in all the testing cities . We give detailed results which record the supply difference between the treatment group and the control group in Fig. 20.

0

5

10

15

20

day number

0.84

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(a) City-A

0

5

10

15

20

day number

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(b) City-B

0

5

10

15

20

day number

0.965

0.970

0.975

0.980

0.985

0.990

0.995

1.000

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(c) City-C

0

5

10

15

20

day number

0.96

0.97

0.98

0.99

1.00

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(d) City-D

0

5

10

15

20

day number

0.92

0.94

0.96

0.98

1.00

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(e) City-E

0

5

10

15

20

day number

0.93

0.94

0.95

0.96

0.97

0.98

0.99

1.00

rescaled five-minute order-taken rate

treatment group (GALILEO)

control group

the day starting A/B test

(f) City-F

√

√

√

0.0

0.2

0.4

0.6

0.8

1.0

action

8

9

10

11

12

13

14

15

averaged response

pred

real

(a) t0_bias2

0.2

0.4

0.6

0.8

1.0

action

8

9

10

11

12

13

14

15

averaged response

pred

real

(b) t0_bias20

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

action

8

9

10

11

12

13

14

15

averaged response

pred

real

0.0

0.2

0.4

0.6

0.8

1.0

action

8

10

12

14

16

18

averaged response

pred

real

(d) t1_bias2

0.2

0.4

0.6

0.8

action

10

11

12

13

14

15

16

17

averaged response

pred

real

(e) t1_bias6

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

action

12

13

14

15

16

17

averaged response

pred

real

(f) t1_bias8

0.0

0.2

0.4

0.6

0.8

1.0

action

10

15

20

25

30

averaged response

pred

real

(g) t2_bias2

0.2

0.4

0.6

0.8

action

12.5

15.0

17.5

20.0

22.5

25.0

27.5

30.0

averaged response

pred

real

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

action

18

20

22

24

26

28

30

averaged response

pred

real

(i) t2_bias8

2

1

0

1

2

3

4

5

action

49

50

51

52

53

54

55

56

averaged response

pred

real

(a) e1_p1

1

0

1

2

3

4

action

50

51

52

53

54

55

averaged response

pred

real

(b) e0.2_p1

1

0

1

2

3

4

action

40

45

50

55

60

averaged response

pred

real

(c) e0.05_p1

2

1

0

1

2

3

4

5

action

49

50

51

52

53

54

55

56

averaged response

pred

real

(d) e1_p0.2

1

0

1

2

3

4

action

50

51

52

53

54

55

averaged response

pred

real

(e) e0.2_p0.2

1

0

1

2

3

4

action

35

40

45

50

55

60

averaged response

pred

real

(f) e0.05_p0.2

1

0

1

2

3

4

action

50

51

52

53

54

55

averaged response

pred

real

(g) e1_p0.05

1

0

1

2

3

4

action

44

46

48

50

52

54

56

averaged response

pred

real

(h) e0.2_p0.05

1

0

1

2

3

4

action

30

35

40

45

50

55

60

65

averaged response

pred

real

(a) t0_bias2

(b) t0_bias20

(d) t1_bias2

(e) t1_bias6

(f) t1_bias8

(g) t2_bias2

(i) t2_bias8

15

pred

pred

onse

14

real

sponse

18

bonse

18

real

lod

13

16

16

vera

D

S

D

6

12

11

10

resp

/era

6

12

14

verag

S

re:

6

14

12

9

8

ave

10

pred

>

10

7

8

real

8

0.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

action

action

action

18

esponse

16

pred

real

response

17

pred

real

sponse

17

pred

real

16

16

res

14

15

res

averaged

12

led

g

D

13

14

aver

6

15

14

10

12

13

8

11

0.0

0.2

0.4

0.6

0.8

1.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

0.2

0.3

0.4

0.5

0.6

0.7

0.8

action

action

action

30

pred

30

pred

30

pred

response

25

real

bonse

26

28

real

spons

S

28

real

resp

24

S

26

averaged

20

15

ged

vera

ave

6

20

18

22

ged

/erage

av

24

22

10

16

20

14

0.0

0.2

0.4

0.6

0.8

1.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

0.2

0.3

0.4

0.5

0.6

0.7

0.8

action

action

action

(a) e1_p1

(b) e0.2_p1

(c) e0.05_p1

(d) e1_p0.2

(e) e0.2_p0.2

(f) e0.05_p0.2

(g) e1_p0.05

(h) e0.2_p0.05

55

pred

56

pred

pred

sponse

54

real

onse

55

real

sponse

54

real

averaged

res

53

52

ged

res

0

S

54

53

ged

resp

53

52

51

6

52

averag

50

>

51

51

49

50

50

2

-1

0

action

1

2

3

4

-1

0

action

1

2

3

-1

0

action

1

2

E

sponse

55

pred

real

54

pred

real

sponse

60

pred

real

54

sponse

53

55

S

53

res

re

S

ged

average

52

51

jed

avera

6

52

51

aver

6

50

45

50

50

40

49

2

-1

0

1

2

3

4

-1

Q

1

2

3

-1

0

1

2

3

action

action

action

sponse

55

54

real

pred

response

54

pred

real

onse

65

60

pred

real

res

S

53

53

re:

0

s

55

jed

verage

52

averagec

52

6

50

45

ave

51

51

D

>

40

50

50

35

30

-1

action

1

2

3

-1

0

action

1

2

-1

0

action

1

2

3

0.0

0.2

0.4

0.6

0.8

1.0

action

8

9

10

11

12

13

14

15

averaged response

pred

real

(a) t0_bias2

0.2

0.4

0.6

0.8

1.0

action

8

9

10

11

12

13

14

15

averaged response

pred

real

(b) t0_bias20

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

action

8

9

10

11

12

13

14

15

averaged response

pred

real

0.0

0.2

0.4

0.6

0.8

1.0

action

6

8

10

12

14

16

18

averaged response

pred

real

(d) t1_bias2

0.2

0.4

0.6

0.8

action

10

11

12

13

14

15

16

17

averaged response

pred

real

(e) t1_bias6

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

action

12

13

14

15

16

17

averaged response

pred

real

(f) t1_bias8

0.0

0.2

0.4

0.6

0.8

1.0

action

10

15

20

25

30

averaged response

pred

real

(g) t2_bias2

0.2

0.4

0.6

0.8

action

12.5

15.0

17.5

20.0

22.5

25.0

27.5

30.0

averaged response

pred

real

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

action

18

20

22

24

26

28

30

averaged response

pred

real

(i) t2_bias8

2

1

0

1

2

3

4

5

action

49

50

51

52

53

54

55

56

averaged response

pred

real

(a) e1_p1

1

0

1

2

3

4

action

50

51

52

53

54

55

56

averaged response

pred

real

(b) e0.2_p1

1

0

1

2

3

4

action

50

52

54

56

58

averaged response

pred

real

(c) e0.05_p1

2

1

0

1

2

3

4

5

action

49

50

51

52

53

54

55

56

averaged response

pred

real

(d) e1_p0.2

1

0

1

2

3

4

action

50

52

54

56

58

60

62

averaged response

pred

real

(e) e0.2_p0.2

1

0

1

2

3

4

action

50

51

52

53

54

55

56

57

58

averaged response

pred

real

(f) e0.05_p0.2

1

0

1

2

3

4

action

46

48

50

52

54

averaged response

pred

real

(g) e1_p0.05

1

0

1

2

3

4

action

50

52

54

56

58

60

62

averaged response

pred

real

(h) e0.2_p0.05

1

0

1

2

3

4

action

44

46

48

50

52

54

averaged response

pred

real

(a) t0_bias2

(b) t0_bias20

(d) t1_bias2

(e) t1_bias6

(f) t1_bias8

(g) t2_bias2

(i) t2_bias8

15

15

15

response

?

14

sponse

14

sponse

14

13

13

13

12

res

12

resp

1.2

ged

11

pa

11

11

avei

6

10

6

ave

6

erag

10

の

av

erag

D

10

6

8

pred

real

8

pred

real

8

pred

real

0.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

action

action

action

18

pred

pred

sponse

17

pred

real

response

16

real

onse

16

17

real

resp

16

averaged

14

ged

res

15

15

1.2

erage

14

6

14

10

ave

13

ave

13

12

8

0.0

0.2

0.4

action

0.6

0.8

1.0

11

0.1

0.2

0.3

0.4

action

0.5

0.6

0.7

0.8

0.9

0.2

0.3

0.4

action

0.5

0.6

0.7

0.8

30

pred

sponse

30

pred

real

esponse

30

28

pred

real

sponse

28

real

25

26

res

26

averaged

resp

20

averagec

res

24

22

verage

24

15

20

18

ave

22

10

16

20

14

0.0

0.2

0.4

action

0.6

0.8

1.0

0.1

0.2

0.3

action

0.4

0.5

0.6

0.7

0.8

0.9

0.2

0.3

0.4

action

0.5

0.6

0.7

0.8

(a) e1_p1

(b) e0.2_p1

(c) e0.05_p1

(d) e1_p0.2

(e) e0.2_p0.2

(f) e0.05_p0.2

(g) e1_p0.05

(h) e0.2_p0.05

55

pred

55

pred

sponse

real

response

real

58

pred

54

54

response

real

56

res

S

53

53

averaged

52

averaged

52

ged

54

51

51

verag

6

52

50

50

av

50

49

48

49

-2

1

0

action

1

2

3

4

-1

0

action

1

2

3

-1

0

action

1

2

3

56

pred

esponse

55

54

pred

real

55

54

real

response

56

pred

real

averaged

res

52

53

51

50

p

S

6

53

52

51

50

averaged

52

54

50

49

49

48

2

-1

0

action

1

2

3

4

-1

0

action

1

2

3

-1

0

action

1

2

m

pred

response

54

real

55

pred

56

pred

53

sponse

54

real

response

55

real

54

averaged

52

averaged

res

53

52

53

52

51

51

51

50

50

50

49

-1

Q

action

1

2

3

49

-1

0

action

1

2

3

-1

0

action

2

3

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAn accurate environment dynamics model is crucial for various downstream tasks, such as counterfactual prediction, off-policy evaluation, and offline reinforcement learning. Currently, these models were learned through empirical risk minimization by step-wise fitting of historical transition data. However, we first show that, particularly in the sequential decision-making setting, this approach may catastrophically fail to predict counterfactual action effects due to the selection bias of behavior policies during data collection. To tackle this problem, the authors 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. They 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.\n\nSTRENGTHS:\n1. This paper addresses the problem of accurate environment dynamics model learning, which exhibits wide impacts in many downstream tasks, like counterfactual prediction, off-policy evaluation, and offline reinforcement learning. This is a very important and meaningful research topic.\n\n2. This is the first research on faithful dynamics model learning in sequential decision-making settings like RL, which demonstrates the novelty of this paper.\n\n3. The analysis to the challenges brought by the conventional empirical risk minimization method is deep and insightful. Especially, the authors use a vivid example to illustrate them. Based on this, the transition to the propose of adversarial weighted empirical risk minimization objective is smooth, which strongly supports the following the Generative Adversarial Offline Counterfactual environment model learning (GALILEO) method.\n\n3. The experiments of this work is sufficient and persuasive. The authors conduct experiments in two synthetic tasks, three continuous-control tasks, and a real-world application. They first verify that GALILEO can make accurate predictions on counterfactual data queried by other policies. Then, they demonstrate that the model learned by GALILEO is helpful to several downstream tasks.\n\nWEAKNESSES:\n1. It can be better if the authors can add several baselines to better validate the superiority of the proposed GALILEO method.\n\n2. The description of the method part is deep and comprehensive. But the authors can consider making it easier to understand if possible. \n\n3. Some typos need to be fixed in the future version, like the subtitle of Section. 5.3.\n\nQUESTIONS:\n1. I am little bit curious to the details of metrics and downstream tasks used in the experiments. For example, the AUUC, value gap, regret@1, and rank correlation metrics and the details behind the off-policy evaluation experiments. Could you please consider adding some corresponding explanations in the future version?\n\nLIMITATIONS:\nSee the Weaknesses and Questions above.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel method for improving the accuracy of environment dynamics models for counterfactual prediction, off-policy evaluation, and offline reinforcement learning. Currently, these models learn via empirical risk minimization (ERM), which the authors show can lead to failures in counterfactual action prediction due to selection bias during data collection. To address this, the authors introduce adversarial weighted empirical risk minimization (AWRM), where an adversarial policy weakens the model's prediction accuracy to encourage improvement. They implement this approach via an algorithm named GALILEO, which is evaluated on synthetic tasks, continuous-control tasks, and a real-world application. Results show that GALILEO can accurately predict counterfactual actions and improve several downstream tasks.\n\nSTRENGTHS:\n1. The concept of an adversarial weighted empirical risk minimization (AWRM) is a novel idea that brings together ideas from adversarial training and reinforcement learning. The use of an adversarial policy to manipulate the model's data distribution and improve its learning process is an innovative approach.\n2. The paper provides a theoretical foundation for AWRM and shows its implementation through the GALILEO algorithm. The authors also conduct a variety of tests to assess GALILEO's performance, including synthetic tasks, continuous-control tasks, and a real-world delivery platform.\n3. The paper is generally well-written and the use of figures to illustrate key concepts also enhances the clarity of the paper.\n\nWEAKNESSES:\n1. The authors didn't cover the preliminaries and related works adequately. I'm not too familiar with counterfactual modelling techniques, and think the authors didn't present enough to situate their work.\n2. Justifications for choosing the three baselines are missing. Why not choose the more recent GAIL based methods?\n3. Typo: line 320 Downstream\n4. Formatting issue: The upper margins in page 2, 7, 8 and 9 seem too small.\n\nQUESTIONS:\nPlease see above. Could you please better situate your work and justify the choice of baseline methods?\n\nLIMITATIONS:\nThe authors didn't address the limitations nor broader societal impacts in the paper. But I didn't see any ethical concerns.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the problem of environment modeling. An adversarial method is proposed that the adversarial counterparts is trained to exploit the model to generate a data distribution that weakens the model’s prediction accuracy, and then the model is trained under the adversarial data distribution with a weighted empirical risk minimization objective. Experiments are conducted on synthetic, control and real-world tasks.\n\nSTRENGTHS:\nThe paper is clearly written. The illustration examples are clear and well motivate the problem.\nApplying IPS and the surrogate optimization step is interesting.\nI also agree that conservative or pessimistic offline model-based RL methods often try to limit policy exploration, which might make it hard to obtain an accurate environment model.\n\nWEAKNESSES:\nComplexity and stability of the adversarial method might be concerned, where two discriminators should be learned.\n\nQUESTIONS:\n1. It is known that training GAN-like methods suffer from instability and often the hyperparameters should be carefully tuned. The proposed method involving training two discriminators and I am worried about the stability of the optimization step in Eq. (6) if it is easy to rise collapse problem.\n\n2. It is mentioned that in HalfCheetah, all policies just keep the cheetah standing. I think similar problems might also exist for other tasks, while they are not shown up since reward curves or numbers are compared more often without accessing how well the policies really do in these tasks. Despite this, could any other intuitive information on the performance of the environment modeling be conducted considering a specific task, instead of focusing on the rewards only? For example, in halfcheetah, which transitions are more likely to have a larger discrepancy?\n\nMinor:\n\nTitle of Section 5.3, ‘Ddownstream’ -> ‘downstream’\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a model-learning approach for counterfactual prediction (CP), off-policy evaluation (OPE), and offline reinforcement learning (ORL). The authors introduce the adversarial weighted empirical risk minimization (AWRM) objective to facilitate learning models that accurately evaluate target policies. Additionally, they present the GALILEO algorithm, a generative adversarial training method that approximates the data distribution induced by the optimal adversarial policy.\n\nSTRENGTHS:\n+ The paper effectively addresses the problem of learning accurate models for CP, OPE, and ORL, which is particularly significant in domains with costly data collection.\n+ The authors provide a comprehensive discussion on the impact of selection bias on CP, which serves as motivation for their objective, AWRM, and the GALILEO algorithm.\n\nWEAKNESSES:\n- However, it is unclear how novel is the weighted version of empirical risk minimization (ERM), compared to prior research. The main contribution lies in the adversarial aspect. Therefore, I recommend revising the introduction to emphasize and motivate the adversarial part.\n\n- The derivations appear reasonable. However, including a brief discussion about Algorithm 1 could enhance the paper's readability. \n\n- While the experimental results are generally ok, one notable limitation is that the authors only consider three tasks from the D4RL and DOPE benchmarks. In addition, the proposed approach does not outperform other methods in one of these tasks (HalfCheetah). To convincingly demonstrate the efficacy of their algorithm, the paper should include more tasks in their evaluation. Additionally, the constraint of limiting the time horizon to {10, 20, 40} is very strong and lacks proper motivation. \n\n- Furthermore, the proposed algorithm lacks comparison with state-of-the-art algorithms for the D4RL benchmark, such as the one mentioned in https://openreview.net/pdf?id=VYYf6S67pQc.\n\nQUESTIONS:\nIn Table 1, it is stated that the returns are normalized to lie between 0 and 100; however, there are negative values in the HalfCheetah column. Is this a typo?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces an adversarial training approach to model learning that improves performance for counterfactual data that may differ widely from the data used to train the model. This is particularly relevant when training from offline data and expecting the model to generalize when deployed later. This paper extensively presents an adversarially weighted empirical risk minimization objective drawing inspiration from inverse propensity scoring/weighting. A practical algorithm, GALILEO, is introduced and extensively tested across synthetic, continuous control, and real-world data. The paper clearly lays out the advantages of unbiased and accurate counterfactual models in a wide array of use cases in RL. The authors do this to contrast to major limitations of a majority of model-learning approaches that use supervised learning to perform empirical risk minimization.\n\nThis is a complete and well written paper. The development of the proposed method are clearly justified with sufficient grounding in the formal exposition of the equations. The derivation of AWRM was easy to follow with the structure put in place by the authors.\n\nThe included experiments clearly lay out the intended contribution of the proposed approach, that GALILEO provides a more accurate and counterfactually correct model. Impressively, this improved model is shown to have clear benefits for downstream performance in tasks beyond the “pre-training” task.\n\nSTRENGTHS:\nThe paper clearly lays out the advantages of unbiased and accurate counterfactual models in a wide array of use cases in RL. The authors do this to contrast to major limitations of a majority of model-learning approaches that use supervised learning to perform empirical risk minimization.\n\nThis is a complete and well written paper. The development of the proposed method are clearly justified with sufficient grounding in the formal exposition of the equations. The derivation of AWRM was easy to follow with the structure put in place by the authors.\n\nThe included experiments clearly lay out the intended contribution of the proposed approach, that GALILEO provides a more accurate and counterfactually correct model. Impressively, this improved model is shown to have clear benefits for downstream performance in tasks beyond the “pre-training” task.\n\nWEAKNESSES:\nI’m not sure I agree that having real-world offline data being biased is a catastrophic problem. Obviously for model-learning it creates major difficulties but perhaps we can focus on identifying better approaches to using expert data than expecting an RL agent to “figure it out” when there aren’t reliable demonstrations of counterfactuals. There’s likely a good reason why alternative or “exploratory” behaviors are not represented in the data? See Fatemi, et al (2021; NeurIPS) and Killian, et al (2023; TMLR) for a formulation of how we may more adequately think about risk and decision making in such environments.\n\nThe spacing between paragraphs and definitions+equations is really tight. This makes the paper difficult to read. I understand that this was likely a response to the NeurIPS template and page limitations but it should be fixed.\n\nWhile the downstream performance of MBRL methods using GALILEO models is promising. I wish that more analysis was done in the learning dynamics + performance in these downstream tasks regarding sample efficiency and deviations from the behavioral data / optimal policies in these environments. Do the GALILEO agents have predetermined action sequences that they exploit early on or are they flexible to the change in domain/task?\n\n>Fatemi, Mehdi, et al. \"Medical dead-ends and learning to identify high-risk states and treatments.\" Advances in Neural Information Processing Systems 34 (2021): 4856-4870\n\n>Killian, Taylor W., Sonali Parbhoo, and Marzyeh Ghassemi. \"Risk Sensitive Dead-end Identification in Safety-Critical Offline Reinforcement Learning.\" Transactions on Machine Learning Research (2022).\n\nQUESTIONS:\nFor the MuJoCo experiments it is unclear between Section 5.1 and 5.2 what the set-up is. It appears that the models are trained with the “medium” datasets but are evaluated on “expert” and “medium-replay” datasets. Is this correct?\n\nFor Figure 5b and 5c, what is meant by “update steps” on the horizontal axis? Is this an evaluation of model performance on the test data while training on the training data? This should be made more clear. I was at first inclined to believe that the models were being fine-tuned on the test data…\n\nTitle for Section 5.3 has typo: “Ddownstream”\n\nI was curious why Causal Curiosity (Sontakke, et al 2021ICML) wasn’t used as a baseline? The paper is included among the References but there is no mention of this work in the Related Work section. I would imagine that it would be a relevant baseline to a counterfactually driven model learning approach like GALILEO.\n\n> Sontakke, Sumedh A., et al. \"Causal curiosity: Rl agents discovering self-supervised experiments for causal representation learning.\" International conference on machine learning. PMLR, 2021.\n\nLIMITATIONS:\nAs stated by the authors, there are several simplications to the modeling process which may be a cause for the deviation in GALILEO performance when applied to downstream tasks. \n\nAdditionally, as mentioned in the “Weaknesses” section, there is an assumption that counterfactual modeling (on action sets outside the support of the dataset) is admissible. This may eliminate the use of GALILEO among safety critical domains, which would be a majority of real-world settings where model-learning could be useful.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-24"}