Causally Correct Partial Models for Reinforcement Learning
Authors: Danilo J. Rezende
Ivo Danihelka
George Papamakarios
Nan Rosemary Ke
Ray Jiang
Theophane Weber
Karol Gregor
Hamza Merzic
Fabio Viola
Jane Wang
Jovana Mitrovic
Frederic Besse
Ioannis Antonoglou
Lars Buesing
Julian Schrittwieser
Thomas Hubert
David Silver
Venue:  (2020)
Keywords: causality
model-based reinforcement learning

Abstract: In reinforcement learning, we can learn a model of future observations and rewards, and use it to plan the agent's next actions. However, jointly modeling future observations can be computationally expensive or even intractable if the observations are high-dimensional (e.g. images). For this reason, previous works have considered partial models, which model only part of the observation. In this paper, we show that partial models can be causally incorrect: they are confounded by the observations they don't model, and can therefore lead to incorrect planning. To address this, we introduce a general family of partial models that are provably causally correct, but avoid the need to fully model future observations.

Under review as a conference paper at ICLR 2020
CAUSALLY CORRECT PARTIAL MODELS FOR REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
In reinforcement learning, we can learn a model of future observations and re- wards, and use it plan the agent's next actions. However, jointly modeling future observations can be computationally expensive or even intractable if the observations are high-dimensional (e.g. images). For this reason, previous works have considered partial models, which model only part of the observation. In this paper, we show that partial models can be causally incorrect: they are confounded by the observations they don't 't model, and can therefore lead to incorrect planning. To address this, we introduce general family of partial models that are provably causally correct, yet remain fast because they do not need to fully model future observations.
INTRODUCTION
The ability predict future outcomes of hypothetical decisions is key aspect of intelligence. One approach to capture this ability is via model-based reinforcement learning (MBRL) (Munro) 1987 Werbos 1987 Nguyen & Widrow] 1990,, Scchiddhuber 1991). In this framework, an agent builds an internal representation st by sensing an environment through observational data y/ (such as rewards, visual inputs, proprioceptive information) and interacts with the environment by taking actions a according to a policy т (a))ss)) The sensory data collected is used to build a model that typically predicts future observations y>v from past actions a-c and past observations Y<r. The resulting model may be used in various ways, e.g. for planning (Oh et al. 2015 Silver et 201 generation of synthetic training data (Weher et al.) 2017), better credit assignment (Heess et al. 2015). learning useful internal representations and belief states (Gregor et al.] 2019 Guo et al.
2018 or exploration via quantification of uncertainty or information gain (Pathak et al. 2017).
Within MBRL, commonll explored methods include action-conditionall next-step models (Oh et al.,
2013 Ha & Schmidhuber 2018 Chiappa et al. 2017, S.1.d//hchiddubrrr 22010) | Xie et al. 2016 Deisen-
roth & Rasmussen 2011 Lin & Mitchell| 1992; Li al.] 2015 Diuk et al.) 2008 [gl et al.) [018]]
Ebert et al. 2018: Kaiser et al., 2019] Janner et al. 2011).. However, it is often not tractable to
accurately model all the available information. This is both due to the fact that conditioning on high-dimensional data such as images would require modeling and generating images in order to plan over several timesteps (Finn & Levine 2017). and to the f'ct that modeling images is chal- lenging and may unnecessarily focus on visual details which are not relevant for acting. These challenges have motivated researchers to consider simpler models, henceforth referred to as partial models, i.e. models which are neither conditioned on, nor generate the full set of observed data (Guo al. 2018 Gregor al. 2019 Amos et al., 2018).
In this paper, we demonstrate that partial models will often fail to make correct predictions under new policy, and link this failure to a problem in causal reasoning. Prior to this work, there has been growing interest in combining causal inference with RL research in the directions of mone based bandit algorithms (Bareinboim et al. 2015 Forney et al., 2017 Zhang & Bareinboim 2017 Lee & Bareinboim, 2018 Bradtke & Barto 1996 and causal discovery with RL (Zhu & Chen 2019). Contrary to previous works, in this paper we focus on model- based approaches and propose novel framework for learning better partial models. A key insight of our methodology is the fact that any piece of information about the state of the environment that is used by the policy to make decision, but is not available to the model, acts confounding variable for that model. As

---

Under review as a conference paper at ICLR 2020
result, the learned model is causally incorrect. Using such model to reason may lead to the wrong conclusions about the optimal course of action as we demonstrate in this paper.
We address these issues of partial models by combining general principles of causal reasoning, probabilistic modeling and deep learning. Our contributions are as follows.
We identify and clarify a fundamental problem of partial models from a causal-reasoning per- spective and illustrate it using simple, intuitive Markov Decision Processes (MDPs) (Section 2.
In order to tackle these shortcomings we examine the following question: What is the minimal information that we have to condition partial model on such that will be causally correct with respect to changes in the policy? (Section 4
We answer this question by proposing a family of viable solutions and empirically investigate their effects on models learned in illustrative environments (simple MDPs and 3D environments). Our method is described in Section 4 and the experiments are in Section S
2 A SIMPLE EXAMPLE: FUZZYBEAR
We illustrate the issues with partial models using simple example. Consider the FuzzyBear MDP shown in Figure ma): an agent initial state sp transitions into an encounter with either teddy bear or grizzly bear with 50% random chance, and can then take an action to either hug the bear or run away. In order to plan, the agent may learn a partial model go(r2|80, ag a1) that predicts the reward r2 after performing actions {ao, a starting from state sp. This model is partial because it conditions on a sequence of actions without conditioning on the intermediate state 81. The model is suitable for deterministic environments, but it will have problems stochastic environments, as we shall see. Such a reward model is usually trained on the agent's 's experience which consists of sequences of past actions and associated rewards.
Start
Start
stay home
visit forest
visit Forest
0.6
Legend
50%
50%
50%
50%
state node
Tedlyy bear
Grizzyy
Teckly bear
Grizzty bear
bear
action node
teminnl state
hug
run
bug
O
run
lug
run
hag
franstion probabilty
0
0.5
0
-0.5
0
(a)
(b)
Figure 1: Examples of stochastic MDPs. (a) FuzzyBear: after visiting forest, the agent meets either teddy bear or a grizzly bear with 50% chance and can either hug the bear or run away. (b) AvoidFuzze Bear: here, the agent has the extra option to stay home.
Now, suppose the agent wishes evaluate the sequence of actions {an visit forest, a hug} using the average reward under the model (0 ( 40, a0,a1). From Figure Ia). a), we see that the correct average reward is 0.5 x 1 0.5 x (-0.5) 0.25. However, the model has been trained on past experience in which the agent has mostly hugged the teddy bear and ran away from the grizzly bear, will learn that the sequence {visit forest, hug} is associated with a reward close to 1, and that the sequence ( (visit forest, run) is associated with reward close to O. Mathematically, the model will learn the following conditional probability:
p(s1 E. Σ. p an) (n) p(rals),
p(ra|sg. ag a,) = Σ p(s)|so, an a )(((rs)))))) = E
where 81 is the state corresponding to either teddy bear or grizzly bear. In the above expression, p(si|so ag) and p(rz|sı, a1) are the transition and reward dynamics of the MDP, and x(a (s1) is the agent's 's behavior policy that generated its past experience. As we can see, the behavior policy affects what the model learns.
The fact that the reward model (F2]sp, ag 01) is not robust to changes in the behavior policy has serious implications for planning. For example, suppose that instead of visiting the forest, the agent
2

---

Under review as a conference paper at ICLR 2020
could have chosen to stay at home as shown in Figure Tb). In this situation, the optimal action is to stay home as it gives a reward of 0.6, whereas visiting the forest gives at most a reward of 0.5x1+0.5x4
0.5. However, agent that uses the above reward model to plan will overestimate the reward of going into the forest as being close to 1 and choose the suboptimal action |
One way avoid this bias is to use behavior policy that doesn't depend the state 81, Lee. [(a](s))= = TTaa). Unfortunately, this approach does not scale well to complex environments as it requires an enormous amount of training data for the behavior policy explore interesting states. A better approach is to make the model robust to changes in the behavior policy. Fundamentally. the problem is due to causally incorrect reasoning: the model learns the observational conditional p(ra|su Su, ao 01) instead of the interventional conditional given by:
p(r2|so, do(ao), do(ai)) = Σ p(si|so, co)p(r2/11,))))
where the do-operator do(.) means that the actions performed independently of the unspecified context (i.e. independently of 81). The interventional conditional is robust to changes in the policy and is more appropriate quantity for planning. In contrast, the observational conditional quantifies the statistical association between the actions Qu, and the reward 72 regardless of whether the
actions caused the reward or the reward caused the actions. In Section 3l we review relevant concepts from causal reasoning, and based on them we propose solutions that address the problem.
Finally, although using p(r2 Sg. do(ag), do(a))) leads causally correct planning, it is not opti- mal either: it predicts reward of 0.25 for the sequence {visir forest, hug) and 0 for the sequence {visit forest, run), whereas the optimal policy obtains reward of 0.5. The optimal policy makes the decision after observing 81 (teddy bear grizzly bear); it is closed-loop as opposed to open-loop. The solution to make the intervention at the policy level instead the action level, as we discuss in the following sections.
3 BACKGROUND ON CAUSAL REASONING
u
u
ul
u
Z
X
x
y
x
X
(c)
(d)
Figure 2: Illustration of various causal graphs. (a) Simple dependence without confounding. This is the prevailing assumption many graph (b) equivalent to setting the value and observing y- (d) Graph with backdoor 2 blocking all paths from u to x. (e) Graph with a frontdoor = blocking all paths from x to y- (f) Graph with a variable 2 blocking the direct path from u y.
applications. (b) Graph with confounding. (c) Intervention on
Many applications of machine learning involve predicting variable y (target) from variable (covariate). standard way to make such prediction is by fitting a model go(y(z)) to a dataset of (1r, u)-pairs. Then, if we are given new ar and the data-generation process hasn't 't changed, we can expect that a well trained (g(g(z) will make an accurate prediction of y.
Confounding: In many situations however, we would like to use the data to make different kinds of predictions. For example, what prediction of y should we make, if something in the environment has changed, or if we set ourselves? In the latter case x didn't 't come from the original data-generation process. This may cause problems in our prediction, because there may be unobserved variables u, known confounders, that affected both and y during the data gettrratnnttionn process. That is, the actual process was of the form p(u) (()p ( (rl))(((aa u) where we only observed a and y as shown in Figure (b.. Under this assumption, a model go(ylar) fitted on (æ, y)-pairs will converge to the target p(y 1) x J p( (()))(( u ()y(y/y |r, ujdu. However, if at prediction time we set ourselves, the actual distribution of will be p(y do(x)) = ((u)p(yla ar u) dudu. This is because setting r ourselves changes the original graph from Figure 2 to the one in Figure
'This problem is not restricted to toy examples. In medical domain, model could learn that leaving the hospital increases the probahility of heing healthy.
3

---

Under review as a conference paper at ICLR 2020
The operation of setting x to fixed value Fy independently of its parents, known as the do-operator (Pearl et al. 2016). changes the lata-generation process to p(u)ō(z Io p x w), where 6(x 2o) is the bution f (p())p(yloc u)du, which we refer to as pyy do = xo)), or simply p(y]do(z) when zo is implied. Let par, be the parents of z..- The do-operator is particular case of the more general concept of an intervention: given a generative process p(æ) = II, P5 (2 (xy/par,), an intervention is defined as a change that replaces one more factors by new factors. For example, the intervention Vafak parl))
Interventions:
As explained above, this results in different target distri-
Ph (T& parg)) + thu (axx parr)) changes p(x) to p(x)
pa(E ppar)) The do-operator is a "hard" interven-
8(mk-1) tion whereby we replace node by delta function; that is, p(z do(zz u)) p(z) pa (prr,) where 2/k denotes the collection of all variables except Tk.
Backdoors and frontdoors: In general, for graphs of the form of Figure 2 b). p(yl2)) does not equal px)) As a consequence, it is not generally possible to recover p(y) dd(()) using ob- servational data, i.e. (at, y)paiirs sampled from p(ar, v). regardless of the amount of data available or the expressivity of the model. However, recovering p(y]do(r)) from observational data alone becomes possible if we assume additional structure in the cata-generation process. Suppose there exists another observed variable z that blocks all paths from the confounder to the covariate x as shown in Figure d). This variable a particular case of the concept of a backdoor (Pearl et al, 2016 Chapter 3.3) and is said to be a backdoor for the pair 2 y. In this case, we can express ply) entirely terms of distributions that can be obtained from the observational data as:
p(y]do(x)) = Epp:) 22, x:)]-
This formula holds as long as p(x|=) > 0 and is referred to as backdoor adjustment. The same formula applies when z blocks the effect of the confounder on y as in Figure More generally, we can use p(z) and p(y|z, 22, æ) from Equation (I) to compute the marginal distribution p(y) under an arbitrary intervention of the form p(xl2z) - (x(=)) on the graph in Figure (()). We refer to the new marginal as Paotw) (y) and obtain it by:
Pdo (v) = (p(a)))((s)) [p]vl=,x)]
(2)
A similar formula can be derived when there is a variable = blocking the effect of x on y, which is known as a fronidoor, shown in Figure e). Derivations for the backdoor and frontdoor adjustment formulas are provided in Appendix A
Causally correct models: Given data generated by an underlying generative process p(x). we say that a learned model (x) is causally correct with respect to a set of interventions I if the model remains accurate after any intervention in I. That is, if (o(x) Re p(x) and (0(x)) is causally correct with respect to I, then (z) = Pao(e) (Ir) for all (1(b)) in I.
Backdoor-adjustment and importance sampling: Given a dataset of N tuples (Zn, ya) gen- crated from the joint distribution p(u)p(2) (pp(2l/)p(at) (/)ppp(()))((()))) x u), we could alternatively approximate the marginal distribution (1)(w)(V) (v) after an intervention p(x[z) + (x(2)) by fitting distribution (y) to maximize the re-weighted likelihood:
1 Σ Ew(en: zn) log go (yn), N
T.(0)= = E Bettoppdtoopeeeeeeesssssoooocodd 2) log (v))
R
where w(2 =) = (xxl=)//((())))) are the importance weights. While this solution is mathematically sound way of obtaining Pao(e) (v), it requires re-fitting of the model for any new (x(z)) Moreover, if (x((2)) is very different from p(x/22) the importance weights w(x,2) 2) will have high variance. By fitting the conditional distribution p(y|2, and using Equation ( we can avoid these limitations.
As we will see in much greater detail in the next section, there is direct connection between partial models in MBRI. and the causal concepts discussed above. In MBRL we are interested in making predictions about some aspect of the future (observed frames, rewards, etc.); these would be the dependent variables y. Such predictions are conditioned on actions which play the role of the covariates x. When using partial models, the models will not have access to the full state of the policy and the policy's 's state will be confounding variable u. Any variable in the computational graph of the policy that mediates the effect of the state in the actions will be a backdoor with respect to the action action-prediction pair.
Connection to MBRL:
4

---

Under review as a conference paper at ICLR 2020
4 LEARNING CAUSALLY CORRECT MODELS
V0
90
$3
S0
$2
(a)
(b)
50
11
h2
-13
S0
90
122
h3
AU
(c)
(d)
(e)
Figure 3: Graphical representations of the environment, the agent, and the various models. Circles are stochastic nodes, rectangles are deterministic nodes. (a) Agent interacting with the environment, generating a trajectory (ye, ) These trajectories are the training data for the models. (b) Same as (a) but also including the (c) Standard autoregressive generative model observations. The model predicts the observation ye which it then feeds into h441. (d) Example of a Mon-Causal Partial Model (NCPM) that predicts the observation without feeding it into fx41. (e) Proposed Causal Partial Model (CPM), with backdoor e, for the actions.
backdoor z₁ in the generated trajectory. The red arrows indicate the locations of the interventions.
We consider environments with a hidden state e and dynamics specified by an unknown transition each step 4, the environment receives an action 4;-1,
probability of the form ((eleet1,,(-1) updates its state to e and produces observable data y/2 p(yele;) which includes reward and potentially other forms of data such as images. An agent with internal state s, interacts with the environment via actions a produced by policy (s) and updates its state using the observations V1+1 by 8t+1 = (St, dt, V1+1), where fs can for instance be implemented with an RNN. The agent will neither observe model the environment state e,; it is a confounder on the data generation process. Figure (]a) illustrates the interaction between the agent and the environment.
Consider an agent at an arbitrary point in time and whose current state is 8o: and assume are interested in generative models that can predict the outcome yr of a sequence of actions (ao...., - on the environment, for an arbitrary time T. A first approach, shown in Figure 3c), would be to use an action-conditional autoregressive model of observations; initializing the model state h a function of (So, on). sample y1 from p(.[hı} update the state h2 = fh1 a1, 31), sam- ple y2 from hb2). and so on until is sampled. In other words, the prediction of observation yr conditioned on all available observations (Sg- F/T)) and actions dşr. This approach for instance found in (Oh et al. 2015).
In contrast, another approach is to predict observation yr given actions but using no observation data beyond sg. This family of models, sometimes called models with overshoot, can for instance be found in (Silver et al. [2017b] Oh et al., 2017 Luo et al. 2019) Guo et al. 2018 Hafner et al. 2018 Gregor et al.) 2019; Asadi et al. 2019) and is illustrated in Figure 3 d). The model deterministically updates its state h+41 = fr (h,, a,), and generates yr from (br). An advantage of those models is that they generate yr directly without generating intermediate observations.
More generally, we define a partial view Vi as any function of past observations y<r and actions ast- We define a partial model partial views Uet and the actions aex: to generate yr, the agent generates e1 from p(-121) updates the state to h2 fix (hi, U1, a1), and on, until it has computed for and sampled yr from p( p(. ). Both previous examples can be seen as special cases of a partial model, with v, = and U2 Ø respectively.
generative model whose predictions are only conditioned 8p, the
We reindex time for notational simplicity, but recall that so is indeed function of past observations ygo- 'For full generality, it may be that the predicted observation is only a subset or simple function of the full observation yr; for instance one could predict only future rewards. For notational simplicity we make no difference between the full observation and the prediction.
5

---

Under review as a conference paper at ICLR 2020
Table 1: Comparison between noncausal partial model and the proposed architecture. The shaded cells indicate the key differences in architectures.
NCPM architecture (overshoot)
CPM architecture
Backdoor Action
St.
(2(2)886)
Agent
Action
Gt ~ x(aa (s))
at
I (asl21)
State Update
824-1 RNN.( ae ve-1)
State Update
8141 RNN, Ct, M641)
State Init.
h1 g(50. ao)
State Init.
h1 g(80.00)
Model
Backdoor
Rt "p(2zl(te)
State Update
h++1= RNN, Vn p(villt))
State Update
h4+1
..RNNA (heg 20, az)
Prediction
Prediction
V
plute (ha)
A subtle consequence of conditioning the model only on partial view it is that the variables yer become confounders for predicting yr. in addition the state of the environment which is always a confounder. In Section 3 we showed that the presence of confounders makes it impossible to correctly predict the target distribution after changes in the covariate distribution. In the context of partial models, the covariates are the actions acr executed by the agent and the agent's initial state So, whereas the targets are the predictions yr we want to make time T. A corollary of this that the learned partial model may not be robust against changes in the policy and thus cannot be used to make predictions under different policies x, and therefore should not be used for planning.
In Section we saw that if there was variable blocking the influence of the confounders on the covariates (a backdoor) or variable blocking the influence of the covariates on the targets (a front- door), it may be possible to make predictions under broad range of interventions if we learn the correct components from data, e.g. using the backdoor-adgustmeee formula in Equation ( In gen- formula because may not
eral it may not be straightforward to apply the packdoor-adjustmen have enough access to the graph details to know which variable is a backdoor. In reinforcement learning however, we can fully control the agent's graph. This means that we can choose any node in the agent's 's computational graph that is between ils internal state 8 and the produced action as a backdoor variable for the actions. Given the backdoor 24, the action ag is conditionally independent of the agent state St.
To make partial models causally correct, we propose choose the partial view Us to be equal to the backdoor 2t This allows us to learn all components we need to make predictions under an arbitrary new policy. In the rest of this paper we will refer to such models as Causal Partial Models (CPM), and all other partial models will be henceforth referred to as Non- Von-Causal Partial Models (NCPM). We assume the backdoor zt sampled from a distribution m=2((08) and the policy is a distribution conditioned on 21, [aal(z)) This is illustrated in Figure (b)) and described in more details in Table rightt) We can perform a simulation under a new policy (aa 24, by directly applying the backdoor-add backdoor-adjustmen adjustment formula, Equation ( to the RL graph as follows:
Paolyi(a, (Bt+1 [/,) = Be(e)hhes)p((s(h)) (p( [p(14+1 [1++)),,
(4)
where the components p(za (2) and p(u (h41)) with h441 (h,, 21, 0z) can be learned from observational data produced by the agent.
Ha & Schmidhu- ber (2018) have complex graphs, which means that there are many possible choices for the backdoor 21. So important question is: what are the simplest choices of 2;? Below we list few of the simplest choices we can use and discuss their advantages and trade-offs; more choices for 2 are listed in Appendix
Modern deep-learning agents (e.g. as in Essphoott et al. 2018); Gregor et al.
(2019); al.
Agent state: Identifying 2n with the agent's state 8 can be very informative about the future, but this comes at cost. As part of the generative model, we have to learn the component p(2)) (x). This may be difficult in practice when 2, = s, due to the high-dimensionality of st, hence and performing simulations would be computationally expensive.

---

Under review as a conference paper at ICLR 2020
Policy probabilities: The 2 can be the vector of probabilities produced by policy when we have discrete actions. The vector of probabilities is informative about the underlying state, different states produce different probabilities.
Intended action: The 2, can be the intended action before using some form of exploration, e.g. s- greedy exploration. This is an interesting choice when the actions are discrete, it is simple to model and, when doing planning, results in a low branching factor which is independent of the complexity of the environment (e.g. in 3D, visually rich environments).
The causal correction methods presented in this section can be applied to any partial model. In our experiments, we will focus environment models of the form proposed by [Grgorr al. (2019). These models consist of a deterministic "backbone" RNN that integrates actions and other contextual information. The states of this RNN are then used to condition generative model of the observed data ye, but the observations are not fed back to the model autoregressively, as shown in Table Ileft). This corresponds to learning a model of the form p(ye|so, .0.....a-11).
We will compare this against our proposed model, which allows us to simulate the outcome of any policy using Equation (4. this setup, a policy network produces 2, before an action at- For example, if the 2, is the intended action before e-exploration, m(2e(()) and the executed action ag will then be sampled from an e-exploration policy π (a,]=;) = (1 - =))6, + E LL, where na is the number of actions and is in (0, 1). Acting with the sampled actions is diagrammed in Figure 3b) and the mathematical description is provided in Table
will be sampled from policy
The model components p(=+| lie) and p(ye (hy) are trained via maximum likelihood on observational data collected by the agent. The partial model does not need to model all parts of the y/ observation. For example, a model to be used for planning can model just the reward and the expected return. In any case, it is imperative that we some form of exploration to ensure that <(a, (2)) 0 for all a, and 2/ as this is a necessary to allow the model to learn the effects of the actions. The model usage is summarized in Algorithms [ and 2 in Appendix D and discuss the model properties in Appendix E
EXPERIMENTS
We analyse the effect of the proposed corrections on variety of models and environments. When the enviroment is an MDP, such as the Fuzzy Bear MDP from Section we can compute exactly both the non-causal and the causal model directly from the MDP transition matrix and the behavior policy. In 3.1 we compare the optimal policies computed from the non-causal and the causal model via value iteration. For this analysis. we used the intended-action backdoor, since it's compatible with a tabular representation. In Section 5.2 we repeat the analysis using a learned model instead. For these experiments,
used the poobey-proaabilities backdoor. The optimal policies corresponding to given model were computed using a variant of the Dyna algorithm (Sutton 1991) or expectimax
(Michic 1966). Finally in Section 5.3 we provide an analysis of the model rollouts in visually rich 3D environment.
5.1 VALUE-ITERATIU ANALYSIS ON MDPs
Given an MDP and behavior policy π., the optimal values VMrG)) of planning based on NCPM and CPM are derived in Appendix[] The theoretical analysis of the MDP does not use empirically trained models from the policy data, but rather assumes that the transition probabilities of the MDP and the policy from which training data are collected are accurately learned by the model. This allows us isolate the quality of planning using the model from how accurate the model is.
Optimal behavior policy: The optimal policy of the FuzzyBear MDP (Figure [a)) is to always hug the teddy bear and run away from the grizzly bear. Using training data from this behavior policy, we show in Figure 7 (Appendix the difference in the optimal planning based on the NCPM (Figure (()) and CPM with the backdoor 2, being the intended action (Figure (()).. Learning from optimal policies with E e-exploration, the converged causal model is independent of the exploration parameter We effects of varying learned models in Figure & (Appendix[]
7

---

Under review as a conference paper at ICLR 2020
1.0
1.0
0.8
0.9
0.6
Vile
0.4
0100
0.7
-10g
0.2
0.6
0.0
250
0.5
Non-Causa Model
Model
-0.2
Gausal Model
-0.2
-0.1
0.0
0.0
Vm
(a)
(b)
Figure 4: MDP Analysis: (a) In the FuzzyBear environment, we randomly generate 500 policies and seatter plot them with x-axis showing the quality of the hehavior policy Verw and y-axis showing corresponding model optimal evaluations Vir For each policy, we derive the corresponding converged model M(x) equivalent to training on data generated by the policy. We then compute the optimal evaluation (M))) using this model. contrast unrealistic optimism of non non-causal model evaluations VNCPM(<> with the more realistic causal model evaluations VCPM(s) for good policies as well as the over-pessimism of the non-causal model compared the causal model for bad policies. Same plot as (a) but for the AvoidFuzzyBear environment.
Sub-optimal behavior policies: We empirically show the difference between the causal and non- causal models when learning from randomly generated policies. For each policy, we derive the corresponding converged model (x) using training data generated by the policy. We then compute the optimal value of (())) using this model. On FuzzyBear (Figure ((a), we see that the causal model always produces a value greater than or equal to the value of the behavior policy. The value estimated by the causal model always be achieved in the real environment. If the behavior policy was already good, the simulation policy used inside the model can reproduce the behavior policy by respecting the intended action. If the behavior policy is random, the intended action is uninformative about the underlying state, so the simulation policy has to choose the most rewarding action, independently of the state. And if the behavior policy is bad, the simulation policy can choose the opposite of the intended action. This allows to find a very good simulation policy, when the behavior policy is very bad. To further improve the policy, the search for better policies should be done also in state 81. And the model can then be retrained on data from the improved policies.
If look the non-causal model, we see that displays the unfortunate property of becoming more unrealistically optimistic as the behavior policy becomes better. Similarly, the worse the pol- icy is, i.e. the lower V is, the non-causal model becomes less able improve the policy. On AvoidFuzzyBear (Figure 4(b),, ((b) the optimal policy is to stay at home. Learning from data generated by random policies, the causal model indeed always prefers to stay home with any sampled inten- tions, resulting in constant evaluation for all policies. On the other hand, the non-causal model gives varied, overly-optimiste evaluations, while choosing the wrong action (visit forest).
5.2 LEARNED MODELS ON MDPs
We previously analyzed the case where the transition probabilities and theoretically optimal policy are known. We will now describe experiments with learned models trained by gradient descent, using the same training setup as described in Section 4
BuzeyFuzzyBear with Dyna: In this experiment we demonstrate that we can learn the optimal policy purely from off-policy experience using general n-step-return algorithm derived from causal model. The algorithm is described detail in Appendix F In short, we simulate experiences from the partial model, and use policy gradient learn the optimal policy these experiences if they were real experiences (this is possible since the policy gradient only needs action probabilities, values, predicted rewards and ends of episodes). We compare non-causal model and causal
8

---

Under review as a conference paper at ICLR 2020
model where the backdoor 2 is the intended action. For the environment use AvoidFuzzyBeai (Figure [b)). We collect experiences that are sub-optimal: half the time the agent visits the forest and half the time it stays home, but once in the forest it acts optimally with probability 0.9. This meant to simulate situations either where the agent has not yet learned the optimal policy but is acting reasonably,
where it is acting with different objective (such as exploration intrinsic reward), but would like to derive the optimal policy. We expect the non-causal model to choose the sub-optimal policy of visiting the forest, since the sequence of actions of visiting the forest and hugging typically yields high reward.
This is what we indeed find, as shown in Figure ((a) We see that the non-causal model indeed achieves a sub-optimal reward (less than 0.6), but believes that it will achieve a high reward (more than 0.6). On the other hand, the causal model achieves the optimal reward and correctly predicts that it will achieve the corresponding value.
0.8
0.6
Search depth
Model
3
reward non-causal
reward causal
Non-causal
/
x
0.2
value cousuausal
Iutended action
4
value causal
Clustered probs
V
V
500 1000 1500 2000 2500 3000 3500 4000 training iterations
Clustered abs
V /
(a)
(b)
Figure 5: (a) Dyna on AvoidFuzzyBear. Models were trained sub-optimal behavior policy that explores both parts of the environment, and evaluation policy was trained purely inside the model. The non-causal model (dotted lines) achieved sub-optimal reward, while expecting large reward. The causal model (solid lines) achieved optimal reward and had correct expectation of the reward. The shaded area indicates 95% confidence intervals from 50 runs. (b) Models solving the AvoidFuzzy Bear MDP with expectimax. The non-causal model misled the agent when using search depth 3 higher.
AvoidFuzzyBear with Expectimax: In this experiment, we used the classical expectimax search (Michie 1966. Russell & Norvig 2009). On the simple wwoidFuzyBBent MDP, it is enough to use a search depth 3: a decision node, a chance node and a decision node. The behavior policy was progressively improving as the model was trained.
In ((b) we see the results for the different models. Only the non-causal model was not able to solve the task. Planning with the non-causal model consistently preferred the stochastic path with the fuzzy bear, as predicted by our thcoretical analysis with value iteration. The models with clustered probabilities and clustered observations approximate modeling of the probabilities or observations. These models are described in Appendix
5.3 VISUALLY RICH 3D ENVIRONMENT
The setup for these experiments is similar to Gregor et al. (2019). where the agent is trained using the IMPALA algorithm (Espeholt et al., 2018), and the model trained alongside the agent via ELBO optimization on the data collected by the agent. The architecture of the agent and model is based on Gregor et al. (2019) and follows the description in Table rright). For these experiments, the backdoor 2, was chosen to be the policy probabilities, and p(2, h+) was parametrized as a mixture of Dirichlet distributions. See Appendix I for more details. We demonstrate the effect of the causal correction on the 3D T-Maze environment where an agent walks around in a 3D world with the goal of collecting the reward blocks (food). The layout of this environment is shown in Figure From our previous results, we expect NCPMs to be unrealistically optimistic. This is indeed what we see in Figure b)). Compared to NCPM. CPM with generated x generates food at the end of a rollout with around 50% chance, expected given that the environment randomly places the food on either side. In Figure []c) and Figure dd. left) we show subsets of rollouts from NCPM and CPM respectively (see Figures 10 12 in Appendix for full rollouts).

---

Under review as a conference paper at ICLR 2020
Environment agent
Probabilitty generate food
1.0
action forced
NCPM, forced action
action gen.2
(reward)
0.8
colors
0.6
Rollouts
0.4
0.2
flood colors
NCPM
CPM
(a)
(b)
(c)
CPM, forced action generated
CPM, forced action + forced
Pallouts
Rollous
Time
(d)
Figure 6: Causal partial model (CPM) and Non-causal partial model (NCPM) rollouts in 3D T-Maze environ- ment. (a) The agent always begins in walled-off corridor, and obtain food reward that spawns randomly on either the left (red) or the right (black) side. The colors of the corridor and side walls are randomized every episode, (b) Probability of the model generating food rollouts for NCPM and CPM. Error bars represent standard eгror over runs, each with 30 episodes. (c)-(d Subset of frames from example rollouts using (c) NCPM and (d) CPM. In all rollouts depicted, the top shows the real frames observed by agent following fixed policy (Ground Truth, GT). Bottom rows indicate model rollouts, conditioned on 3 previous frames without revealing the location of the food. CPM and NCPM differ their state-update formula and action generation (see T. but frame generation y ~ p(y /lx) is the same for both, as introduced in Gregor et al. (2019). For CPM, we compare rollouts with forced actions and generated to rollouts with forced actions and forced from the ground truth. We can observe that rollouts with the generated (left) respect the randomness the food placement (with and without food), while the rollouts with forced (right) consistently generate food blocks, following actions consistent with the backdoor from the well-trainch ground truth policy.
CONCLUSION
We have characterized and explained some of the issues of partial models in terms of causal reason- ing. We proposed simple, yet effective, modification to partial models so that they still make correct predictions under changes in the behavior policy, which we validated theoretically and exper- imentally. The proposed modifications address the correctness of the model against policy changes, but don't address the comrecnees/rohustness against other types of intervention in the environment. We will explore these aspects future work.
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Brandon Amos, Laurent Dinh, Serkan Cabi, Thomas Rothörl, Sergio Gómez Colmenarejo, Alistair Muldal, Tom Erez. Yuval Tassa, Nando de Freitas, and Misha Denil. Learning awareness models. arXiv preprint SEiv:1880400611, 2018.
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michel L. Littman. Combating the compounding-errve problem with multi-step model. arXiv preprint arXiv:1905.13320, 2019.
Elias Bareinboim, Andrew Forney, and Judea Pearl. Bandits with unobserved confounders: causal approach. In Advances in Neural Information Processing Systems 28, pp. 1342-1350, 2015.
Steven J Bradtke and Andrew G Barto. Linear least-squares algorithms for temporal difference learning. Machine learning, 22(1-3):33-57, 1996.
Silvia Chiappa, Sébastien Racanière, Daan Wierstra, and Shakir Mohamed. Recurrent environment simulators. arXiv preprint 2rXiv1774.002554 2017.
Mare P. Deisenroth and Carl E. Rasmussen. PILCO: A model- and data-efficient approach to policy search. In International Conference of Machine Learning, 2011.
Carlos Diuk, Andre Cohen, and Michael L. Littman. An object-oriented representation for efficient reinforcement learning. In Proceedings of the 25th international conference Machine learning. pp. 240-247, 2008.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual fore- sight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv:/8/2.00568, 2018.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, IMPALA: Scalable distributed deep-RL with importance weighted learner architectures. arXiv preprint arX::
2018.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In IEEE Inter- national Conference on Robotics and Automation, pp. 7786--793, 2017.
Andrew Forney, Judea Pearl, and Elias Bareinboim. Counterfactual data-fusion for online reinforce- ment learners. In Proceedings of the 34th International Conference Machine Learning, pp-
1156-1164, 2017.
Karol Gregor, Danilo Jimenez Rezende, Frédéric Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for RL. arXiv preprint 1rXiv:1906.09237, 2019.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A. Pires, Toby Pohlen, and Rémi Munos. Neural predictive belief representations. arXiv preprint 1rXiv:1811.0644 2018.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1803 101222003..2...... 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint arXiv:1811.00 1.04551, 2018.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learn ing continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems, pp. 2944-2952, 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Maximilian Igl. Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for POMDPs. arXiv preprint A.06:/800.022 2018.
Michael Janner. Justin Fu. Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in Neural Information Processing Systems, 2019.

---

Under review as a conference paper at ICLR 2020
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv:1903.0000 2019.
Sanghack Lee and Elias Bareinboim. Structural causal bandits: Where to intervene? In Advances in Neural Information Processing Systems 31, pp- 2568-2578 2018.
Xiujun Li, Lihong Li, Jianfeng Gao, Xiaodong He, Jianshu Chen, Li Deng, and Ji He. Recurrent reinforcement learning: hybrid approach. arXiv preprint arXiv: 1509033444 2015.
Long Lin and Tom Mitchell. Memory approaches to reinforcement learning in non-Markovian domains. Technical report, Carnegie Mellon University, 1992.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. In Interna- tional Conference on Learning Representations, 2019.
Donald Michie. Game-playing and game-learingg automata. In Advances programming and non-numerical computation, pp. 183-200. Elsevier, 1966.
Paul Munro. dual back-propagation scheme for scalar reward learning. 9th Annual Conference of the Cognitive Science Society, pp. 165-176, 1987.
Derrick Nguyen and Bernard Widrow. The truck backer-upper: example of self-learning in neural networks. In Advanced neural computers, pp- 1 11-19. Elsevier, 1990.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard I.. Lewis, and Satinder Singh. cenion-conditional video prediction using deep networks in atari games. In Advances in neural information process- ing systems, pp. 2863-2871, 2015.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In Advances in Neural Information Processing Systems, pp. 6118-6128, 2017.
Deepak Pathak, Pulkit Agrawal, Alexci A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 16-17, 2017.
Judea Pearl, Madelyn Glymour, and Nicholas Jewell. Causal inference in statistics: A primer.
John Wiley & Sons, 2016.
Stuart Russell and Peter Norvig, Artificial Intelligence: A Modern Approach. Prentice Hall Press, Upper Saddle River, NJ, USA, 3rd edition, 2009. ISBN 0136042597, 9780136042594.
Jürgen Schmidhuber. Curious model model-building control systems. In Proceedings 1991 IEEE Interna- tional Joint Conference on Neural Networks, pp. 1458-1463, 1991.
Jürgen Schmidhuber. Formal theory of creativity, fun, and intrinsic motivation (1990-2010). IEEE Transactions on Autonomous Mental Development, 2(3):230-244 247, 2010.
David Silver, Richard S. Sutton, and Martin Müller. Sample-based learning and search with per- manent and transient memories. In Proceedings of the 25th international conference on Machine learning, pp. 968 975, 2008
David Silver, Julian Schritwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang. Arthur Guez, Thomas Hubert, Lawrence R. Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy P. Lill- icrap, Fan Fong Celine Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 150:354-359 2017a.
David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Arold, David P. Reichert, Neil C. Rabinowitz, André Barreto, and Thomas Degris. The predictron: End learning and planning. In International Conference on Machine Learning, 2017b.
12

---

Under review as a conference paper at ICLR 2020
Richard S. Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. In Machine Learning Proceedings 1990, pp. 216-224. Elsevier, 1990.
Richard Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-163, 1991.
Théophane Weber, Sébastien Racanière, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al. Imagination-augmentee agents for deep reinforcement learning. arXiv preprint arXiv: 1707.06203, 2017.
Paul J. Werbos. Learning how the world works: Specifications for predictive networks in robots and brains. In Proceedings of IEEE International Conference on Systems, Man and Cybernetics, NY, 1987.
Chris Xie, Sachin Patil, Teodor Moldovan, Sergey Levine, and Pieter Abbeel. Model-based rein- forcement learning with parametrized physical models and optimism-driven exploration. In 2016 IEEE international conference on robotics and automation, pp. 504-511, 2016.
Junzhe Zhang and Elias Bareinboim. Transfer learning in multi-armed bandits: A causal approach. In Proceedings of the 26th International Joint Conference on Artificial Intelligence, pp. 1340- 1346, 2017.
Shengyu Zhu and Zhitang Chen. Causal discovery with reinforcement learning. arXiv preprint arXi:1906. 04477v, 2019.
13

---

Under review as a conference paper at ICLR 2020
A BACKDOOR AND FRONTDOOR ADJUSTMENT FORMULAS
Starting from a data-generation process of the form illustrated in Figure 2 p(z,y.u) p(u)p(z(uup(())z u), we can use the do-operator to compute p(y do(x)) = ( x u)du.
Without assuming any extra structure in p(r/u) in p(y|z, it is not possible compute p(y]do(x)) from the knowledge the joint p(a, y) alone.
If there was variable blocking all the effects of on as illustrated in Figure dd), then p(y)do(x)) can be derived as follows:
Joint density Intervention
((u)p(//uppplllpppol u)
(5)
p(z|2) + (x) ptupp(elwww@)p)ppp..
(6)
Joint after intervention Conditioning the new joint
[ / u)
plyld(())
(8)
(xx)
f p(=))(ulae
= Epccspllk
(10)
where we used the formula
f pluuplelw)plole.uup u)du p(z) /
p(ulz)p(ulz, u)du
(11)
= =p(2)p(u/2z,
(12)
If instead of just fixing the value of x. we perform a more general intervention p(z|2) + (x)22)
then (do(p(e(=)) (y) can be derived follows:
Joint density Intervention
pfubu(ala)p(alssapeee u)
(13)
p(x - (x]=))
(14)
Joint after intervention
wwwpptelwwWaaebpppee.... u)
(15)
/ p(u)p(elaa)(elee)oooo u)dudzdr
New marginal
(oo()(ale))
(16)
/ (p))() =)dadr
(17)
Epte)ocats (y =)].
(18)
Applying the same reasoning to the graph shown in Figure e), we obtain the formula
ply)do(x)) = E [epcete)(((sddd(
pppele)p(e (p(y [/],
(19)
where p(z[z), p(x') and p(y x², 22) can be directly measured from the available (x, y, 2) data. This
formula holds long p( (2)2:) 0, Vi,, and it is simple instance of frontdoor adjustment (Pearl et al. 2016).
B DERIVATION OF CAUSAL CORRECTNESS FOR THE MODELS IN FIGURE 3
Here, we will show in more detail that the models (c) and (e) in Figure are causally correct, whereas model (d) is causally incorrect. Specifically, we will show that given an initial state sy and after setting the actions ag ar to specific values, models (c) and (e) make the same prediction about the future observation yr+1 as performing the intervention in the real world, whereas model (d) does not.
14

---

Under review as a conference paper at ICLR 2020
Model (c) Using the do-operator, hard intervention in the model is given by:
go(uT+ 1800 do(ao:)) go (yT441550 (g:7)
II go (y(e) dyn.r,
(20)
where ht is deterministic function of so, do.r-1 and y1a--1 The same hard intervention in the real world is given by:
/ p(yI:T+1 do(ao:)) dyn.r
P(UT+1 1100 do(aa:)))
(21)
1I pl so. do(an:))) V1--11 dy
(22)
T+1
II
plys so. a0:t-1 V1:1 1)) dyi.r.
(23)
If the model is trained perfectly, the factors qo ye h2) will become equal the conditionals p(ye so; d0a--1, M1a-1). 1). Hence, intervention in a perfectly trained model makes the same pre- diction in the real world, which means that the model causally correct.
Model (d) The interventional conditional in the model is simply:
go ( do(ao.r)) qo (yT++ a0.T) (o(U++ [1T+1);
(24)
where h141 is a deterministic function of and du:r- In a perfectly trained model, we have that go (o(y==1/hrz1) p dp:r). However, the observational conditional p(yr+ +1 (so ou:r) is
not generally equal to the inverventional conditional p(yr+1 (Sp. do(o)))) which means that the model causally incorrect.
Model (e) Finally, the interventional conditional in this model is:
I
qo (y(441 So, )) 4o (gr+1|$u, ag:r) go( (144)) II
go d2i.
(25)
where h2 is deterministic function of so, au:-11 and 21--1- The same intervention in the real world can be written as follows:
= /
p(yr+1 (s. do(auz))
p plur+ilsu, do(au:)), 21.1 )p( su, do(ao:)) d21:r
(26)
T
/ plur+ilsu, do(au:)) 21.1)
II
pl 2 ss, do(ao:)), =14-1) d2i.r (27)
f plurT+lse ao.r, ) II p(z+ So. ao.-1. 314-1) džy.r.
T
(28)
In perfectly trained model, we have that go (yT [bT+1) ) = p(ur+1 do:T; 21:T) and go (h)) = p(2e|50, So, 00:t-1, 21:t-1). That means that the intervention in perfectly trained model makes the same prediction as the same intervention in the real world, hence the model is causally correct.
C ADDITIONAL CHOICES OF THE BACKDOOR 2t
The first alternative backdoor consider is the empty backdoor:
Empty backdoor 2, = Ø: This backdoor is in general not appropriate; it is however appropriate when the behavior policy does in fact depend no information, i.e. is not function of the state
15

---

Under review as a conference paper at ICLR 2020
St- For example, the policy can be uniformly random (or any non-state dependent distribution actions). This severely limits the behavior policy. Because the backdoor contains no information about the observations, the simulations open-loop. ie. can only consider plans which consist of a sequence of fixed actions, not policies.
be any layer from the policy. To model the layer with a p(2,(h,) distribution, we would need to know the needed numerical precision of the considered layer. For example, quantized layer can be modeled by discrete distribution. Alternatively, if the layer is produced by a variational encoder or variational information bottleneck, we can train p(zy(he) to minimize the [L((encoder (z st) hx)).
An intermediate layer: In principle, the
Finally, if backdoor is appropriate, we can combine it with additional information:
Combinations: It is possible combine layer with information from other layers. For example, the intended action can be combined with extra bits from the input layer. Such can be more informative. For example, the extra bits can hold downsampled and quantized version of the input layer.
D ALGORITHMS FOR TRAINING AND SIMULATING FROM THE MODEL
Algorithms ! and describe how the model is trained and used simulate trajectories. The algo- rithm for training assumes distributed actor-learner setup (Espeholt et al. 2018).
Algorithm Model training
Data collection on an actor: For each step:
2t ~ m s... the backdoor (e.g., the partial view with the intended action) dg ~ π (a, zt) ... sample the executed action (e.g., add e-exploration)
Collect:
8t ... agent state
2t partial view
..... executed action
3i+1 targets (rewards, returns,...
Model training on a learner:
Require trajectory: Su, aer, Zer, VeT
h1 = g(So, ....aan...iiiiiiiiiii. initiallze the model state For each trajectory step:
Train p(yel (b,) to model yr-
Train p(2; h2) to model žt.
h++1 RNN, (ha, 24, a+) update the model state
Algorithm Using the model to generate a simulation under a new policy v
Require an agent state: Sp
an vi (o)) ... choose the first action
h; = g(0,, do) initialize the model state
For each trajectory step:
Predict the wanted targets p (h)) (e.g., rewards, returns,
2, p(z, (h)) ... sample the partial view
0g ~ (a, h4, 21) choose the next action
h411 RNN,, a4) update the model state
E DISCUSSION OF MODEL PROPERTIES
Jublei[ an overview of properties of autoregressive models, deterministic non-causal mod- els and the causal partial models. The causal partial models have generate only partial view. The
16

---

Under review as a conference paper at ICLR 2020
partial view can be small and easy to model. For example, partial view with the discrete intended action can be flexibly modeled by a categorical distribution.
The causal partial models are fast and causally correct in stochastic environments. The causal partial models have a low simulation variance, because they do not need to model and generate unimportant background distractors. If the environment has deterministic regions, the model can quickly learn to ignore the small partial view and collect information only from the executed action.
It is interesting that the causal partial models are invariant of the TTa (2) distribution. For example, if the partial view Σạ is the intended action, the optimally learned model would be invariant of the used s-exploration: x [z)) Analogously, the autoregressive models are invariant of the whole pol- icy π (ar (st). This allows the autoregressive models to evaluate any other policy inside of the model. The causal partial model can run inside the simulation only policies conditioned on the starting state So, the actions det and the partial views zge- If we want to evaluate policy conditioned on different features, we can collect trajectories from the policy and retrain the model. The model can always evaluate the policy used to produce the training data. We can also improve the policy, because the model allows to estimate the return for an initial (so, ao) pair, so the model can used as a critic for policy improvement.
Table 2: Models and their properties.
Autoregressive
Deterministic
Causal Partial Model
Generates
observation
nothing
partial view: 21
Speed Causally correct
slow
fast
fast
always
in deterministic environments or with on polpolic simulations
always
Simulation variance
high (distracted)
lowest
low
Extra branching
huge
0
controlled by 21 size (
Invariant of
π (ay]s)
Evaluable policies Training
any
(0 80, aet))
1/ 04 80, act, Fcr)
once
iterative with policy
iterative with policy
F DYNA-STYLE POLICY-GRADIENT ALGORITHM
In this section we derive an algorithm for learning an optimal policy given (non-optimal) expe- rience that utilizes n-step returns from partial models presented in this paper. In general, model of the environment can be used in a number of ways for reinforcement learning. In Dyna (Sutton) 1990), we sample experiences from the model, and apply a model-free algorithm (Q-learning in the original implementation, but more generally we could consider SARSA or policy gradient) as these were real experiences. In Dyna-2 (Silver al. 2008), the same process is applied but in the context the agent is currently in in-starting the simulat iimulations from the current state-and and adapting the policy locally (for example through separate fast weights). In MCTS, the model is used to build tree of possibilities. Can we apply our model directly in these scenarios? While we don't 't have full model of the environment, we can produce a causally correct simulation of rewards and values; one that should generalize to policies different from those the agent was trained on. Policy probabilities, values, rewards and ends of episodes are the only variables that the above RL algorithms need.
Here we propose specific implementation of Dyna-style policy-gradien algorithm based on the models discussed in the paper. This is meant as a proof of principle, and more exploration is left for future work.
As the agent sees an observation yt+1, it forms an internal agent state 8, from this observation and the previous agent state: 81+1 = RNN, ae, 11+1). The agent state in our implementation is the state of the recurrent network, typically LSTM (Hochreiter & Schmidhuber 1997). Next, let us assume that at some point in time with state so the agent would like to learn to do simulation from the model. Let h be the state of the simulation at time t. The agent first sets h = g(8.. Qp) and proceeds with n-steps of the simulation recurrent network update h4+1 RNN 21, a.). The agent learns the model p(2)(h,) which it can use to simulate forward. We assume that the model was trained on some (non-optimal) policy/experience We would like derive an optimal policy and
17

---

Under review as a conference paper at ICLR 2020
value function. Since these need to be used during acting (if the agent were then act optimally in the real environment), they are functions of the agent state sy: [(a]);)) ((t)). Now in general, he o 8, but we would like use the simulation to train optimal policy and value function. Thus we define a second pair of functions TT (oe (4, z)), Vn 2)). Here the extra 2, are needed, since the he has seen z's only up to point 2--1-
Next we are going to train these functions using policy gradients on simulated experiences. We start with some state 5, and produce a simulation be+1...... /tr by sampling 2t from the model each step and action ag ~ Tu(a; Me, zx). However at the initial point I, we sample from т, not The and compute the value V. not Vn. Sequence of actions, values and policy parameters are the quantities needed to compute policy gradient update. We use this update to train all these quantities.
There is one last element that the algorithm needs. The values and policy parameters are trained at the start state and along the simulation by n-step returns, computed from simulated rewards and the bootstrap value at the end of the simulation. However this last value is not trained in any way because it depends on the simulated state (hr) not the agent state sT. We would like this value to equal to what the agent state would produce: (ST). Thus, during training of the model, we also train Vi (br) to be close to ((a) by imposing Lg penalty. In our implementation, we actually impose a penalty at every point f during simulation but we haven't 't experimented with which choice is better.
G SIMPLE EXTENSIONS
Variance reduction. To reduce the variance of a simulation, is possible to sample the 2 from a proposal distribution (h)) The correct expectation can be still recovered by using importance weight: w p(=,
Data efficient training. Usually, we know the distribution of the used partial view: ~ m(z]ss))- When training the p(2g(h+) model, we can then minimize the exact (L(m(Z))s:) p(Z)(h))).
H MODELS TRAINED BY CLUSTERING
When using a tree-search, we want to have small branching factor at the chance nodes. A good & variable would be discrete with a small number of categories. This is satisfied, the žt is the intended action and the number of the possible actions is small. We do not have such compact discrete žt. using as 2n the observation, the policy probabilities or some other modeled layer. Here, will present a model that approximates such causal partial models. The idea is to cluster the modeled layers and use just the cluster index as 24. The cluster index is discrete and we can control the branching factor by choosing the the number of clusters.
Concretely, let's call the modeled layer 2t. We will model the layer with a mixture of components. The mixture gives discrete latent variable 2 represent the component index. To train the mixture, we use clustering loss to train only the best component to model the xt, given h, and 2;:
Lelustering = min( Belustering log p(=, (h)) logp(xeelung 2x))
(29)
2t
where p(ze(hx) is model of the categorical component index and Belustering (0, 1) is hyper- parameter to encourage moving the information bits to the latent 2,. During training, we use the index of the best component as the inferred 2z. In theory, better inference be obtained by smoothing.
In contrast training by maximum likelihood, the clustering loss uses just the needed number of the mixture components. This helps to reduce the branching factor in a search.
In general, the cluster index is not guaranteed to be sufficient loss log pæx, hh, 21) is not zero. For example, if x, is the next observation, the number of mixture components may need to be unrealistically large, if the observation can contains many distractors.
backdoor, if the reconstruction
18

---

Under review as a conference paper at ICLR 2020
1 VALUE-ITERATION ANALYSIS ON MDPS
I.1 OPTIMAL VALUE DERIVATIONS
We derive the following two model-based evaluation metrics for the MDP environments.
(CPM(()) (so): optimal value computed with the non-causal model, which is trained with training data from policy n, starting from state 80-
()) (so): optimal value computed with the causal model, which trained with training data from policy æ, starting from state 8n-
The theoretical analysis of the MDP does empirically trained models from the policy data but rather assumes that the transition probabilities p(5++1 8;, a:) of the MDP, and the policy, π Ta; 8;) or π (2, s,), from which training data are collected are accurately learned by the model.
Computation of V (N) For the non-causal model,
k
V
KCCPM(+)
max Bo
Ex, (1 az) 80, ao a a:]
&
Σ == s,
Σ
p(s; su, " (a).....)))))(ss
max
Notice that the probability of 8; is affected by a; here, because the network gets ag as an input, when predicting the 11+1- This will introduce the non-causal bias. The network implements the expectation implicitly by learning the mean of the reward seen in the training data. We can compute the expectation exactly, if we know the MDP. The p(8; 80, dg,-.., a.) can be computed recursively in two-steps as:
p(1; 8p, ao. ...,, 0i (TT(a: 8i) Ex p(s) so, ao.
p(s: 80, ao..... 0.))
(30)
ai () s(
Here, we see the dependency of the learned model on the policy n. The remaining terms can be expressed as:
p(s: 80, oo, (a,ii)
Σ p(si, 8:-1 80, ao. .o.....,4.... a;-1)
(31)
8L-1
Σ p(8-11
aa-1)p((s; (1-1,a-11)
(32)
5;-1
Denoting p(s:) so, ao....,,.y az) by Si.j. have the two-step recursion
S1,-1 x(a (s)
S Ex (a; ss)
(33)
T(a;
S = Σ S p(s: 8;-1, 0-11)
(34)
81-1
with S1,0 =p011 00). We then compute V (50) maxag..... Σt E., (1,741181,05)) 0z).
19

---

Under review as a conference paper at ICLR 2020
Computation of ())(())
For the causal model,
VCPM(=) (50) = max dg
p(21 Su, ao) max Σ
p(22 so. ao. 21, a1)
(35)
21
22
max
p(zk 50. ao, 21, a...... 2k-1, ax-1)
(36)
2k
B(r 80: 00, 21, ;0,,.,,,,,, ;a)].
(37)
where for any E k1,
p(zs do: 21, .. #1-1, 0;-1) = Σ p(si, 100 00, 21, 01,..., Fi 0;-1)
(38)
Σ p(s: 88, 00, -,, 0; <<;; [x)..
(39)
(40)
Denoting p(s; | so, ao 2411, 0/-1) by Z,, we have
Z =
Σ p( 81
su. av. 21 25-1, 0s-1)
(41)
9:-1
Σ p(s; |1-11,--1)p(ss111 so, ao, 21 .... 21 a;-1)
(42)
Σ p(s, (=-1,,--))p(s-11 so, ao, 21 ... 21-1),
(43)
where we used the fact that 8 is independent of a;-1, given 2:-1- Furthermore,
p(3-11 so, ao: 31 ..... 23-1) p 1 2
8o, ao 21 ... a4-2)
(44)
p
80: ao: 21 a:-2)
3/-1 )p 9/-1
80. do, ...,, 21-2, a1-2)
(45)
Ex__, <(21-1
6/ (1
80, 00, 21 .... 24 04 04-2)
(2
8i-1)
(46)
Ex__, (
12,1)
Therefore we can compute Z; recursively,
T π (i-1)
Z Σ p(s: s:-1, 0s-1)
(47)
Exe_, ( 16/1))2,11
20

---

Under review as a conference paper at ICLR 2020
with Z p(s1 su, ao). The last term to compute in the definition of VCPM(s) (50) is
E Birr41 ai) s0: 0o, 21, ....... 0:)]
B Bleet1 (88, (s), du, 21, 01,....,,,,.....................
(48)
==
==0
Σ p(s; 80, ao, 21, ,,,,,,,., (,) (s)
1=
(49)
π [ zi 8:)
T4+ (5),
(50)
==
1.2
PLANNING WITH CON-CAUSAL VS CAUSAL MODELS ON FUZZYBEAR
M/s
0.5
*b
1/
Xi
bug mn 0
run
hug
rin
hug
TUn
My. Ho X1
Al 60
My HD
Fry-
01
B/slau. R0. 0,ay - I.
Eirgle, on. to 0, 1)
Egrlan no,
Elrshe, An. > 1,01 -0.8
1, Dff
Kirz/ko Ri.
1|
(a) Non-Causal Model
(b) Causal Model
Figure 7: FuzzyBear decision trees evaluated with NCPM and CPM based on the optimal policy data with the intended action. Decision paths through red action nodes that give the maximum expected rewards are highlighted in red. The blue nodes are states and the gray ones are chance nodes.
In Figure 7(a) the non-causal agent always chooses hug at step since has learned from the optimal policy that reward of +1 always follows after taking a hug. Thus from the non- causal agent's point of view, the expected reward always after hugging. This wrong since only hugging teddy bear gives reward 1. Morcover it exceeds the maximum expected reward 0.5 of the FuzzyBear MDP. In Figure [bb] the causal agent first samples the intention 21 from the optimal policy, giving equal probability of landing in either of the two chance nodes. Then it chooses hug 21 = 0, indicating a teddy bear since the optimal policy intends to hug only it observes a teddy bear. Likewise, it chooses if ž 1. indicating grizzly bear. While the non-causal model expects unrealistically high reward, the causal model never over-estimates the expected reward.
1.3 LEARNING WITH OPTIMAL POLICY AND VARYING E-EXPLORATION:
We analyze learning from optimal policy with varying amounts of e-exploration for models on FuzzyBear (Figure ( and AvoidFuzzyBear (Figure 8(b)). As the parameter e-exploration varies in range (0, 1]. the causal model has constant evaluation since the intended action is not affected by the randomness in exploration. The non-causal model, on the other hand, evaluates based on the deterministic optimal policy data (i.e. at ε 0) at an unrealistically high value of 1.0 when the maximum expected reward is 0.5. As - 1, the training data becomes more random and its optimal evaluation expectantly goes down to match the causal evaluation based on a uniformly random policy. The causal evaluation based the optimal policy VeM(r)) converges to the ground truth environmental evaluation V - 0.
21

---

Under review as a conference paper at ICLR 2020
NCPM leamed from with exp with exploracion
NCPM learned from exploration
any polic
poliy fiom
0.8
0.9
Vite)
1/2
0.7
0.4
0.6
onn
0.0
0.6
0.0
0.2
0.4
0.6 0.8
Parameter (a)
Parameter e (b)
Figure 8: aa) In FuzzyBear, we optimal policy with E-exploration to generate training data for the Non- Causal Partia Model (NCPM) and Causal Partial Model (CPM). We vary the exploration parameter (0, I] observe differences found optimal values VM(m) under the model (m). where (xx denotes either CPM NCPM trained on hchavior policy The NCPM evaluation (CCMM(T)) gives unrealistically high 0). Expectantly, decreases the level of CPM optimal value VCPM(s...))) Iearned from the uniformly random policy as + 1. The CPM optimal values
value 1.0 learned from the deterministic optimal policy
(CPM(T>) constant for any value of based the theoretical analysis in Section L.I (b) shows the same
plots as [aj the AvoidFuzzy environment. Learning from any policy x, opprop value always equals the maximum expected reward 0.6, by correctly choosing to stay home.
J DETAILS FOR 3D EXPERIMENTS
J.1 CONDITIONAL DIRICHLET MIXTURE
When the backdoor variable 2, was chosen to be the action probabilities, the distribution p(=: was chosen as mixture-network with N₂ Dirichlet components. The concentration parameters org (hg) of each component were parametrized as ag (lt) ssooftmax(f((h)) where J& is the output of relu-MLP with layer sizes [256, 64, No X Naj, is total concentration parameter and N₂ is the number of actions.
J.2 HYPER PARAMETERS AND TRAINING
The ryper-parametes value ranges used in our 3D experiments are similar to Gregor al. (0019) and are shown in Table
To speed up training, we interleaved training on the T-maze level with simple "Food" level, in which the agent simply had to walk around and eat food blocks (described by Gregor et al. (2019)).
J.3 ANALYSIS OF ROLLOUTS
For each episode, rollouts are generated after having observed the first frames from the environ- ment. For the rollouts, processed the first 25 frames to classify the presence of food blocks by performing color matching RGB values, using K-means and assuming 7 clusters. Rollouts were generated shortly after the policy had achieved ceiling performance (15-20 million frames seen), but before the entropy of the policy reduces to the point that there is no longer sufficient exploration. See Figure for these same results for later training.
22

---

Under review as a conference paper at ICLR 2020
Table 3: Hyper-paramelerss used. Each reported experiment repeated at least 5 times with different random seeds.
Hyper-parameter
Description
Value
Mpolicy
policy learning rate
0.0001
Pmodel
model learning rate
0.0005
c
policy entropy regularization
0.0004
B1
Adam $1
0.9
82
Adam 32
0.999
Lo
Overshoot Length
8
La
Unroll Length
100
Number points used evaluate the generative loss per trajectory
N₂
6
Number of points used to evaluate the generative loss per overshoot
Ng
2
N,
Number of ConvDRAV Steps
8
N₂
Number of units in LSTM
512
a
Total concentration of Dirichlet distributions
100
N
Number of components of Dirichlet mixture
10
forced action gen.
Probability generate food
0.8
0.6
2.0
Entropy of policy
1.0
0.2
0.0
15M-20M
25M Data [frames]
10M
20M
30M
40M
50M
Training frames seen
Figure 9: While earlier in training. CPM generates a diverse range of outcomes (food or no food), as the policy becomes more deterministic (as seen in the right plot of the policy entropy over training). CPM starts to generate more food and hecomes overoptimistie, similar to NCPM. This can be avoided by training the model with non-zero s-exploration.
23

---

Under review as a conference paper at ICLR 2020
K
MODEL ROLLOUTS
Table Different types model rollouts considered. Blue cells indicate variables require interaction with the real environment, depending on the agent's state 8t. Orange cells indicate variables that can be computed from the model's state ht.
Rollout
Rollout with forced actions
Rollout with forced actions and backdoor
Backdoor Action
ktg 7((122))
p(2
mee 7 0, (as 24)
0, (a Ba) & ml2. st)
State
RNN, (h, 01-1, 7111
Prediction
Me plye (4)
rolouts
obnarred fames
niserved Tramte
Figure 10: Full rollouts for NCPM, conditioned on forced actions.
24

---

Under review as a conference paper at ICLR 2020
molloo
odnarasd Sarress
olserved
Figure 11: Full rollouts for CPM, conditioned on forced actions and generated z.
25

---

Under review as a conference paper at ICLR 2020
molooos
adnarand Samris
olearved
Figure 12: Full rollouts for CPM, conditioned on forced actions and forced z.
26


Reviews:
[{"review_id": "rkez7nP55B", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "SUMMARY:\nThe authors apply ideas from causal learning to the problem of model learning in the context of sequential decision making problems. They show that models typically learned in this context can be problematic when used for planning. The authors then reformulate the model-learning problem using a causal learning framework and propose a solution to the above-mentioned problem using the concept of \"backdoors.\" They perform experiments to demonstrate the advantages of doing so.\n\nMAJOR COMMENTS:\nOverall, I'm positive about the submission, though I think there's room for improvement.\n\nThe paper is well-written and does a very nice job of formulating the model-learning problem through the lens of causal learning, and it is convincing that the baseline model-learning procedure suffers from confounds. Moreover, the proposed modified procedure for model learning would seem to have the potential to fundamentally and positively impact model-learning in general. Finally, I found that limited experiments supported the points made in the paper.\n\nThat said, I think the paper as written could be substantially improved with a little extra effort. First, it does not -- in my opinion -- pass a basic reproducibility test, i.e., I am not confident that I could implement the proposed modified model-learning procedure even after reading the paper. The authors should seriously consider adding an algorithm box with pseudocode to both show the flow of training data to the model-learning steps and also how planning can be accomplished with the learned models.\n\nSecond, the paper would greatly benefit from any discussion (theoretical, experimental, or -- ideally -- both) of the tradeoffs that would arise when considering the proposed technique. What price is paid for the gain in planning accuracy? Is more data required for learning compared to non-causal approaches? If such a price is incurred, how was this made clear in the presented results?\n\nFinally, an important aspect of the experiments seems to have gone un-discussed. Namely, what is the reason for the behavior of the proposed technique in Figure 4(a)? From much of the paper, I would have expected the red dots to live entirely on the horizontal dotted line $V^*_{env}$, but instead this only happens when the behavior policy has the same value. Why is this? Moreover, why do *better* behavior policies result in *worse* performance for the optimal model evaluation policies?\n\nMINOR COMMENTS:\n    * The authors should define \"par_k\", first referenced at the bottom of p3, before using it. Similarly for $\\psi_k$\n    * p5, second paragraph. I think there's a typo in \"... sample $y_2$ from $q(\\cdot | h_1)$ ...\" (should probably be conditioned on $h_2$ not $h_1$).\n\nPOST-RESPONSE COMMENTS:\nIn my opinion, the authors adequately addressed both my own concerns, and also several valid concerns from the other reviewers. Therefore, I'm raising my score to \"accept.\"", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "rJgKAsG09B", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper considers the problem of predicting a variable y given x where x may suffer from a policy change, e.g., x may follow a different distribution than the original data or suffer from a confounding variable. The flow of the paper proceeds in learning a causally correct model in the sense that the model is robust to any intervention changes. Specifically, the paper considers a setting called \"partial model\" meaning a generative model conditioned on functions of past observations. To make the partial model causally correct, the paper considers the partial model conditioned on the backdoor that blocks all paths from the confounding variables. \n\n1. The problem that this paper addresses seems to be new and interesting. The approach makes much sense: the problem is due to the confounding effect which can be addressed by introducing some other variables that implicitly blocks the confouders.  \n\n2. The  paper assumes the existence of the backdoor variable which is crucial for causal correctness. Does the backdoor always exist? Pearl's book may have some discussions on this. It will be still useful to include some materials in case of unfamiliar readers. \n\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper."}, "extra_text": {}}, {"review_id": "HJl1FwQaFH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "*Summary*\n\nThis paper considers the effect of partial models in RL, authors claim that these models can be causally wrong and hence result in a wrong policy (sub optimal set of actions). Authors demonstrate this issue with a simple MDP model, and emphasize the importance of behavior policy and data generation process. Then authors suggest a simple solution using backdoors (Pear et al) to learn causally correct models.They also conduct experiments to support their claims.\n\n*Decision*\n\nI vote for rejection of this paper, based on the following argument:\n\nTo my understanding authors are basically solving the “off-policy policy evaluation” problem, without relating to this literature. For example, the MDP example is just an off-policy policy evaluation problem, and it is very well known that in this case you need to consider the behavior policy, for example with importance sampling.\nEven authors definition of the problem at the end of the page 4, and beginning of page 5, is the problem of “off-policy policy evaluation” when y_t = r_t \nAuthors have not cited any paper in this literature, and did not situate their work with respect to this literature. To my understanding, the proposed solution is basically importance sampling, that is very well known and studied in the field.\n\nAdditionally, I suggest that authors be more careful with their citations, for example, authors cited Silver et al 2017 [The Predictron: End-To-End Learning and Planning], as one recent paper using the method; however Silver et al 2017  is in MRP setting (Markov Reward process) where there is no action, so the described problem setting doesn't apply.  \n\nImprovement\n\nThe current manuscript needs a major revision, mainly 1) situate the work with respect to off-policy policy evaluation literature, and then 2) Considering step 1, a clarification for what is the novelty/ contribution of the current paper is needed. \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "H1eOrcwaKB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper tackles the issue of identifying the causal reasoning behind why partial models in MBRL settings fail to make correct predictions under a new policy. The novel contribution was a framework for learning better partial models based on models learning an interventional conditional, rather than an observational conditional. The paper tried to provide both theoretical and experimental reasoning for this framework. \n\nI vote to (weak) reject the paper due to the major issues with section 2. Furthermore, the paper hard or at times almost impossible to understand as too many assumptions are made and too little is explained.\n\nRecommendations\n\nBecause your graphs are not MDPs, you are not framing your example as an RL problem. This is causing a number of issues with notation and lack of clarity in the argument you're making.\n\n1. It is unclear to me that the FuzzyBear example is correctly constructed as a RL example, reasons being that:  \n- Figure 1 (a) & (b) do not correspond to an MDP as two different states, teddy vs grizzly, are both designated as s_1 and similarly, the two possible actions, hug or run, are both designated as a_1 and thus are not distinct.\n- Note your terminal state for the episodes\n- Have a reward for (s0, a0) as every s-a pair should have a reward\n- It would be helpful to note that the environments in Figure 1 are stochastic\n\n2. Clarify notation. There are a number of assumptions about what background knowledg the reader should have. Given the bridging of disciplines in the paper, it would be useful to provide more detail on notation in Section 3.\n\n3. Add a section on reinforcement learning in Section 3. If it's the last subsection in section 3, you could describe the relationship between the various causal reasoning and RL principles. This would further clarify how you're bridging these subtopics.\n\n4. For sentence,\n\n\"Fundamentally, the problem is due to causally incorrect reasoning: the model learns the observational condi- tional p(r|s0, a0, a1) instead of the interventional conditional given by p(r|s0, do(a0), do(a1)) = s1 p(s1|s0, a0)p(r|s1, a1).\"\n\nAs you don't cover the meaning of the do() operator until a later paragraph, provide a quick description of it as it is not common knowledge to a general AI audience, e.g., where do() indicates that the action was taken.\n\n5. Correct the following sentences,\n\n\"Mathematically, the model with learn the following conditional probability:\"\n\n\"In Section 3, we review relevant concepts from causal reasoning based on which we propose solutions that address the problem.\"\n\n6. I recommend putting the interventional conditional equation, p(r|s0, do(a0), do(a1)) = 􏰀s1 p(s1|s0, a0)p(r|s1, a1), on its own line as the reader is doing a comparison of it with the previous equation, p(r|s0, a0, a1), given on page 2.\n\n7. Strengthen your abstract by aligning more with claims you make in your conclusion.\n\n8. The experiments in Figure 5 are averaged over 5 seeds. This is not enough to be statistically significant - furthermore, there are no error bars in the Figure.\n\nQuestion(s):\n1. You've indicated two policies for Figure 1 (a):\n- pi1: the agent knows it is encountering a teddy bear, so it will hug\n- pi2: the agent knows it is encountering a grizzly bear, so it will run\nIs this the \"change in the behavior policy\" that you're referring to? If so, make this clearer, this currently requires a lot of work by the reader to make sense of it.\n\n2. What are the partially observable parts of the environments in Figure 1 (a) & (b)? Make this clear.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors show that in a reinforcement learning setting, partial models can be causally incorrect, leading to improper evaluation of policies that are different from those used to collect the data for the model.  They then propose a backdoor correction to this problem that allows the model to generalize properly by separating the effects of the stochasticity of the environment and the policy.  The reviewers had substantial concerns about both issues of clarity and the clear, but largely undiscussed, connection to off-policy policy evaluation (OPPE).  

In response, the authors made a significant number of changes for the sake of clarity, as well as further explained the differences between their approach and the OPPE setting.  First, OPPE is not typically model-based.  Second, while an importance sampling solution would be technically possible, by re-training the model based on importance-weighted experiences, this would need to be done for every evaluation policy considered, whereas the authors' solution uses a fundamentally different approach of causal reasoning so that a causally correct model can be learned once and work for all policies.

After much discussion, the reviewers could not come to a consensus about the validity of these arguments.  Futhermore, there were lingering questions about writing clarity. Thus, in the future, it appears the paper could be significantly improved if the authors cite more of the off policy evaluation literature, in addition to their added textual clairifications of the relation of their work to that body of work.  Overall, my recommendation at this time is to reject this paper.
Decision: Reject