Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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.
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. 

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

[34]: Counterfactual data augmentation using locally factored dynamics. NeurIPS 2020

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.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.
strengths_and_weaknesses: Strengths:
1. The presentation is very clear and easy to follow.
2. 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.
3. The counterfactual data generation is interesting.
4. The experimental results clearly demonstrate the effectiveness of the proposed framework.

Weakness:
1. 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.
2. 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.
questions: Please refer to the "Weakness"
limitations: Please refer to the "Weakness"
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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.
strengths_and_weaknesses: *Strengths*

It'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.

*Weaknesses*

- 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.
- 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)
- I think the paper's clarity would greatly benefit from a clear pseudocode of their entire framework.
- 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.
- 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.

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


I 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.
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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. 
strengths_and_weaknesses: ## *Strengths*

#### **1. Technical soundness and significance** 

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

#### **2. Empirical verification** 

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

#### **3. Presentation**

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

## *Weaknesses*

#### **1. About causality**

##### *1.1. The concept of "counterfactual" is vague in the paper*

Usually, 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].

##### *1.2. Identifiability of the (learned or given) causal graphs*

Similarly, 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. 

#### **2. Missing comparisons with similar approaches**

Some 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.**

#### **3. Missing references**

**Disclaimer: All mentioned works are critical and representative works in different related areas, and none of them is done by me.**

##### *3.1 Factored MDP for deep RL: [4-10]*

#### *3.2 Causal representation learning: [11]*

#### *3.3 Sparse/independent  (causal) mechanisms: [12-13]*

#### *3.4 Causal discovery in dynamic systems [14-15]*


#### **4. Presentations**

##### 4.1 *About CoDA* 

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

##### 4.2 *About masked dynamic models* 

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

##### 4.3 **Adding more related works sections** 

Adding more related works sections on causality (structural causal model) and causal RL, data augmentation for RL, and OOD generalization for RL would be helpful.

##### 4.4 **Minor typos**

a. Line 43: they too learn to generalize ... [check grammar]

b. Line 237: missing the right bracket






#### **References**

[1]. Pearl, Judea. Causality. Cambridge university press, 2009.

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

[3] Lu, Chaochao, et al. "Sample-efficient reinforcement learning via counterfactual-based data augmentation." NeurIPS workshop on Offline RL (2020).

[4] Hallak, Assaf, et al. "Off-policy model-based learning under unknown factored dynamics." International Conference on Machine Learning. PMLR, 2015.

[5] Balaji, Bharathan et al. Factoredrl: Leveraging factored graphs for deep reinforcement learning. NeurIPS 2020 Deep Reinforcement Learning Workshop, 2020.

[6] Loynd, Ricky, et al. "Working memory graphs." International conference on machine learning. PMLR, 2020.

[7] Wang, Tingwu, et al. "Nervenet: Learning structured policy with graph neural networks." International conference on learning representations. 2018.

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

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

[10] Huang, Biwei, et al. "Action-sufficient state representation learning for control with structural constraints." International conference on machine learning, 2022.

[11] Schölkopf, Bernhard, et al. "Toward causal representation learning." Proceedings of the IEEE 109.5 (2021): 612-634.

[12] Goyal, Anirudh, et al. "Recurrent independent mechanisms." ICLR 2021.

[13] Madan, Kanika, et al. "Fast and slow learning of recurrent independent mechanisms."  ICLR 2021.

[14] Bongers, Stephan, Tineke Blom, and Joris Mooij. "Causal modeling of dynamical systems." arXiv preprint arXiv:1803.08784 (2018).

[15] Yao, Weiran, Guangyi Chen, and Kun Zhang. "Learning Latent Causal Dynamics." arXiv preprint arXiv:2202.04828 (2022).

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

#### **1. About the entangled long horizon structure**

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

#### **2. Graphs in long horizon structure**

The authors said that temporal entanglement might happen. Why not consider using the non-stationary graphs with temporal changes to describe the interaction skeletons? 

#### **3. More complicated evaluation senarios**

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


#### **References**

[1] Huang, Biwei, et al. "Action-sufficient state representation learning for control with structural constraints." International conference on machine learning, 2022.

[2] Cobbe, Karl, et al. "Leveraging procedural generation to benchmark reinforcement learning." International conference on machine learning. PMLR, 2020.

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.
review_title: null

Meta-review: metareview: The paper suggests to improve sample efficiency and out-of-distribution generalization in RL by learning locally factored world models, and use these models to generate counterfactual data to train on. The key assumption is that the environment model is the right model to factorize (as opposed to, say a policy or value function) and that this model will generalize out of distribution when performing the relevant interventions. All reviewers were in agreement the paper was well written and presented an interesting idea with sound empirical verification. Several comments pointed to an unclear definition of 'counterfactual' used in the paper (as the authors point out, it means different things depending on whether adopting a potential outcome or DAG framework) - please make sure this is clear in the final version, as well as a clear explanation of the distinction with CODA.
confidence: Certain
Decision: Accept