MoCoDA: Model-based Counterfactual Data Augmentation
Authors: Silviu Pitis
Elliot Creager
Ajay Mandlekar
Animesh Garg
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
off-policy reinforcement learning
model-based reinforcement learning
offline reinforcement learning
data augmentation

Abstract: The number of states in a dynamic process is exponential in the number of objects, making  reinforcement learning (RL) difficult in complex, multi-object domains. For agents to scale to the real world, they will need to react to and reason about unseen combinations of objects. We argue that the ability to recognize and use local factorization in transition dynamics is a key element in unlocking the power of multi-object reasoning. To this end, we show that (1) known local structure in the environment transitions is sufficient for an exponential reduction in the sample complexity of training a dynamics model, and (2) a locally factored dynamics model provably generalizes out-of-distribution to unseen states and actions. Knowing the local structure also allows us to predict which unseen states and actions this dynamics model will generalize to. We propose to leverage these observations in a novel Model-based Counterfactual Data Augmentation (MoCoDA) framework. MoCoDA applies a learned locally factored dynamics model to an augmented distribution of states and actions to generate counterfactual transitions for RL. MoCoDA works with a broader set of local structures than prior work and allows for direct control over the augmented training distribution. We show that MoCoDA enables RL agents to learn policies that generalize to unseen states and actions. We use MoCoDA to train an offline RL agent to solve an out-of-distribution robotics manipulation task on which standard offline RL algorithms fail. 

MoCoDA: Model-based Counterfactual Data
Augmentation
Silviu Pitis" 1 Elliot Creager' Ajay Mandlekar² Animesh Garg
University of Toronto and Vector Institute, 2NVIDIA
Abstract
The number of states in dynamic process is exponential in the number of objects, making reinforcement learning (RL) difficult in complex, multi-object domains. For agents to scale to the real world, they will need to react to and reason about unseen combinations of objects. We argue that the ability to recognize and use local factorization in transition dynamics is key element in unlocking the power multi-object reasoning. To this end, we show that (1) known local structure the environment transitions is sufficient for exponential reduction in the sample complexity of training a dynamics model, and (2) locally factored dynamics model provably generalizes -i-or--distiibutionn to unseen states and actions. Know- ing the local structure also allows us to predict which unseen states and actions this dynamics model will generalize to. We propose leverage these observations in a novel Model-based Counterfactual Data Augmentation (MoCoDA) framework. MoCoDA applies learned locally factored dynamics model to an augmented distribution of states and actions to generate counterfactual transitions for RL. MoCoDA works with a broader set of local structures than prior work and al- lows for direct control over the augmented training distribution. We show that MoCoDA enables RL agents learn policies that generalize to unseen states and actions. We use MoCoDA to train an offline RI. agent to solve an out-of- distribution roboties manipulation task on which standard offline RI. algorithms fail
Introduction
Modern reinforcement learning (RL. algorithms have demonstrated remarkable success in several different domains such games [42 53 and robotic manipulation [3 4. By repeatedly attempting a single task through trial-and-error, these algorithms can learn to collect useful experience and eventually solve the task of interest. However, designing agents that can generalize in off-task and multi-task settings remains an open and challenging research question. This is especially true the offline and zero-shot settings, in which the training data might be unrelated to the target task, and may lack sufficient coverage over possible states.
One way to enable generalization in such cases is through structured representations of states, transition dynamics, or task spaces. These representations can be directly learned, sourced from known learned abstractions over the state space, or derived from causal knowledge of the world. Symmetries present in such representations enable compositional generalization to new configurations of states tasks, either by building the structure into the function approximator or algorithm or by using the structure for data augmentation 33. 51].
28 58 15 43)
In this paper, we extend past work on structure-driven data augmentation by using a locally factored model of the transition dynamics to generate counterfactual training distributions. This enables agents to generalize beyond the support of their original training distribution, including to novel
"Correspondence to spitisecs toronto. edu
VVisualizations & code available at attttpp: wttphhtpt///tteeeee :/siiess...e./eiiiie google www.le..c...iee////ee http.///ccee//cccccaaaeeeeip---p
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

city
highway
(T)
drive
right
2
trojectory samplex
drive on left side
X
3
city
highway
PMoCODA(T)
drive
right side
laail driving miles come sule
iceee soae (H+1)
vebiile
cover actiogs
drive on left side
4
transition structure
Figure 1: Out-of-Distee Distriiiitiittionn Generalization using MOCODA: A driver can use MOCoDA to quickly adapt to driving in the left lane during UK trip. Their prior experience Ptwr(T) (top left) contains mostly right-driving experience (e.g-
and limited amount of left-driving experience after renting the car in the UK (c.g. 3 p. A locally factored model that captures the transition structure (bottom left) allows the agent to accurately sample counterfactual experience from Po (r) (bottom center), including novel left-lane city driving maneuvers (e.g. 4 This enables fast adaptation when learning an optimal policy for the new task (UK driving). Our framework MoCoDA draws single-step transition samples from PMocons (T) given PEMe (T) and knowledge of the causal structure; several realizations of this framework are described Section
tasks where learning the optimal policy requires access to states never seen in the experience buffer. Our key insight is that learned dynamics model that accurately captures local causal structure (a "locally factored" dynamics model) will predictably exhibit good generalization performance outside the empirical training distribution. We propose Model-based Counterfactual Data Augmentation (MoCoDA), which generates an augmented state-action distribution where its locally factored dynamics model is likely perform well, then applies its dynamics model to generate new transition data. By training the agent's policy and value modules on this augmented dataset, they too learn to generalize well out-of-distribution. To ground this in an example, use MOCODA to adapt to driving on the left side of the road while on vacation in the UK (Figure D Given knowledge of the target task, we can even focus the augmented distribution on relevant areas of the state-action space (e.g., states with the car on the left side of the road).
consider how US driver might
Our main contributions
A. Our proposed method, MoCoDA, leverages a masked dynamics model for data-augmentation in locally-factored settings, which relaxes strong assumptions made by prior work on factored MDPs and counterfactual data augmentation.
B. MoCoDA allows for direct control of the state-action distribution on which the agent trains; we show that controlling this distribution in task relevant way can lead improved performance.
C. We demonstrate "zero-shot" generalization of policy trained with MoCoDA to states that the agent has never seen. With MoCoDA, we train an offline RL agent to solve an -usttf--istribution robotics manipulation task which standard offline RL algorithms fail.
2
Preliminaries
2.1 Background
We model the environment infinite-horizon, reward-free Markov Decision Process (MDP), described by tuple (S, A, P.r) consisting the state space, action space, transition function, and discount factor, respectively [52] 15]. We use lowercase for generic instances and uppercase for

---

Me++1
L
2t
0
Pa(yu:) ) {vl, ar
Ls
M(
0
1
V-+11
B1+1 Pz (.130, a.)
1
S x A
G
The state-action space S x A is divided into local subsets, L1, L2, Ls. which each have their own factored causal structure, g The local transition model p is factored according to :eg.gg. in the example shown, pC (Ire, Mr. a.) = (2x) Py (vs. a.)]
Figure 2: Locally Factored Dynamies:
variables (e.g., s E rSnge(S) - 8, though we also abuse notation and write S E S). A task is defined as tuple (r. Po), where r : S x A + I is a reward function and Po is an initial distribution over S. The goal of the agent given a task is to learn a policy : S + AA that maximizes value Ep. Ee (r((st, at). Model- based RL is one approach to solving this problem, in which the agent learns a model Ps of the transition dynamics P. The model is "rolled out" to generate "imagined" trajectories, which are used either for direct planning [1] g or as training data for the agent's ' policy and value functions [56] 20
Factored MDPs, A factored MDP (FMDP) is a type of MDP that assumes a globally factored transition model, which can be used to exponentially improve the sample complexity of RL [16 24 451. In an FMDP, states and actions are described by set of variables {X*} J. so that S S = X x2xXX x.....x X"", and each state variable X X (X* is a subspace of S) is dependent on subset of state-action variables (its "parents" P)(X')) at the prior timestep. X ~ ()))a()))). We call a set {X²} of state-action variables a "parent set" if there exists state variable X such that [X*} Pa(X*). We say that X' is "child" of its parent set Pa(Xt). We refer to the tuple (X Pa(X)) (.) as a "causal mechanism".
Local Causal Models. Because the strict global factorization assumed by FMDPs rare, recent work on data augmentation for RL and object-oriented RL suggests that transition dynamics might be better understood in a local sense, where all objects may interact with each other over time, but in locally sparse manner 115 28 39). Our work uses an abridged version of the Local Causal Model (LCM) framework [5., as follows: We assume the state-action space decomposes into a disjoint union of local neighborhoods: SxA = & L L2 L... U Ln. A neighborhood L is associated with its own transition function pc, which is factored according to its graphical model G (29]. We assume no two graphical models share the same structure (i.e., the structure of G uniquely identifies f.). Then, analogously to FMDPs, if (st, 0t) E f., each state variable X1+1 at the next time step is (X;
dependent its parents (X; at the prior timestep, X (P (11+1) 1)). We define mask function M S x A -> {La} that maps (s, a) E L the adjacency matrix of G This formalism is summarized in Figure 2 and differs from FMDPs in that each L has its own factorization.
Given knowledge of M, the Counterfactual Data Augmentation (CoDA) framework (51) allowed agents to stitch together empirical samples from disconnected causal mechanisms to derive novel transitions. It did this by swapping compatible components between the observed transitions to create new ones, arguing that this procedure can generate exponentially more data samples as the number of disconnected causal components grows. CoDA was shown to significantly improve sample complexity in several settings, including the offline RL setting and a goal-conditioned robotics control setting. Because CoDA relied empirical samples of the causal mechanisms generate data in a model-free fashion, however, it required that the causal mechanisms be completely disentangled The proposed MoCoDA leverages dynamics model to improve upon model-free CoDA in several respects: (a) by using learned dynamics model, MOCoDA works with overlapping parent sets, (b) by explicitly modeling the parent distribution, MOCoDA allows the agent to control the overall data distribution, (c) MOCODA demonstrates zero-shot generalization to new areas of the state space, allowing the agent to solve tasks that are entirely outside the original data distribution.
"This assumption is matter of convenience that makes counting local subspaces in Section 3 slightly easier and simplifies our implementation of the locally factored dynamics model in Section To accommodate cases where subspaces with different dynamics share the same causal structure, one could identify local subspaces using latent variable rather than the mask itself, which we leave for future work.
3

---

2.2
Related Work
RL with Structured Dynamics. A growing literature recognizes the advantages that structure can provide in RL, including both improved sample efficiency [37 [9] and generalization performance Some of these works involve sparse interactions whose structure changes over time which is similar to and inspires the locally factored setup assumed by this paper. Most existing work focuses on leveraging structure to improve the architecture and generalization capability of the function approximator [62]. Although MoCoDA also uses the structure for purposes of improving the dynamics model, our proposed method is among the few existing works that also use the structure for data augmentation [38 40 51].
62
59
54
15
28
Several past and concurrent works aim tackle unsupervised object detection 36 12
(i.e., learning
an entity-oriented representation of states, which is prerequisite for learning the dynamics factor- ization) and learning the dynamics factorization 277 60] These are both open problems that run orthogonal to MOCoDA. We expect that as solutions for unsupervised object detection and factored dynamics discovery improve, MoCoDA will find broader applicability.
RL with Causal Dynamics. Adopting this formalism allows one to cast several important problems within RL as questions of causal inference, such as off- off--poliy evaluation 7 44 learning baselines
for model-free RL [41], and policy transfer [25). Lu et al. (38) applied SCM dynamics to data augmentation in continuous sample spaces, and discussed the conditions under which the generated transitions are uniquely identifiable counterfactual samples. This approach models state and action variables as unstructured vectors, emphasizing benefit in modeling action interventions for settings such as clinical healthcare where exploratory policies cannot be directly deployed. We take a com- plementary approach by modeling structure within state and action variables, and our augmentation scheme involves sampling entire causal mechanisms (over multiple state or action dimensions) rather than action vectors only. See Appendix Forr more detailed discussion of how MoCoDA sampling relates to causal inference and counterfactual reasoning.
3
Generalization Properties of Locally Factored Models
3.1 Sample Complexity of Training Locally Factored Dynamics Model
In this subsection, we provide original adaptation of elementary result from model-based RL to the locally factored setting, to show that factorization can exponentially improve sample complexity. We note that several theoretical works have shown that the FMDP structure can be exploited to obtain similarly strong sample complexity bounds in the FMDP setting. Our goal here is not to improve upon these results, but adapt small part (model- general locally factored setting and show that local factorization enough for (1) exponential gains in sample complexity and (2) out-of-distributiinn generalization with respect to the empirical joint, to a set of states and actions that may be exponentially larger than the empirical set. Note that the following discussion applies to tabular RL, but we apply our method to continuous domains.
generalization) to the significantly more
Notation. We work with finite state and action spaces (S)) [AA] < oo) and assume that there are m local subspaces L of size [C], such that m|C| = SS ||. For each subspace L. we assume transitions factor into causal mechanisms {P.}, each with the same number of possible children, |c.). and the same number of possible parents, |Pa.|- Note m[I;/c; |S| (child sets mutually exclusive) but mII.|Pa 2 S||AA| (parent sets may overlap).
Theorem 1. Let n be the number of empirical samples used to train the model of each local causal mechanism, PfC at each configuration of parents Pa; = x. There exists constant such that, if
ck2 |cd] (0)((5)
n N
e²
then, with probability at least 1 8, we have:
max [P(s (o(s. a)ll Se. (s.a)
Sketch of Proof. We apply a concentration inequality to bound the & error for fixed parents and extend this to a bound on the l error for a fixed (s, a) pair. The conclusion follows by a union bound across all states and actions. See Appendix for details.
D

---

To compare to full-state dynamics modeling, we can translate the sample complexity from the per- parent count n. to total count N. Recall ml1;/a = |S], so that |cs| = (|$1/m) 1/*, and mPl,|Pa|l 2 [S] We assume a small constant overlap factor 2 1, so that |Pa. = =[[]]]||mmm We need the total number of component visits to be nı Pa; (km, for a total of =u[[S|]A|/m ||//sm/////// state-action visits, assuming that parent set visits are allocated evenly, and noting that each state-action visit provides k parent set visits. This gives:
1//m
Corollary 1. To bound the error as above, we need to have
2 cmkk (|S)) log (S|(A1))))
total train samples, where we have absorbed the factor into constant Comparing this to the analogous bound for full-state model learning (Agarwal et al. 0. Prop. 2.1):
c|S|a| log( (og((5)1A)))))
see that have gone from super-linear 0([8||A |A| log log([S[[AA)) sample complexity in terms of [S]|A], the exponentially smaller (mnk² ((S) ([ss(()m//m))) 1/8 log([S]|AA))
This result implies that for large enough |S|| our model must generalize to unseen states and actions, since the number of samples needed (N) is exponentially smaller than the size of the state-action space ((S][A). In contrast, if it did not, then sample complexity would be (([))[AA))
The global factorization property of FMDPs is a strict assumption that rarely holds in reality. Although local factorization is broadly applicable and significantly more realistic than the FMDP setting, it is not without cost. In FMDPs, we have single subspace (m 1). In the locally factored case, the number of subspaces m is likely to grow exponentially with the number of factors k. as there exponentially many ways that k factors can interact. To be more precise, there are k:2h possible bipartite graphs from i nodes to nodes. Nevertheless, by comparing bases (2 < |S)) see that we still obtain exponential gains in sample complexity from the locally factored approach.
Remark 3.1.
3.2 Training Value Functions and Policies for citto--DDittrbuutonn Generalization
In the previous subsection, we saw that a locally factored dynamics model provably generalizes outside of the empirical joint distribution. A natural question is whether such local factorization can be leveraged to obtain similar results for value functions and policies?
We will show that the answer is yes, but perhaps counter-intuitively, it is not achieved by directly training the value function and policy on the empirical distribution, as is the case for the dynamics model. The difference arises because learned value functions, and consequently learned policies, involve the long horizon prediction Ep,,, Σ 7"r(st, oz). which may not benefit from the local sparsity of G When compounded over time, sparse local structures can quickly produce an entangled long horizon structure (cf. the "butterfly effect"). Intuitively, even if several pool balls are far apart and locally disentangled, future collisions are central to planning and the optimal policy depends on the relative positions of all balls. This applies even if rewards are factored (e.g., rewards in most pool variants) [54].
We note that, although temporal entanglement may be exponential in the branching factor of the unrolled causal graph, it's possible for the long horizon structure to stay sparse (e.g., & independent factors that never interact, or long-horizon disentanglement between descision relevant and decision irrelevant variables (19).. It's also possible that other regularities in the data will allow for good out-of-distribution generalization. Thus, cannot claim that value functions and policies will never generalize well cuttof-ddisrrbutiion (see Veerapaneni et al. [58] for an example when they do). Nevertheless, we hypothesize that exponentially fast entanglement does occur in complex natural systems, making direct generalization of long horizon predictions difficult.
Out-of-distribute generalization of the policy and value function can be achieved, however, by leveraging the generalization properties of a locally factored dynamics model. We propose to do this by generating out-of-distribution states and actions (the "augmented parent distribution"). and then applying our learned dynamics model to generate transitions that are used to train the policy and value function. We call this process Model-based Counterfactual Data Augmentation (MoCoDA).
5

---

Iudly frord crnvacos P.
angieeeed data
(..0,]] - Rug
Eoreer destristion 2lix Pils.o
ogeiit
korc dat
Figure 3: RL training with MOCODA: We the empirical dataset train parent distribution model (s, a) and locally factored dynamics model Po(s' s, a). both informed by the local structure. The dynamics model applied to the augmented parent distribution P' (s, a) produce augmented dataset PsP The augmented & empirical datasets are labeled with target task reward, r(s, a) and fed into RL algorithm as training data.
4 Model-based Counterfactual Data Augmentation
In the previous section, we discussed how locally factored dynamics model can generalize beyond the empirical dataset to provide accurate predictions on an augmented state-action distribution we call the "parent distribution". We now seek to leverage this cut-of-distribution generalization in the dynamics model to bootstrap the training of an RL agent. Our approach is to control the agent's training distribution P(s, s') via the locally factored dynamics (s (s,s,, a) and the parent distribution (s a) (both trained using experience data). This allows us to sample augmented transitions (perhaps unseen in the experience data) for consumption by downstream RL agent. We call this framework MOCoDA, and summarize it using the following three-step process:
Step 1
Given known parent sets, model the parent distribution Pø(s, a) and generate an appropriate augmented parent distribution Po(s, a).
Step 2
Apply a learned dynamics model (s to augmented parent distribution to generate "augmented dataset" of transitions (s, a, s').
Step 3
Use augmented dataset s, a, s' ~ Po (alongside experience data, if desired) train an off-policy RL agent on the (perhaps novel) target task.
Figure illustrates this framework in a block diagram. An instance of MoCoDA is realized by specific choices at each step. For example, the original CoDA method [51) is instance of MOCoDA, which (1) generates the augmented parent distribution by uniformly swapping non-overlapping parent locally factored dynamics model. CoDA works when local graphs have non-overlapping parent sets, but does not allow for control over the parent distribution and does not work in cases where parent sets overlap. MoCoDA generalizes CoDA, alleviating these restrictions and allowing for significantly more design choices.
sets, and (2) uses subsamples of empirical transitions
4.1 Augmenting the Parent Distribution
How should the parent distribution be augmented (Step 1) to generate the augmented dataset? In other words, after fitting P& (s, a) to experience. how should we realize Pa(s, a)? We describe some options below, noting that our proposals (MOCODA, MOCODA-U, MOCODA-P) rely on knowledge of (possibly local) parent sets sets-i.e. i.e., they require the state to be decomposed into objects.
Baseline Distributions. If we restrict ourselves to states and actions in the empirical dataset (EMP) or short-horizon rollouts that start in the empirical state-action distribution (DYNA), as typical in Dyna-style approaches [57 20), we limit ourselves to a small neighborhood of the empirical state- action distribution. This forgoes the opportunity to train our off-policy RI. agent on oft-of- data that may be necessary for learning the target task.

---

Another option is to sample random state-actions from S x A (RAND). While this provides coverage of all () relevant to the target task. there is no guarantee that our locally factorized model generalizes well in RAND. The proof of T shows that our model only generalizes well to a particular (s, a) if each component generalizes well on the configurations of each parent set in that (s,a). In context of Theorem I this occurs only if the empirical data used to train our model contained at least n samples for each set of parents in (s, a). This suggests focusing on data whose parent sets have sufficient support in the empirical dataset.
The MOCODA distribution. We do this by constraining the marginal distribution of each parent set (within local neighborhood L) in the augmented distribution to match the corresponding marginal in the empirical dataset. As there are many such distributions, in absence of additional information, is sensible to choose the one with maximum entropy [21]. We call this maximum entropy, marginal matching distribution the MOCODA augmented distribution. Figure P an illustrative example of going from EMP (driving primarily on the right side) to MOCODA (driving on both right and left). We propose an efficient way to generate the MOCODA distribution using a set of Gaussian Mixture Models, one for each parent set distribution. We sample parent sets one at a time, conditioning any previous partial samples due to overlap between parent sets. This process is detailed in Appendix B
Weaknesses of the MOCODA distribution. Although our locally factored dynamics model likely to generalize well on MOCODA, there are a few reasons why training our RL agent on MOCODA in Step 3 may yield poor results. First, if there are empirical imbalances within parent sets (some parent configurations more common than others). these imbalances will appear in MOCODA. Moreover, multiple such imbalances will compound exponentially, so that (s, a) tuples with rare parent combinations will be extremely rare in MOCODA, even if the model generalizes well to them. Second, Support(Mocon)
may be so large that it makes training the RL. algorithm in Step 3 inefficient. Finally, the cost function used in RL algorithms is typically an expectation over the training distribution, and optimizing the agent in irrelevant areas of the state-action space may hurt performance. The above limitations suggest that rebalancing MOCODA might improve results.
MOCODA-U and MOCODA-P. To mitigate the first weakness of MOCODA we might skew MOcoDA toward the uniform distribution over its support, 24 ((lppor((MooOD))) Although this is possible to implement using rejection sampling when l is small, exponential imbalance makes impractical when i large. A more efficient implementation reweights the GMM components used in our MOCODA sampler. We call this approach (regardless of implementation) MOCODA-U. To mitigate the second and third weaknesses of MOCODA, we need additional knowledge about the target task-e.g.. domain knowledge or expert trajectories. We can use such information to define which can also be
prioritized parent distribution MOCODA-P with support in Support(MocoD) obtained via rejection sampling (perhaps on MOCODA-U to also relieve the initial imbalance).
4.2 The Choice of Dynamics Model and RL Algorithm
Once we have augmented parent distribution, (s a), we generate our augmented dataset by applying dynamics model Pa ss, a). The natural choice in light of the discussion in Section 3s a locally factored model. This requires knowledge of the local factorization, which is more involved than the parent set knowledge used to generate the MOCODA distribution and its reweighted variants. We note, however, that locally factored model may not be strictly necessary for MoCoDA, so long as the underlying dynamics are factored. Although unfactored models do not perform well in our experiments,
hypothesize that a good model with enough in-distribution data and the right regularization might learn to implicitly respect the local factorization. The choice of model architecture not core to our work, and we leave exploration of this possibility future work.
Masked Dynamics Model. In our experiments, we assume access to a mask function M : S {0, 1 (((++)))) (perhaps learned (T))) which maps states and actions to the adjacency map of the local graph gc, Given this mask function, we design a dynamics model Ps that accepts M(s, a) as an additional input and respects the causal relations in the mask (i.e., mutual information T(X]; X 1++11 (Se, Ad) \ x;) = 0 if M(:,, az)is = 0). There are many architectures that enforce this constraint. In our experiments we opt for simple one, which first embeds each of the k parent sets: f = [Fa.)a and then computes the j-th child as a function of the sum of the masked embeddings, 9j (M(
f). See Appendix B for further implementation details.
The RL Algorithm. After generating an augmented dataset by applying our dynamics model the augmented distribution, we label the data with our target task reward and use the result to train an RL
7

---

agent. MoCoDA works with wide range of algorithms, and the choice of algorithm will depend on the task setting. For example, our experiments are done an offline setup. where the agent given buffer of empirical data, with no opportunity to explore. For this reason, it makes sense to use offline RI. algorithms, as this setting has proven challenging for standard online algorithms (34.
Remark 4.1. The rationales for (1) regularizing the policy toward the empirical distribution in offline RI. algorithms, and (2) training on the MOCODA distribution, are compatible: in each case, we want to restrict ourselves to state-actions where our models generalize well. By using MOCoDA we expand this set beyond the empirical distribution. Thus, when we apply offline RL algorithms in our experiments, we train their offline component (e.g. the action sampler in BCQ (14] or the BC constraint in TD3-BC (I3)) on the expanded MOCODA training distribution.
5 Experiments
Hypotheses Our experiments are aimed finding support for critical hypotheses:
HI Dynamics models, especially ones sensitive to the local factorization, are able to generalize well the MOCODA distribution.
H2 This out-of-distribution generalization can be leveraged via data augmentation to train an RL agent to solve out-of-distribttion tasks.
Note that support for H2 provides implicit support for H1.
Domains We test MoCoDA on continuous control domains. First is simple, but controlled, 2D Navigation iagatiionnn domain, where the agent must travel from one point in a square arena to another. States are 2D (z, y) coordinates and actions are 2D (Aa, Ay) vectors. In most of the state space, the sub-actions Ax and Ay affect only their respective coordinate. In the top right quadrant, however, the Ax and Ay sub-actions each affect both I and y coordinates, so that the environment is locally factored. The agent has access to empirical training data consisting of left-to-right and bottom-to-tor trajectories that are restricted to a J shape of the state space (see the EMP distribution in Figure 4 We consider target task where the agent must move from the bottom left to the top right. In this task there is sufficient empirical data to solve the task by following the J shape of the data, but learning the optimal policy of going directly via the diagonal requires out-of-disti generalization.
Second, test MoCoDA in a challenging HookSweep2 robotics domain based on Hook-Sweep (32), in which a Fetch robot must use a long hook to sweep two boxes one side of the table (either toward or away from the agent). The boxes are initialized near the center of the table, and the empirical data contains trajectories of the agent sweeping exactly one box one side of the table, leaving the other in the center. The target task requires the agent to generalize to states that it has never seen before (both boxes together on one side of the table). This is particularly challenging because the setup is entirely offline (no exploration), where poor -utto--distribution generalization typically requires special offline RL algorithms that constrain the agent's 's policy to the empirical distribution (34]]
Directly comparing model generalization error. In the 2D Navigation domain we have access to the ground truth dynamics, which allows us to directly compare generalization error on variety of
EMP
DYNA
MOCODA
MOCODA-U
RAND
Figure 4: 2D Navigation Visualization. (Best viewed with 2x zoom) Blue arrows represent transition samples as vector from (Ze, 31) to (Ze+1, 1t+1). Shaded red areas mark the edges of the initial states of empirical trajectories and the center of the square. We see that 5-step rollouts (DYNA) do not fill in the center (needed for optimal policy), and fail to constrain actions those that the model generalizes well For MOCODA, see the effect of compounding dataset imbalance discussed in Subsection 4.1 which is resolved by MOCODA-U.
8

---

Table 1: 2D Navigation Dynamies Modeling Results: Mean squared error + std. dev. over seeds, scaled by le2 for clarity (best model boldfaced). The locally factored model experienced less performance degradation out-of- -i------ittiiutinn, and performed better on all distributions, except for empirical distribution (EMP) itself.
Generalization Error (MSE x1e2) (lower is better)
Model Architecture
EMP
DYNA
RAND
MoCoDA
MOCODA-U
Not Factored
0.14 + 0.04
2.41 +:2 0.29
4.4 1 0.31
0.95 + 0.06
1.29 1 0.15
Globally Factored
0.36 - 0.01
2.09 0.28
3.17 + 0.3
0.41 = 0.02
0.51 0.02
Locally Factored
0.23 - 0.1
1.47 + 0.27
2.03 + 0.19
0.33 = 0.11
0.46 - 0.11
Table 2: 2D Navigation Offline RL Results: Average steps to completion + std. dev. over seeds for various RI. algorithms (best distribution in each row boldfaced), where average steps was computed over the last 50 training epochs. Training MOCODA and MOCODA-U improved performance in all cases. Interestingly, even using RAND improves performance, indicating the importance of training on -uitt--distiiuttion data. Note that this is an offline RL task, and so SAC (an algorithm designed for online RL) is not expected to perform well.
Average Steps to Completion (lower better)
RL Algorithm
EMP
RAND
MoCoDA
MOCoDA-U
CoDA 51
SAC (online RL)
53.1 + 9.8
27.6 1.1
38.8 + 18.3
41.3 17.7
35.1 18.1
BCQ
58.5 + 10.1
31.7 2.4
22.8 1 0.4
24.8 + 4.2
25.0 + 0.4
CQL
45.8 4.0
27.6 + 1.3
22.8 + 0.2
22.7 = 0.3
23.6 + 0.5
TD3-BC
40.0 16.1
26.1 + 0.8
21.0 + 0.7
20.7 = 0.8
21.4 + 0.6
distributions, visualized in Figure 4 We compare three different model architectures: unfactored, globally factored (assuming that the (zr, Az) and (v. Ay) causal mechanisms are independent ev- erywhere, which is not true in the top right quadrant), and locally factored. The models are each trained on a empirical dataset of 35000 transitions for up to 600 epochs, which is early stopped using We find strong support for HI:
a validation set of 5000 transitions. The results are shown in Table I even given the simple dynamics of 2d Navigation, on, is clear that the locally factored model is able to generalize better than fully connected model, particularly on the MOCODA distribution, where performance degradation minimal. We note that the DYNA distribution was formed by starting in EMP and doing 5-step rollouts with random actions. The random actions produce out-of-distributin data to which no model (not even the locally factored model) can generalize well to.
Solving diitiiiiiiiiiiio tasks. We apply the trained dynamics models to several base distribu- tions and compare the performance of RL agents trained on each dataset. To ensure improvements are due to the augmented dataset and not agent architecture, we train several different algorithms, including: SAC [17. BCQ [4] (with DDPG 35)). CQL 31 and TD3-BC [1]
The results on 2D Navigation are shown Table We see that for all algorithms, the use of the MOCODA and MOCODA- U augmented datasets greatly improve the average step count, providing support for H2 and suggesting that using these datasets allows the agents to learn to traverse the diagonal of the state space, even though it is out-of duttof-distribution with respect to EMP. This is consistent with a qualitative assessment of the learned policies, which confirms that agents trained on the J-shaped EMP distribution learn a J-shaped policy, whereas agents trained on MOCODA and MOCODA-U Iearn the optimal (diagonal) policy.
The results on the more complex HookSweep2 environment, shown in Table 3 provide further support for H2. On this environment, only results for BCQ and TD3-BC are shown, as the other algorithms failed on all datasets. For HookSweep2 we used a prioritized MOCODA-P parent distribution, as follows: knowing that the target task involves placing two blocks, we applied rejection sampling to MOCODA to make the marginal distribution of the joint block positions approximately uniform over its support. The effect to have good representation in all areas of the most important state features for the target task (the block positions). The visualization in Figure makes clear why training on MOCODA or MOCODA-P was necessary in order to solve this task; the base EMP distribution simply does not have sufficient coverage of the goal space.

---

Figure 5: HookSweep2 Visualization: Stylized visualization of the distributions EMP (left), MOCODA (center), and MOCODA-P (right). Each figure can be understood as a top down view of the table, where a point is plotted if the two blocks are close together on the table. The distribution EMP does not overlap with the green goal areas on the left and right, and so the agent is unable to leam. In the MOCODA distribution, the agent gets some success examples. In the MOCODA-L P distribution, cnteenactioss are reweighted so that the joint distribution of the two block positions is approximately uniform, leading more evenly distributed coverage of the table.
Table 3: HookSweep2 Offline RL Results: Average success pereentage (+ std. dev. over 3 secds), where the average was computed over the last 50 training epochs. SAC and CQL (omitted) were unsuccessful with all datasets. We see that MOCODA was necessary for learning, and that results improve drastically with MOCODA-P, which re-balances MOCODA toward a uniform distribution the box coordinates (see Figure Additionally, we show results from an ablation, which generates the MoCoDA datasets using fully connected dynamics model. While this still achieves some success, it demonstrates that using a factordd model is important for OOD generalization. In this case the more OOD MoCoDA-P -P distribution does not help, suggesting that the fully connected model is failing to produce useful OOD transitions.
Average Success Rate (higher is better)
RLAlgorithm
MoCoDA (not factored)
MoCoDA-P (not factored)
EMP
MoCoDA
MOCODA-P
BCQ TD3-BC
2.0 + 1.6
20.7 + 4.1
64.7 + 4.1
14.0 = 3.3
15.3 + 4.1
0.7 0.9
38.7 + 7.5
84.0 + 2.8
29.3 + 3.8
26.0 + 1.6
6
Conclusion
In this paper, we tackled the challenging yet common setting where the available empirical data provides insufficient coverage of critical parts of the state space. Starting with the insight that locally factored transition models are capable of generalizing outside of the empirical distribution, we proposed MoCoDA, a framework for augmenting available data using a controllable "parent distribution" and locally factored dynamics model. We find that adding augmented samples from MoCoDA allows RL agents to learn policies that traverse states and actions never before seen in the experience buffer. Although our data augmentation is "model-based", the transition samples it produces are compatible with any downstream RL algorithm that consumes single-step transitions.
Future work might (1) explore methods for learning locally factorized representations, especially in environments with high-dimensional inputs (e.g., pixels) 22 28 , and consider how MOCoDA might integrate with latent representations, (2) combine the insights presented here with learned predictors of out-of- -isdistribution generalization (e.g., uncertainty- baaedd prediction) [46], (3) create benchmark environments for entity-based RL [6 so that object-oriented methods and models can be better evaluated, and (4) explore different approaches to re-balancing the training distribution for learning downstream tasks. With regards to direction (1), note that asserting (or not) certain independence relationships may have fairnsss implications for datasets [7 9 that should be kept in mind or explored. This relevant also regards to direction 4, dataset re-balancing may result in (or fix) biases in the data (30]. Re-balancing schemes should be sensitive to this.
Acknowledgments and Disclosure of Funding
We thank Jimmy Ba, Marc-Etienne Brunet, and Harris Chan for helpful comments and discussions. We also thank the anonymous reviewers for their feedback, which significantly improved the final manuscript. Silviu Pitis is supported by an NSERC CGS-D award. Animesh Garg is supported as CIFAR AI chair, and by an NSERC Discovery Award, University of Toronto XSeed Grant and NSERC Exploration grant. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada, and companies sponsoring the Vector Institute.
10

---

References
[1] Alekh Agarwal, Nan Jiang, Sham M Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. CS Dept., UW Seattle, Seatile, WA, USA. Tech. Rep. 2019.
[2] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, pages 104-114. PMLR, 2020.
[3] Marcin Andrychowicz. Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbecl, and Wojeiech Zaremba. Hindsight experience replay. In Advances in neural information processing systems, pages 5048-5058 2017.
OpenAl: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowiez, Bob McGrew, Jakub Pachocki, Arthur Petron, Mathhia Plappert. Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Roboticx Research, 39.1),220, 2020.
[5] Bharathan Balaji, Petros Christodoulou, Xiaoyu Lu, Byungsoo Jcon, and Jordan Bell- Dell-Maaterson. Factoredrl: Leveraging factored graphs for deep reinforcement learning. In NeurIPS workshop on Offline RL, 2020.
[6] Christopher M Bishop. Mixture density networks. 1994.
[7] Lars Buesing, Theophanc Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean- han-Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactually-guie policy search. International Conference on Learning Representations, 2019.
[8] Guillaume Chaslot, Sander Bakkes. Istvan Szita, and Pieter Spronek. Monte-carlo tree search: A new framework for game ai. In Proceedings of the AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment, volume 4, pages 216-211 217, 2008.
[9] Elliot Creager, David Madras, Toniann Pitassi, and Richard Zemel. Causal modeling for fairnsss in dynamical systems. In Hal Daumé III and Aarti Singh, editors. Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 2185-2195. PMLR, 13-18 Jul 2020.
[10] Alexander D' Amour, Hansa Srinivasan, James Atwood, Pallavi Baljekar, David Sculley, and Yoni Halpern. Fairness is not static: deeper understanding of long term fairness via simulation studies. In Proceedings of the 2020 Conference on Fairness, Accountability, and Transparency, pages 525-534, 2020.
[11] Pieter-Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Y Rubinstein. A tutorial on the cross- entropy method. Annals of operations research, 134(1):19667, 2005.
[12] Andrea Dittadi, Samuele Papa, Michele De Vita, Berhhard Schölkopf, Ole Winther, and Francesco Locatello. Generalization and robustness implications object-centric learning. Proceedings of the international conference on Machine learning, 2022.
[13] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[14] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without explo- ration. In International Conference on Machine Learning, pages 2062. PMLR, 2019.
[15] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
[16] Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificcal Intelligence Research, 19:399-468, 2003.
[17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870, PMLR, 2018
[18] Lily Hu and Issa the 2020 Conference on Fairness, Accountability, and Transparency, 2020.
What's sex got to do with fair machine lcarning? In Proceedings of
[19] Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández Lobato, Clark Glymour, Bernhard Schölkopf. and Kun Zhang. Action-sufficcin state representation learning for control with structural constraints. In International Conference Machine Learning, pages 226099279, PMLR, 2022.
11

---

[20] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, pages 12498-12509, 2019.
[21] Edwin T Jaynes. Information theory and statistical mechanics. Physical review, 106(4):620, 1957.
[22] Jindong Jiang, Sepehr Janghorbani, Gerard De Melo, and Sungjin Ahn. Scalor: Generative world models with scalable object representations. In International Conference on Learning Representations, 2019.
[23] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning, pages 631-673. PMLR, 2018.
[24] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In I/CAI, volume 16, pages 740-747, 1999.
[25] Taylor W Killian, Marzych Ghassemi, and Shalmali Joshi. Counterfactually guided policy transfer in clinical settings. In Conference on Health, Inference, and Learning, pages 5-31, PMLR, 2022.
[26] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 141269880, 2014.
[27] Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pages 2688 2697. PMLR, 2018.
[28] Thomas Kipf, Elise van der Pol, and Max Welling, Contrastive learning of structured world models. In International Conference on Learning Representations, 2020.
[29] Daphne Koller and Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009.
[30] Emmanouil Krasanakis, Eleftherios Spyromitros-Xioute Xiouffs,, Symeon Papadopoulos, and Yiannis Kompatsiaris. Adaptive sensitive reweighting to mitigate bias in fairness-awan classification. In Proceedings of the 2018 world wide web conference, pages 853- 353-862, 2018
[31] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement lcarning. Advances in Neural Information Processing Systems, 33:1 179-1191, 2020.
[32] Andrey Kurenkov, Ajay Mandlekar, Roberto Martin-Martin, Silvio Savarese, and Animesh Garg. Ac-teach: A bayesian actor actor-critic method for policy learning with an ensemble of suboptimal teachers. In Conference on Robot Learning, pages 717-734. PMLR, 2020.
[33] Michael Laskin, Kimin Lec, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas, Reinforce- ment learning with augmented data. In Advances in Neural Information Processing Systems, 2020.
[34] Sergey Levine. Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01666 2020.
[35] Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning, 2016.
[36] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold. Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances Neural Information Processing Systems, 33:1 1525-11555155 2020.
[37] Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan, and Matthew Hausknecht. Working memory graphs. International conference on machine learning, pages 6404 6414. PMLR. 2020.
[38] Chaochao Lu, Biwei Huang, Ke Wang. José Miguel Hernández-Lobatd Kun Zhang, and Bernhard Schölkopf. Sample--fficient reinforcement learning via counterfactual based data augmentation. In NewrIPS Workshop on Offline Reinforcement Learning, 2020.
[39] Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio. Fast and slow lcarning of recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
[40] Ajay Mandlekar, Danfei Xu, Roberto Martín Martín-Martín Silvio Savarese, and Li Fei-Fei. Learning to generalize across long- horizon tasks from human demonstrations, 2020.
12

---

[41] Thomas Mesnard. Theophane Weber, Fabio Viola, Shantanu Thakoor. Alaa Saade, Anna Harutyunyan. Will Dabney, Thomas S Stepleton. Nicolas Heess, Arthur Guez, et al. Counterfactual credit assignment in model-free reinforcement learning. International Conference on Machine Learning, pages 7654-7664. PMLR, 2021.
[42] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 5187754)::52955555 -33322999333233 2015.
[43] Geraud Nangue Tasse, Steven James, and Benjamin Rosman. A boolean task algebra for reinforcement learning. Advances in Neural Information Processing Systems, 33:9497-9507, 2020.
[44] Michael Oberst and David Sontag. Counterfactual off- off-policy evaluation with gumbel-ma structural causal models. In International Conference on Machine Learning, pages 4881-4890. PMLR, 2019.
[45] lan Osband and Benjamin Van Roy. Near-optimal reinforeement learning in factored MDPs. Advances in Newral Information Processing Systems, 27, 2014.
[46] Feiyang Pan, Jia He, Dandan Tu, and Qing He. Trust the model when it is confident: Masked model-based actor-critic, Advances in neural information processing systems, 13:11057-10046, 2020.
[47] Ji Ho Park, Jamin Shin, and Pascalc Fung. Reducing gender bias in abusive language detection. In Conference on Empirical Methods Natural Language Processing, 2018
[48] Judea Pearl. Prohabilities of causation: Three counterfactual interpretations and their identification. Synthese, pages 93-149, 1999.
[49] Judea Pearl. Causality. Cambridge university press, 2009.
[50] Silviu Pitis, Harris Chan, and Stephen Zhao. mrt: modular rl. https:///githuu.commm/ comssptiismmrr 2020.
[51] Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33:3976-3990, 2020.
[52] Martin I. Puterman, Markoy decision processes: discrete stochastic dynamic programming. John Wiley Sons, 2014.
[53] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-35) 2017.
[54] Shagun Sodhani, Sergey Levine, and Amy Zhang. Improving generalization with approximate factored value functions. In ICLR Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
nww/forummm/ het.:aaeeeettoooommddddeeeeeeeee P1d-B4x3BrUUceq
[55] Alexander L Strehl. Model-based reinforcement learning in factored-state mdps. In 2007 IEEE Interna- tional Sympoxium on Approximate Dynamic Programming and Reinforcement Learning, pages 103-110. IEEE, 2007.
[56] Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting, ACM Sigart Bulletin, 264):160-163, 1991.
[57] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018
[58] Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-basee reinforcement learning. In Conference on Robot Learning, pages 1439-145 PMLR, 2020.
[59] Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International conference on learning representations, 2018.
[60] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task- independent state abstraction. International Conference on Machine Learning, pages 23151-23111 PMLR, 2022.
[61] Clemens Winter. Huang Costa. Bamford Chris, and Matricon Theo. Entity gym. https: htap::///gitaan com/ entity-neural wwt.oeeeeeoottteeeeeeeeeee gym 2022.
[62] Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compositional generalization in control. arXiv preprint arXiv:2203. 05960, 2022.
13

---

Checklist
A. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Subsection 4.1
(C) Did you discuss any potential negative societal impacts of your work? [Yes] See Section 6 and Appendix D
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
B. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Appendix
(b) Did you include complete proofs of all theoretical results? [Yes) See Appendix A
C. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] http://gihubbb.oomsspti/t//oto//
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] Yes in Section and in the Appendix C
(c) Did report error bars (e.g.. with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] Yes in Appendix C
D. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cited [00 40 for our RL framework Hook environment assets; both are open-source.
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either the supplemental material or as URL? [Yes] The changes will be included with the code.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
E. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "pEhmm4jz8ae", "reviewer": "Reviewer_1vMv", "summary": "The approach in the paper is to learn a locally-factored dynamics model and the parent distribution of experience data, and use both to generate samples to augment experience data, towards helping the agent reduce sample complexity, and better generalize out of distribution. The paper also provides theoretical arguments justifying why this approach should work, an later considers a couple of tweaks on top of MoCoDa to address some limitations.", "questions": "- Lines 82-84 state an assumption that no two graphical models share the same structure. If the authors could spell out the implications (possible benefits and limitations) of this, and why this assumption is important, it would help readers better understand the framework.", "limitations": "Yes, this has been addressed adequately.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper builds on the work of ref. [34], and provides interesting insights into the theoretical aspects of how a dynamics model helps generalize out-of-distribution with reduced sample complexity. Fig. 4 help understand how the factorized causal model helps overcome implicit data biases that might have otherwise hurt generalization. It is interesting to see how MoCoDa helps generalize beyond the empirical distribution, and across different architectures/algorithms. \n\nThe paper demonstrates an effective implementation of a simple/natural idea, is well-written, and should be of great interest to the audience at this conference.\n\n[34]: Counterfactual data augmentation using locally factored dynamics. NeurIPS 2020\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "jfFOqyFrUg", "reviewer": "Reviewer_155Q", "summary": "This paper shows, both theoretically and experimentally, that a locally factored dynamics model improves sample efficiency and out-of-distribution generalisation when learning a dynamics transition model. In addition, this paper further proposes a Model-based Counterfactual Data Augmentation model to generate counterfactual transitions for RL, allowing offline RL agent to solve an out-of-distribution tasks.", "questions": "Please refer to the \"Weakness\"", "limitations": "Please refer to the \"Weakness\"", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The presentation is very clear and easy to follow.\n2. Despite the fact that many existing works demonstrate that the known factored dynamics model can improve data efficiency, this paper may be the first to demonstrate that the factored dynamics model is also beneficial to out-of-distribution generalization.\n3. The counterfactual data generation is interesting.\n4. The experimental results clearly demonstrate the effectiveness of the proposed framework.\n\nWeakness:\n1. The framework assumes that the factored dynamics model is known. However, the factored model is unknown in most practical cases, and this paper does not propose solutions to this problem, limiting the framework's application scenarios.\n2. This framework only evaluates a simple environment with four dim state spaces. It would be preferable to work in more complex environments, such as mujoco robot control tasks.", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "RkCBML_1Me", "reviewer": "Reviewer_8Yp5", "summary": "This paper proposes a model-based data augmentation approach named MoCoDa. The approach exploits the local factorization in state-action space, inspired by causality and factored MDP literature, to generate transitions for policy learning. The authors provided the theoretical analysis of sample complexity on the local factored model, and meanwhile, the empirical results on two continuous control benchmarks verify the hypotheses made by the main paper. ", "questions": "**Please note that running any of the mentioned ideas or datasets is not a must during the rebuttal phase. Some discussion is highly appreciated.**\n\n#### **1. About the entangled long horizon structure**\n\nThe authors said the long horizon prediction could no longer leverage the sparsity of causal graphs. I can roughly understand what the authors meant here. However, suppose we can learn a good disentangled representation of \"decision-making-relevant\" and \"decision-making-irrelevant\" state decomposing in the graph. In that case, we can leverage the sparsity of the causal graphs. Similar ideas have been proposed in [1], where they also learn the connections between states and rewards using masks. And those variables (states) having direct connections to future rewards are decision-making-relevant. Then we can leverage these states for down-streaming policy learning.\n\n#### **2. Graphs in long horizon structure**\n\nThe authors said that temporal entanglement might happen. Why not consider using the non-stationary graphs with temporal changes to describe the interaction skeletons? \n\n#### **3. More complicated evaluation senarios**\n\nCan the proposed method be applied to other complicated OOD generalization tasks, where the distribution shift is more natural and practical? e.g., Procgen Benchmark [2], different Atari games, etc. \n\n\n#### **References**\n\n[1] Huang, Biwei, et al. \"Action-sufficient state representation learning for control with structural constraints.\" International conference on machine learning, 2022.\n\n[2] Cobbe, Karl, et al. \"Leveraging procedural generation to benchmark reinforcement learning.\" International conference on machine learning. PMLR, 2020.\n", "limitations": "Most of the limitations are mentioned in the weakness parts. I think the authors can improve the work by giving more clear clarifications on causality contexts. I think the paper is a good submission, and I will increase my score if the authors provide supporting analysis on the above weaknesses. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## *Strengths*\n\n#### **1. Technical soundness and significance** \n\nIn general, the paper delivers a novel approach, inspired by causality and factored MDP, to effectively generate transition samples for RL. The method is technically sound and would potentially benefit the fields of RL, causal representation learning, etc. \n\n#### **2. Empirical verification** \n\nThe paper includes two evaluation tasks on continuous control domains. The results indicate that the model can learn the underlying dynamic process and improve the OOD generalization on novel tasks. The quantitative results (Table 1-3) and visualization (Fig. 4-5) can verify the effectiveness. \n\n#### **3. Presentation**\n\nThe overall presentation is clear and easy to follow, although it can be further improved by addressing a few minor points (mentioned below in the weakness sections). The appendix supports the main paper and provides most of the required details. \n\n## *Weaknesses*\n\n#### **1. About causality**\n\n##### *1.1. The concept of \"counterfactual\" is vague in the paper*\n\nUsually, in the context of causality, counterfactual reasoning means you must reconstruct the exogenous noise factors in the data generating process (see Pearl's Causal ladder in [1-2]). While in this paper, some critical and crucial points are missing. It would be highly appreciated if the authors could make some discussion and analysis (if any) in the rebuttal and final version of this paper. The authors can refer to one of the existing works on data augmentation using counterfactual learning for RL [3].\n\n##### *1.2. Identifiability of the (learned or given) causal graphs*\n\nSimilarly, we usually need to guarantee that the graph is identifiable in the causality literature. It would be better to give proof or analysis of identifiability since the authors assume the given or learned graphs are essentially causal graphs. \n\n#### **2. Missing comparisons with similar approaches**\n\nSome other approaches [4-10] also aim to learn the factored masks for RL to improve the sample efficiency or few/zero-shot generalization. It would be better for the authors to discuss them in the updated version. **Please note that running them in the rebuttal phase is not a must. The discussion and analysis would be highly appreciated in the rebuttal.**\n\n#### **3. Missing references**\n\n**Disclaimer: All mentioned works are critical and representative works in different related areas, and none of them is done by me.**\n\n##### *3.1 Factored MDP for deep RL: [4-10]*\n\n#### *3.2 Causal representation learning: [11]*\n\n#### *3.3 Sparse/independent  (causal) mechanisms: [12-13]*\n\n#### *3.4 Causal discovery in dynamic systems [14-15]*\n\n\n#### **4. Presentations**\n\n##### 4.1 *About CoDA* \n\nSince the work is built upon CoDA, it is better to have a short background or preliminary section on CoDA (about motivation, model, and important remarks/conclusions) and the relevance and differences with MoCoDA. I noticed that the authors had provided analysis in Section 4, but giving them in the early sections will help authors to understand these two works better. \n\n##### 4.2 *About masked dynamic models* \n\nThe current content on deriving the masks is vague. The authors can consider moving some contents in Appendix B.2 (well-written and easy to understand) into the main paper for better readability.\n\n##### 4.3 **Adding more related works sections** \n\nAdding more related works sections on causality (structural causal model) and causal RL, data augmentation for RL, and OOD generalization for RL would be helpful.\n\n##### 4.4 **Minor typos**\n\na. Line 43: they too learn to generalize ... [check grammar]\n\nb. Line 237: missing the right bracket\n\n\n\n\n\n\n#### **References**\n\n[1]. Pearl, Judea. Causality. Cambridge university press, 2009.\n\n[2]. Powell, Stephen. \"The Book of Why: The New Science of Cause and Effect. Pearl, Judea, and Dana Mackenzie. 2018. Hachette UK.\" Journal of MultiDisciplinary Evaluation 14.31 (2008): 47-54.\n\n[3] Lu, Chaochao, et al. \"Sample-efficient reinforcement learning via counterfactual-based data augmentation.\" NeurIPS workshop on Offline RL (2020).\n\n[4] Hallak, Assaf, et al. \"Off-policy model-based learning under unknown factored dynamics.\" International Conference on Machine Learning. PMLR, 2015.\n\n[5] Balaji, Bharathan et al. Factoredrl: Leveraging factored graphs for deep reinforcement learning. NeurIPS 2020 Deep Reinforcement Learning Workshop, 2020.\n\n[6] Loynd, Ricky, et al. \"Working memory graphs.\" International conference on machine learning. PMLR, 2020.\n\n[7] Wang, Tingwu, et al. \"Nervenet: Learning structured policy with graph neural networks.\" International conference on learning representations. 2018.\n\n[8] Artem Zholus, Yaroslav Ivchenkov, and Aleksandr Panov. Factorized world models for learning causal relationships. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.\n\n[9] Shagun Sodhani, Sergey Levine, and Amy Zhang. Improving generalization with approximate factored value functions. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.\n\n[10] Huang, Biwei, et al. \"Action-sufficient state representation learning for control with structural constraints.\" International conference on machine learning, 2022.\n\n[11] Schölkopf, Bernhard, et al. \"Toward causal representation learning.\" Proceedings of the IEEE 109.5 (2021): 612-634.\n\n[12] Goyal, Anirudh, et al. \"Recurrent independent mechanisms.\" ICLR 2021.\n\n[13] Madan, Kanika, et al. \"Fast and slow learning of recurrent independent mechanisms.\"  ICLR 2021.\n\n[14] Bongers, Stephan, Tineke Blom, and Joris Mooij. \"Causal modeling of dynamical systems.\" arXiv preprint arXiv:1803.08784 (2018).\n\n[15] Yao, Weiran, Guangyi Chen, and Kun Zhang. \"Learning Latent Causal Dynamics.\" arXiv preprint arXiv:2202.04828 (2022).\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "JwZHzkINsST", "reviewer": "Reviewer_FUqL", "summary": "This paper presents a novel model-based framework, MOCODA, as a generalization of the CODA framework for augmenting the training dataset of an offline-RL agent with samples generated using the parent distribution and locally factored dynamics, to solve zero-shot out-of-distribution (OOD) tasks. MOCODA is evaluated on continuous control tasks from the 2D Navigation and the HookSweep2 environments.", "questions": "- Why are there no CODA baselines? If I understand correctly, it makes sense to compare with CODA, right? And if not, it will be useful to have this addressed in the paper.\n- One of the advantages of MOCODA is that it is applicable even when the parent sets overlap: are there cases where they do and does MOCODA perform well there? \n\n\nI am happy to discuss and increase my score if these questions are answered. ", "limitations": "Yes, the authors discuss the limitations of their work-- especially its potential inefficiency in training RL agents and suggest mitigating this by rebalancing MOCODA. The authors also discuss the societal impact of their work with respect to the handling of the causal assumptions made in this paper-- it would be good to see these listed out explicitly.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths*\n\nIt's interesting to see the usage of locally factored dynamics models towards generalizing to unseen areas of the state space, and it does seem to reduce generalization error on some easy tasks.\n\n*Weaknesses*\n\n- It's not clear to me why the authors do not compare their method with CODA when they have presented SAC as an online RL model-free baseline.\n- The authors say MOCODA works even when the parent sets overlap. Some analysis of the parent sets and the parent distribution would be useful, as it is not very clear right now why and how this is an advantage over CODA right now. (I have read B.2)\n- I think the paper's clarity would greatly benefit from a clear pseudocode of their entire framework.\n- The authors do not present any results to show the significance of using locally-factored models on the HookSweep2 tasks, which are considerably more difficult than the 2D Navigation tasks, and it's not clear why not.\n- For improving the quality of the paper from the perspective of causal modeling, I think it would be helpful to have the causal assumptions made listed down specifically.\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept