Paper: Robust agents learn causal world models (ICLR-2024)

Reviewer comment (Theory):
Only necessary condition is proved but not the sufficient condition. It will be stronger to prove something like, if the agent has learned some "approximate" causal relationship, it can efficiently learn the optimal decision under distribution shift.

Author response:
We thank the reviewer for their encouraging review and helpful comments. 

As the reviewer has requested, we have derived a new theorem which proves the sufficiency condition corresponding to the necessity condition proven in the original submission. This new theorem is included in the updated manuscript (Theorem 3).

Reviewer comment (Theory):
How to identify and prove the sufficient condition for learning the causal model for learning the optimal decision making under distribution shift?

Author response:
We thank the reviewer for their encouraging review and helpful comments. 

As the reviewer has requested, we have derived a new theorem which proves the sufficiency condition corresponding to the necessity condition proven in the original submission. This new theorem is included in the updated manuscript (Theorem 3).

Reviewer comment (Experiments):
They do not conduct an experiment for justifying their results.

Author response:
We thank the reviewer for their helpful comments, and have made the following changes to the manuscript in response to their recommendations

1. Experiments section (appendix F.1, end of appendices). We now empirically validate our results using a simulation example as suggested by the reviewer. This involves converting the proof of theorem 1 into an algorithm for learning the underlying CBN from the agents policy under distributional shifts, and testing it on randomly generated CIDs. We also explore how the accuracy of the learned CBN scales with the agent’s regret bound.

Reviewer comment (Theory):
Their results can only be applied to a small range of scenarios, where we need to reach small regret for all mixture of local interventions. However, most applicable tasks, such as transfer learning, only consider interventions on a subset of variables.

Author response:
We agree our results would be stronger and more applicable to current systems if we could extend theorems 1 & 2 to shifts on a small subset of environment variables. Unfortunately, this would be a significant extension of the current proofs and we were unable to do this extension in the time available. We agree this is an important direction for future work.

Reviewer comment (Writing):
The statement of the result in the introduction felt pretty slippery; although this was resolved in later sections, the paper would benefit from sharpening that first section.

Author response:
We want to very much thank the reviewer for their encouraging review and helpful comments. We have updated the introduction to be clearer and in line with the reviewers suggestions. On top of this we have also introduced a simplified overview of the proof and an experiments section (appendix F).

Reviewer comment (Writing):
The terms "distributional shifts", "distributionally shifted environments", "target domains", and "causal modelling and transfer learning" should be explicitly defined.

Author response:
"Please define these ...". We have added definitions for these terms in sections with the corresponding titles, and removed the term `causal modelling’. If the reviewer would prefer definitions in the form e.g. Definition 1 (Bayesian networks) we will include these.

Reviewer comment (Writing):
There is a typo in the phrase "used to derive out results".

Author response:
Typos fixed

Reviewer comment (Presentation):
The sentence "Our analysis focuses on distributional shifts that involve changes to the causal data generating process, and hence can be modelled as interventions (Schölkopf et al., 2021)" would have been nice to include in the introduction.

Author response:
Have clarified in the introduction that we are focusing on interventional shifts.

Reviewer comment (Writing):
The sentence "This does not assume that all shifts an agent will encounter can be modelled as interventions, but requires that the agent is at least capable of adapting to these shifts" is unclear.

Author response:
"This does not assume that all shifts... I don't know that I understand this sentence". Apologies, this was confusingly written. What we are trying to say is, imagine an agent that is robust to `all’ distributional shifts (including those that cannot be represented with local interventions, such as changing the set of variables). This agent will also be subject to our theorems, because it is robust to all distributional shifts which includes interventional shifts as a subset. So our results apply to agents that are robust to a larger set of shifts. Will remove if the reviewer recommends.

Reviewer comment (Writing):
There is a typo in "By cCreftheorem: main,theorem: main approx agents".

Author response:
Typos fixed

Reviewer comment (Writing):
The reviewer would like to hear what changes to the introduction might look like.

Author response:
We want to very much thank the reviewer for their encouraging review and helpful comments. We have updated the introduction to be clearer and in line with the reviewers suggestions. On top of this we have also introduced a simplified overview of the proof and an experiments section (appendix F).

Reviewer comment (Experiments):
As the authors acknowledge, the results are mainly theoretical. Even a minimal empirical validation of the key insights would strengthen the paper. For example it would be great even if you turn the informal overview (appendix C) into a simple simulation example rather than remain a thought experiment.

Author response:
1. Experiments section (appendix F.1, end of appendices). We now empirically validate our results using a simulation example as suggested by the reviewer. This involves converting the proof of theorem 1 into an algorithm for learning the underlying CBN from the agents policy under distributional shifts, and testing it on randomly generated CIDs. We also explore how the accuracy of the learned CBN scales with the agent’s regret bound.

Reviewer comment (Experiments):
The scope is currently limited to unmediated decision tasks. Extending the results to broader RL settings would increase applicability (although I acknowledge that seems significantly more challenging task and out of scope of this work - it’s just a personal curiosity at this point and would be excited to see the next paper already).

Author response:
We tried to extend the results to unmediated decision tasks, but unfortunately this was too challenging in the time available. We agree this is an important direction for future work.

Reviewer comment (Theory):
On a similar note, the implications of the assumptions are not discussed. (e.g. I’d like to see things like, “assumption 2 implies that there exist distribution shifts for which the optimal policies are different”.

Author response:
3. Discussed implications of assumptions immediately after their statement and suggested.

Reviewer comment (Writing):
"The environment is described by a set of random variables C..." this sentence belongs to the main text since it you don't explain C random variable although it's heavily used.

Author response:
4. Added missing description of chance nodes “C” on page 3.

Reviewer comment (Presentation):
Although discussed in the appendix, i'd like to see the description of what squares, circles and diamonds mean in the CID.

Author response:
5. Added description of CID nodes (circles, squares diamonds) to caption of figure 1.

Reviewer comment (Theory):
In assumption 1 you stated $\text{Desc}_D \cap \text{Anc}_U = \emptyset$ but this doesn't exclude the trivial setting (which you state in the appendix is not of focus). Can you either extend the assumption or comment in the main text that it's not of interest the trivial setting? (i know it's a nitpick but got me wondered while reading it in the main text and i feel that since you thought about it you could have mentioned it earlier in the text).

Author response:
7. The reviewer also points out that it was not made clear how our assumptions deal with the trivial case (i.e. where the agent’s action does not cause the utility). In this case, all policies are optimal, and hence the assumption of domain independence is violated (there is a single policy that is optimal under all domain shifts). We have added a note making this clear.

Reviewer comment (Theory):
Can you please give a simple sketch of the proof? This would help the readability significantly. Also i feel there is a simple sentence that can be written on each theorem that explains its implications.

Author response:
6. Added sentence description immediately after / before each theorem.