Paper: Bad Habits: Policy Confounding and Out-of-Trajectory Generalization in Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
The theoretical analysis introduces useful formalisms but is very dense. For readers less familiar with this notation, could you provide more intuitive explanations of key results like Proposition 1 and Theorem 1?

Author response:
We acknowledge that the theory may appear somewhat intimidating. We have tried to enhance accessibility by providing numerous examples and figures. Unfortunately, defining the phenomenon mathematically requires introducing several new notations and definitions. 

Proposition 1 simply demonstrates that certain policies may induce state representations that require fewer observation variables than the minimal (Markov) state representation (Definition 5) while still being $\pi$-Markov (Markov for the states visited when following such policies; Definition 7). This is somewhat obvious if you consider policies that visit only a subset of the state space—think of a policy that makes the agent stay put. However, what is not so obvious is that this occurs even for well-performing (optimal) policies. The underlying reason is what we term policy confounding.  The policy introduces spurious correlations between past and future observation variables. This is illustrated by Example 1 and demonstrated by Theorem 1.

Reviewer comment (Writing):
The distinction between out-of-trajectory and out-of-distribution generalization is somewhat unclear. Could you clarify this difference with explicit examples?

Author response:
Certainly! Out-of-trajectory (OOT) generalization is a particular instance of the more general out-of-distribution (OOD) generalization problem in RL. The objective of out-of-trajectory generalization is not to generalize to environments with different rewards and transitions (as is the case in all the previous works that target generalization in RL) but simply to alternative trajectories within the same environment. 

Consider a robot that is trained to go from your office to the coffee machine and back, and from your office to the printer and back. There are two different ways to go to the coffee machine: either through the printer room or through a corridor that leads directly to the coffee machine. The path through the printer room is longer, so the robot never takes it when you order coffee. However, one day you order coffee, but the corridor is blocked. Hence, the robot tries to go through the printer room and comes back with a copy of a new paper titled "Bad Habits" instead of the coffee. This is an example of out-of-trajectory generalization. Since the robot is used to getting copies when in the copy room, it ignored the order you gave it. An example of the more general problem of OOD generalization could be to ask the robot to navigate your office when the floor is wet or to order something different, such as a glass of water. The key difference is that in these two examples, the states the robot visits or the rewards it receives are different. The robot has never been trained on a wet floor, nor has it ever gotten you a glass of water before. However, in the first example, we would expect the robot to be aware that being in the copy room does not imply getting copies. To be fair, the blocked corridor represents a change in the environment; however, this change is intended to prompt the agent to choose an alternative path. Note that this alternative path was also viable in the original environment.

Reviewer comment (Novelty):
How does policy confounding relate to prior work on spurious correlations and generalization in RL? Are there clear differences in causes and solutions?

Author response:
Our response to the previous comment addresses the differences concerning prior works on generalization in RL. As for works investigating spurious correlations, we are only aware of the paper by Langosco et al. (2022). The key difference in Langosco’s paper is that spurious correlations are already present in the environment. The setting is as follows: a certain object is always next to the actual goal, so the agent learns to associate the object with the goal. However, in the test environments, the object is no longer present. In our case, it is the policy itself that introduces the spurious correlations, such as always getting copies when at the printing room. As for solutions, some of the methods that work for the general RL generalization problem would likely mitigate the OOT generalization problem, such as domain randomization (Tobin et al., 2017; Peng et al., 2018; Machado et al., 2018) and exploration (Eysenbach & Levine, 2022; Jiang et al., 2022). However, we believe that coming up with more targeted strategies is crucial.

Reviewer comment (Theory):
You cite causal representation learning as a promising direction for future work, but do not provide specifics on how these techniques could be applied to address policy confounding. Could you expand on how invariant risk minimization or other causal inference tools could help mitigate the effects you demonstrate? Are there any concrete steps you propose for integrating causal representations into solving this problem?

Author response:
Although this is beyond the scope of this paper, we are happy to provide some hints regarding what we believe are the most promising strategies. One option would be to exploit the notion of prediction invariance (Peters et al., 2016; Arjovsky et al., 2019). Since state representations that are robust to trajectory deviations must be invariant across policies, one could enforce this property by modifying the objective function (Zhang et al., 2020). Another alternative is to actively encourage the agent to stress-test its representations in the environment by trying out alternative trajectories. This could be achieved by adding intrinsic rewards (Sontakke et al., 2021). We would be happy to include this information in the appendix if the reviewer deems it important.

Reviewer comment (Writing):
In Example 1 (Section 6), X_t and x_t seem to not be defined when they are first used. Later, it is stated that X_0 is the signal received at t=0, but even then, it is not clear how X differs from G. These should be better defined.

Author response:
> **In Example 1 (Section 6), X_t and x_t seem to not be defined when they are first used. Later, it is stated that X_0 is the signal received at t=0, but even then, it is not clear how X differs from G. These should be better defined.**

Thanks for pointing this out. The reviewer’s interpretation is correct. $X$ represents the signal which takes the same value as $G$ at the start location. At any other location, $X$ takes a dummy value. We will clarify this in the text.

Reviewer comment (Theory):
The do operator is not well-defined. Based on the informal description (just below Definition 9), I thought that it meant that there exists an s_t, such that we consider all possible values of Phi(s_t), and adversarially pick one to try to meet one of the conditions in Definition 9 if possible (even if this is the correct interpretation, it is not clear). However, based on a few readings for Section 6.0, I suspect the correct interpretation is something different.

Author response:
> **The do operator is not well-defined.**

The do-operator is commonly used in the causal inference literature to distinguish cause-effect relationships from mere statistical associations (Pearl, 2000). We define do-operator in the text below Definition 9: $\text{do}(\Phi(s_t))$ means setting the variables forming the state representation $\Phi(s_t)$ to a particular value and considering all possible states in the equivalence class, $s'_t \in \{s_t\}^\Phi$ (i.e., all states that share the same value for the state variables that are "selected” by $\Phi$). We understand this definition may be hard to parse and will add the following example below the definition for clarification:

In the T-maze environment, $R^{\pi^*}(\text{do}(L_8 = \text{green goal})) \neq R^{\pi^*}(L_8 = \text{green goal})$, since $R^{\pi^*}(\text{do}(L_8 = \text{green goal})) = \pm1$ and $R^{\pi^*}(L_8 = \text{green goal}) = +1$. This is already explained in the sentence before Definition 9. However, we will move it after the definition to make use of the $\text{do}(\cdot)$ operator.

Reviewer comment (Experiments):
Justifying/motivating the work: The examples given focus on the idea that we train on one MDP, and then evaluate on another. This approach illustrates the issue well, but is too contrived to provide motivating examples that show why we should care about policy confounding in practice (and the experiments are all set up this way as well). Section 6.2 attempts to address this by discussing when we should worry in practice: function approximation and narrow trajectory distributions. However, these two paragraphs are far too short and high-level to truly justify this work. I believe that this is an interesting topic worthy of study, but I think this paper in its current form does a poor job of showing the reader that this is an interesting topic worthy of study. Perhaps showing that policy confounding can be a problem in less contrived settings (for example, when the MDP or distribution of MDPs does not shift between training and evaluation) would help address this weakness. Alternatively, another approach could be to focus a future version of the paper on a setting where the training and evaluation MDPs are inherently different, such as "sim-to-real" robotics.

Author response:
> **The examples given focus on the idea that we train on one MDP, and then evaluate on another. (...) showing that policy confounding can be a problem in less contrived settings (for example, when the MDP does not shift between training and evaluation) would help address this weakness.**

As explained in the introduction the objective of OOT generalization is not to generalize to other MDPs (as is the case in the standard RL generalization problem). The objective is to generalize to different trajectories within the same environment. In order to test this we must force the agent to take alternative trajectories. The only way to do so is by modifying the MDP. However, as explained in the last paragraph of section 7.2: "these alternative trajectories are both possible and probable in the original environments, and thus one would expect well-trained agents to perform well on them."

> **6.2 attempts to address when we should worry about policy confounding in practice. However, these two paragraphs are far too short and high-level to truly justify this work.**

We do not intend to justify our work on these two paragraphs. These two paragraphs are meant to clarify when the phenomenon of policy confounding described by all the previous theoretical results and examples is more likely to occur in practice. They also bridge the gap between the theoretical sections and the experiments.

Reviewer comment (Novelty):
The main theorem seems almost trivial, and the empirical work is based on extremely simple toy gridworlds. So even aside from the issues above, the contribution may be a bit light.

Author response:
> **The main theorem seems almost trivial.**

The theorem may seem trivial after having illustrated the phenomenon with an example and having provided the definition of policy confounding. However, it is certainly not trivial that the underlying reason why in some situations agents learn representations that are not fully Markov is policy confounding.

## Rebuttal 2 (ID: YRxTjxv3pW)
>  **The empirical work is based on extremely simple toy gridworlds.**

The environments are meant as pedagogical examples to illustrate the problem in practice and to clarify the ideas introduced by the theory. Our main contribution is the characterization of policy confounding. 

That being said, the fact that standard RL methods fail already in such simple settings gives great cause for concern. Moreover, we have found plenty of empirical evidence in prior works of particular forms of policy confounding in high dimensional domains (Machado et al., 2018; Zhang et al., 2018a; Zhang et al., 2018b; Song et al., 2020; Nikishin et al., 2022; Lan et al., 2023). A review of all these works along with explanations on how the phenomena they highligt are particular instances of policy confounding is provided in Appendix C.

Reviewer comment (Experiments):
The primary weakness is that the experiments focus on toy examples specifically designed to elicit the problem. It's not entirely clear if this problem is relevant in more realistic settings. Appendix C does contain examples of failures from previous works which may be related to policy confounding though.

Author response:
We would like to point out that the environments are meant as pedagogical examples to illustrate the problem in practice and to clarify the ideas introduced by the theory. Our main contribution is the characterization of policy confounding. Moreover, the fact that standard RL methods fail already in such simple settings gives great cause for concern.

Reviewer comment (Evaluation):
T-Maze, sec2. It could be more clear to mention why the avg return is around 0.2. Is this around the level that is to be expected? Why is there a little bump around step 30k for the eval env?

Author response:
0.2 is the average return over 10 runs. It is hard to expect anything from neural networks when being evaluated out-of-distribution. In some of the runs, the agent is stuck doing nothing for a while and then goes to one of the two goals. In some others, the agent simply goes to the wrong goal.

Reviewer comment (Theory):
About the formal definition (def 9) of policy confounding: It would be helpful to understand this definition by explaining how it applies to the T-maze. Could you clarify the meaning of the do() operator here from a mathematical point of view? Would it be the same as writing that the equality has to hold for all $s_t$ s.t. $\phi(s_t)$ rather than only the ones visited by $\pi$?

Author response:
In the T-maze environment, $R^{\pi^*}(\text{do}(L_8 = \text{green goal})) \neq R^{\pi^*}(L_8 = \text{green goal})$, since $R^{\pi^*}(\text{do}(L_8 = \text{green goal})) = \pm1$ and $R^{\pi^*}(L_8 = \text{green goal}) = +1$. This is already explained in the sentence before Definition 9. However, we will move it after the definition to make use of the $\text{do}(\cdot)$ operator.

That is right! The equality should hold for all states in $s_t$'s equivalence class under $\Phi$ (i.e., all states that share the same value for the state variables that are "selected” by $\Phi$).

Reviewer comment (Experiments):
About "Narrow trajectory distributions". If the environment transitions leads to a diverse set of states, would we still observe policy confounding? What if the agent is incentivized to explore through exploration bonuses? Could this avoid the issue?

Author response:
That depends on how diverse the set of states is and how frequently each state is visited. In general, if the environment is random enough such that the agent visits all states sufficiently frequently, we would most likely avoid policy confounding. The same goes for the policy. If the policy is stochastic enough such that the agent visits all states sufficiently frequently, the issue may be avoided. So, yes, exploration does help! Both domain randomization and exploration are discussed in Section 6.

Reviewer comment (Theory):
The paper makes a link between policy confounding and causality (or lack of it). Are there any tools from causality that could be used to address the problem?

Author response:
Not directly as far as we know. The question of how tools from causal inference can be adapted to prevent policy confounding will be investigated in future work. However, one promising option would be to exploit the notion of prediction invariance (Peters et al., 2016; Arjovsky et al., 2019) since state representations that are robust to trajectory deviations are those that are invariant across policies.