Bad Habits: Policy Confounding and Out-of-Trajectory Generalization in Reinforcement Learning

We appreciate the time the reviewer invested in reviewing our work. We hope our answers below can help clear up the reviewer’s concerns.

The paper argues policy confounding poses a distinct challenge from general RL generalization, but does not extensively benchmark against recent generalization methods like invariant risk minimization, data augmentation, dynamics randomization, robust policy learning, and meta RL. These approaches could be highly relevant given the claims about out-of-trajectory generalization.

We would like to emphasize that this is not an experimental paper. As the reviewer correctly points out, the goal of this paper is to formalize the concept of policy confounding and create awareness about its risks. We demonstrate that the issue occurs when using two of the most popular RL methods: PPO and DQN. Whether the issue manifests itself in each specific RL method is beyond the scope of this paper. Moreover, the RL methods the reviewer mentions (robust policy learning and meta RL) target other types of issues that are not directly related to policy confounding, so there is no a priori reason why these methods would help mitigate its effects. As for other generalization techniques (such as invariant risk minimization or data augmentation), it is not clear how these can be adapted to the RL setting. We do show that dynamics randomization can help mitigate its effects (last paragraph, Section 7.2).

Generally, we observe the tendency in RL research to run numerous experiments and then draw conclusions based on a subset of the results. This practice is inefficient in terms of time and energy. In this paper, our approach has been to first analyze the issue analytically and then conduct targeted experiments to validate that they indeed support the proposed theory.

While theory and simple experiments demonstrate policy confounding, how confident are you this poses a problem not already addressed by the latest techniques for improving generalization in RL? Comparisons to some of these state-of-the-art methods could better situate your claims within the broader context of research on robustness and generalization.

As mentioned in the previous response the paper highlights a new, and unexplored problem, in RL. Whether or not the latest techniques for improving generalization in RL inadvertently solve this problem does not make it less of a problem. We believe the RL community should still be warned about the risks that policy confounding poses.

You demonstrate policy confounding in simple domains. Do you have evidence this manifests in more complex, high-dimensional problems? Testing on complex benchmarks could better showcase significance.

We did not experiment on high-dimensional domains. Our environments are meant as pedagogical examples to illustrate the problem in practice and to clarify the ideas introduced by the theory. However, 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.

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?

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.

What directions seem most promising for future work to address policy confounding?

As previously mentioned, our main contribution is the characterization of policy confounding. The proposed solutions, rather than serving as definitive solutions, aim to exemplify instances when the phenomenon is most likely to occur. Regarding what directions seem most promising for future work, we think that a complete reformulation of the RL problem is necesary. The agent should not only focus on maximizing rewards but should also prioritize learning robust state representations grounded in causal relations, as opposed to mere statistical associations. Hence, integrating tools from causal representation learning into the RL framework seems to be a promising path forward.

Could you clarify the difference between OOT and OOD generalization with explicit examples?

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.

How does policy confounding relate to prior work on spurious correlations and generalization in RL?

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.

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?

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.

We appreciate the time the reviewer invested in reviewing our work. We hope our answers below can help clear up the reviewer’s concerns.

Overall, the paper is very awkward for reading. The authors propose many definitions and setup notations in the early pages before finally arriving to the definition of policy confounding (Sec. 6).

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.

I would suggest for the authors to reformat the paper to first give a high level sketch or intuition of the definition, perhaps aided by a visual figure, and then explaining the numerous notations and definitions required for defining policy confounding.

Both the introduction and the example in Section 2 are intended to provide a high-level intuition of policy confounding. We would be happy to correct the text and add more explanations where needed if the reviewer could point out what specific parts of the text are unclear.

In section 2, the authors provide a very lengthy (almost page long) textual description of policy confounding. This is very difficult to read, I would advise the authors to move explanations into visual aids, split up into paragraphs, highlight key steps, etc.

Section 2 is, in fact, an example of policy confounding in a gridworld environment. The example is accompanied by Figure 2, which depicts a sketch of the environment and a plot showing its practical effects. We kindly ask the reviewer to specify any additional visual aids they would like us to include. Regarding the text, we will split it into another paragraph beginning with "At first sight" to enhance readability. Any other suggestions regarding which specific parts of the text are unclear would be greatly appreciated.

The authors can improve in two ways. First, they can connect their theory more to their toy experiments. For example, is it possible to come up with a toy environment where all possible representations can be enumerated and tracked? Then it could be interesting to see what kinds of representations RL algorithms learn, and if they indeed are minimal and do not generalize.

We did analyze the internal representations learned by the RL algorithms in the Frozen T-Maze environment. The results of this experiment are reported in Appendix E. The experiment demonstrates that indeed the agent does learn a representation that does not generalize beyond the trajectories followed by the optimal policy.

Next, it would be interesting to investigate confounding in existing deep RL tasks, especially in well-known benchmarks (Atari, DMC, etc.).

We would like to emphasize that this is not an experimental paper. Our main contribution is the characterization of the phenomenon of policy confounding. The goal of our experiments is to verify that the phenomenon described by the theory does occur in practice. To that end we have designed as set of environments, which are intended as pedagogical examples and serve to to illustrate the problem in practice and to clarify the ideas introduced by the theory.

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.

and to see if some RL algorithms are better than others. For example, would MBRL algorithms fare better or worse than model-free RL algorithms?

We do experiment with two types of RL methods: on-policy (PPO) and off-policy (DQN). The reason for this experiment is that, according to the theory, off-policy methods should be less prone to policy confounding than on-policy methods as the former train on a more diverse set of trajectories. This is confirmed by the experiments (See Figure 5). There is no reason why model-based RL (MBRL) methods would not suffer from policy confounding as much as model-free methods.

Generally, we observe the tendency in RL research to run numerous experiments and then draw conclusions based on a subset of the results. This practice is inefficient in terms of time and energy. In this paper, our approach has been to first analyze the issue analytically and then conduct targeted experiments to validate that they indeed support the proposed theory.

We appreciate the input provided by the reviewer, which will help us enhance the quality of the paper. We will address these points in the next revision.

However, we would like to emphasize that a rejection recommendation based solely on presentation issues seems excessive, especially considering that the quality of the presentation has been highly praised by reviewers BRdu and kJx4.

Are there any remaining concerns that are impeding the reviewer's recommendation for acceptance?

We would like to thank the reviewer for taking the time to review our work. We believe there are certain misconceptions about the paper that we hope our responses below can clear up.

It is not clear how Example 1 corresponds to the ideas before and after.

The reviewer seems to be confusing state (observation) variables with states and state representations. In Example 1, the sequence of loctions $L_0,... L_t$, signals $X_0,..., X_t$, goals $G_0, …, G_t$, and actions $A_0, …, A_t$ are the state variables from which only $L_0,... L_t$, $X_0,..., X_t$, and $A_0, …, A_t$ are observed, we call the latter the set observation variables and denote it by $\Theta_t$. States are abstract entities that can be represented by the observation variables (see Definitions 1 and 2).

Now, let's clarify how the example relates to the definitions and theoretical results. There are two scenarios in Example 1, which correspond to the two DBNs in Figure 2.

In the first scenario (Figure 2, left), actions are exogenous variables and thus taken at random. Hence, the only valid state representation is $\langle X_0, L_t \rangle$ (i.e., the agent needs to condition ob its location and the signal received at the start location). This representation is Markov (Definition 4) and minimal (Definition 5).

In the second scenario (Figure 2, right), actions are sampled from the optimal policy. In this case, states can be represented simply as $\langle L_t \rangle$. This representation is $\pi$-Markov (Definition 7) and $\pi$-minimal (Definition 8). The reason for this is that under the optimal policy, the agent’s location $L_t$ becomes a proxy for the signal at the start location $X_0$ ($L_t$ and $X_0$ are perfectly correlated). Hence, the agent can ignore the signal and still predict transitions and rewards. In particular, $R^{\pi^*}(L_8 = \text{green goal}) = +1$ when following $\pi^*$, while in the first scenario since $L_8$ and $X_0$ are not correlated, $R(L_8 = \text{green goal}) = \pm1$.

As demonstrated by Theorem 1, and highlighted by the diagram in Figure 3, the underlying phenomenon that makes $L_t$ and $X_0$ become correlated is policy confounding.

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.

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.

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.

We sincerely appreciate the time the reviewer invested in reviewing our work and providing useful feedback. We are very much encouraged by the overall positive evaluation and especially by the reviewer’s acknowledgment of the significance of our work.

Please find the answers to your questions below:

The primary weakness is that the experiments focus on toy examples.

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.

mention why the avg return is around 0.2. Is this around the level that is to be expected?

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.

As an alternate demonstration, instead of adding the ice to the environment, why not simply place the agent directly on the state above (or below) the original starting location and see how it fares on the train env?

We agree with the reviewer that placing the agent directly outside its usual path would make everything much simpler. Yet, some may argue that the agent has never been trained on such trajectories and, therefore, it is reasonable that it is not able to perform well. To address this type of criticism, we made sure that the trajectories the agent is evaluated on are both possible and probable in the original environments.

The reviewer may find the experiment reported in Appendix E.1 interesting. In this experiment we manually permuted the value of the signal variable (X in Figure 2) to see whether the agent attends to it or not.

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.

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.

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 states rather than only the ones visited by pi?

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$).

About "Narrow trajectory distributions". If the environment transitons leads to a diverse set of states, would we still observe policy confounding?

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.

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?

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.

Is policy confounding inevitable to some extent? It seems to be a consequence of having policies focus on certain parts of the state space.

Good point! The way the RL objective is normally formulated makes avoiding policy confounding challenging. However, we don't think the issue is inevitable. We believe that, apart from maximizing reward, the agent should strive for learning representations grounded in causal relationships rather than mere statistical associations. Incorporating auxiliary objectives may help in this regard. However, this is again beyond the scope of this paper. Our goal in this paper was to highlight the phenomenon of policy confounding and raise awareness about its risks.

Dear reviewer,

As the discussion period comes to a close, we kindly request your feedback regarding our rebuttal. Have the responses to your questions effectively addressed the concerns you raised?