Seeing is not Believing: Robust Reinforcement Learning against Spurious Correlation
Authors: Wenhao Ding
Laixi Shi
Yuejie Chi
Ding Zhao
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
robustness
causality
spurious correlation
TL;DR: We propose Robust State-Confounded Markov Decision Processes to make RL agent robust against spurious correlation, where different portions of the state do not have causality but have correlations induced by unobserved confounders.

Abstract: Robustness has been extensively studied in reinforcement learning (RL) to handle various forms of uncertainty such as random perturbations, rare events, and malicious attacks. In this work, we consider one critical type of robustness against spurious correlation, where different portions of the state do not have correlations induced by unobserved confounders. These spurious correlations are ubiquitous in real-world tasks, for instance, a self-driving car usually observes heavy traffic in the daytime and light traffic at night due to unobservable human activity. A model that learns such useless or even harmful correlation could catastrophically fail when the confounder in the test case deviates from the training one. Although motivated, enabling robustness against spurious correlation poses significant challenges since the uncertainty set, shaped by the unobserved confounder and causal structure, is difficult to characterize and identify. Existing robust algorithms that assume simple and unstructured uncertainty sets are therefore inadequate to address this challenge. To solve this issue, we propose Robust State-Confounded Markov Decision Processes (RSC-MDPs) and theoretically demonstrate its superiority in avoiding learning spurious correlations compared with other robust RL counterparts. We also design an empirical algorithm to learn the robust optimal policy for RSC-MDPs, which outperforms all baselines in eight realistic self-driving and manipulation tasks.

Seeing is not Believing: Robust Reinforcement
Learning against Spurious Correlation
Wenhao Ding
Laixi Shi
Yuejie Chi
Ding Zhao'
CCaneeiie Mellon University (wenhaod, laixis, faaxxi,,,
CCliformia Institute of Technnooog yuejiec) Candrew . cmu edu, dingzhao@cmu, edu
Abstract
Robustness has been extensively studied in reinforcement learning (RL. handle various forms of uncertainty such as random perturbations, rare events, and ma- licious attacks. In this work, we consider one critical type of robustness against spurious correlation, where different portions of the state do not have correlations induced by unobserved confounders. These spurious correlations are ubiquitous in real-world tasks, for instance, self-driving car usually observes heavy traffic the daytime and light traffic at night due to unobservable human activity. model that learns such useless or even harmful correlation could catastrophically fail when the confounder in the test case deviates from the training one. Although motivated, enabling robustness against spurious correlation poses significant chal- lenges since the uncertainty set, shaped by the unobserved confounder and causal structure, is difficult characterize and identify. Existing robust algorithms that assume simple and unstructured uncertainty sets are therefore inadequate to address this challenge. To solve this issue, we propose Robust State-Confounler Markov Decision Processes (RSC-MDPs) and theoretically demonstrate its superiority avoiding learning spurious correlations compared with other robust RL counter- parts. We also design an empirical algorithm to learn the robust optimal policy for RSC-MDPs, which outperforms all baselines in eight realistic self-driving and manipulation tasks. Please refer to the website for more details.
1 Introduction
Reinforcement learning (RL), aiming to learn policy to maximize cumulative reward through interactions, has been successfully applied to a wide range of tasks such as language generation [I]. game playing (2, autonomous driving (3], etc. While standard RL has achieved remarkable success in simulated environments, a growing trend in RI. is to address another critical concern robustness with the hope that the learned policy still performs well when the deployed (test) environment deviates from the nominal one used for training [4). Robustness is highly desirable since the performance of the learned policy could significantly deteriorate due to the uncertainty and variations of the test environment induced by random perturbation, rare events, even malicious attacks (5.6].
Despite various types of uncertainty that have been investigated RL, this work focuses on the uncertainty of the environment with semantic meanings resulting from some unobserved underlying variables. Such environment uncertainty, denoted as structured uncertainty, innumerable real-world applications but still receives little attention in sequential decision-making tasks 77] To specify the phenomenon of structured uncertainty, us consider a concrete example (illustrated in Figure I in driving scenario, where shift between training and test environments caused by an unobserved confounder can potentially lead to a severe safety issue. Specifically, the observations brightness and traffic density do not have cause and effect on each other but are
motivated by
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

controlled by confounder (i.e. sun and human activity) that is usually unobserved to the agent. During training, the agent could memorize the spurious state correlation between brightness and traffic density, i.e., the traffic is heavy during the daytime but light at night. However, such correlation could be problematic during testing when the value of the confounder deviates from the training one, e.g., the traffic becomes heavy at night due to special events (human activity changes), as shown at the bottom of Figure Consequently, the policy dominated by the spurious correlation in training fails out-of-distribution samples (observations of heavy traffic at night) in the test scenarios.
The failure of the driving example in Figure T is attributed to the widespread and harmful spurious correlation, namely, the learned policy is not robust the structured uncer- tainty of the test environment caused by the unobserved confounder. However, ensuring robustness to structured un- certainty is challenging since the targeted uncertain region- the structured uncertainty set of the environment is carved by the unknown causal effect of the unobserved confounder, and thus hard to characterize. concerning robustness in RL 8 usually consider homo- geneous and structure structure-agnostic uncertainty set around the state 19 6 [0), action [T] [2), or the training environment 13 mcasured by some heuristic functions (9] [3] 8 to
Unobserved Confounder
Heavy Traffic
Sun
Human Activity
Light Traffic
Train
Bright- ness
Traffic Density
Test (fail)
Light Traftic
contrast, prior works
Spunous
Correlation
Heavy Traffic
account for unstructured random noise small perturba- tions. Consequently, these prior works could not cope with the structured uncertainty since their uncertainty set is dif- ferent from and cannot tightly cover the desired structured uncertainty set, which could be heterogeneous and allow for potentially large deviations between the training and test environments.
Figure 1: A model trained only with heavy traffic in the daytime learns the spurious correlation between brightness and traffic density and could fail to drive in light traffic in the daytime.
In this work, to address the structured uncertainty, we first propose a general RL formulation called State-confounded Markov decision processes (SC-MDPs), which model the possible causal effect of the unobserved confounder in an RL task from causal perspective. SC-MDPs better explain the reason for semantic shifts in the state space than traditional MDPs. Then, we formulate the problem of seeking robustness to structured uncertainty as solving Robust SC-MDPs (RSC-MDPs), which optimizes the worst performance when the distribution of the unobserved confounder lies in some uncertainty set. The key contributions of this work summarized as follows.
We propose new type of robustness with respect to structured uncertainty to address spurious correlation in RL and provide a formal mathematical formulation called RSC-MDPs, which are well-motivated by ubiquitous real-world applications.
We thcoretically justify the advantage of the proposed RSC-MDP framework against structured uncertainty the prior formulation in robust RL without semantic information.
We implement an empirical algorithm to find the optimal policy of RSC-MDPs and show that outperforms the baselines on eight real-world tasks in manipulation and self-driving.
2 Preliminary and Limitations of Robust RL
In this section, we first introduce the preliminary formulation of standard RL and then discuss a natural type of robustness that is widely considered in the RL literature and most related to this work robust RL.
Standard Markov decision processes (MDPs). An episodic finite-horizon standard MDP is represented by M (S. A. T.r,P), where C Rm and A C Rd are the state and action spaces, respectively, with n/da being the dimension of state/action. Here, T is the length of the horizon; P = {PiJTEeST, where P. S x A + A(S) denotes the probability transition kernel at time step for all < t I T: and r {ri]isET denotes the reward function, where rz S x A + 00, 1] represents the deterministic immediate reward function. A pol- icy (action selection rule) denoted by π {m {misIISTT namely, the policy at time step is Te S + A(A) based on the current state 8t Tel: [t). To represent the long-term cumulative
sometimes they observed but ignored given many variables be considered in neural networks.

---

reward, the value function V.T.P S + R and Q-value function Q : 8 8 x A + R associ- ated with policy at step are defined V. (s) Er,r | (s Q7 (s, = P Ei= (sk: ax) St s, a, = a), where the expectation is taken over the sample trajectory {(s,, 4:) hisgr generated following ae- ~ Te(- 8t) and 8t41 ~ Pt ( 8t, 8t,()).
az) 8k =s and
Robust Markov decision processes (RMDPs). As a robust variant of standard MDPs motivated by distributionally robust optimization, RMDP is a natural formulation to promote robustness to the uncertainty of the transition probability kernel [3 I5 represented as Mrah = {S. A, r, {P Here, we reuse the definitions of S, A, T", in standard MDPs, and denote 2t (P") as an uncertainty set of probability transition kernels centered around a nominal transition kernel pD = (PT)T measured by some 'distance' function P with radius o. In particular, the uncertainty set obeying the (s, a) recttttnggll [16] can be defined over all (s, a) state-action pairs at each time step as
(P%) = @ 2 (Pt.a)
U (P == (Pr.s,a E =A(S): (Pr.s,a PrssapPee Plu Plina) S so),
(1)
where @ denotes the Cartesian product. Here, Pt.ca = Pc(- (() E A(S) and Ph (- (-18,a)) E A(S) denote the transition kernel P. or P each state-action pair a) respec- tively. Consequently, the next state 8t41 follows 81+1 ~ Pa(- 8t, (t) for any P. E &t (PP namely, St+1 can be generated from any transition kernel belonging to the uncertainty set 41 4 (PE., (.s...))
rather than a fixed one in standard MDPs. As a result, for any policy л, the corresponding robust value function and robust Q funcciin are defined as
V V.Ta (s)
inf (PO)
V. P (s);
(s,
Oi P PEll" (po)
inf
(2)
which characterize the cumulative reward in the worst case when the transition kernel is within the uncertainty set set (P'). Using samples generated from the nominal transition kernel po, the goal of RMDPs is to find an optimal robust policy that maximizes Vi when = 1, i.e., perform optimally in the worst case when the transition kernel of the test environment lies in a prescribed uncertainty set lm (P').
Lack of semantic information in RMDPs. In spite of the rich literature on robustness in RL, prior works usually hedge against the uncertainty induced by unstructured random noise or small perturbations, specified as a small and homogeneous uncertainty set around the nominal one. For instance, in RMDPs, people usually prescribe the uncertainty set of the transition kernel using heuristic and simple function & with relatively small a. However, the unknown uncertainty in the real world could have complicated and semantic structure that cannot be well-covered by a homogeneous ball regardless of the choice of the uncertainty radius a,, leading to either over conservative policy (when is large) insufficient robustness (when o is small), Altogether, we obtain the natural motivation of this work: How to formulate such structured uncertainty and ensure robustness against it?
3 Robust RL against Structured Uncertainty from Causal Perspective
To describe structured uncertainty, we choose to study MDPs from a causal perspective with a basic concept called the structural causal model (SCM). Armed with the concept, we formulate State-confounded MDPs - a broader set of MDPs in the face of the unobserved confounder in the state space. Next, we provide the main formulation considered in this work robust state-confounded MDPs, which promote robustness to structured uncertainty.
Structural causal model. We denote a structural causal model (SCM) [17 by a tuple (X, F, P= } where X the set of exogenous (unobserved) variables, is the set of endogenous (observed) variables, and pa is the distribution of all the exogenous variables. Here, F is the set of structural functions capturing the causal relations between X and Y such that for each variable y/4 E Y. fs E F is defined as yi + fi (P(ss). 2;), where S X and PA(μs) C Y denotes the parents of the node yı- We say that a pair of variables yi and yj are confounded by a variable C' (confounder) they are both caused by C, i.e., C' E PA(j:) and C' E PA(y)). When two variables yı and yj do not have direct causality, they are still correlated if they are confounded, which case this correlation called spurious correlation.
3

---

(a) SC SC-MPP (Ours)
(b) Standard MDP
(G) POMDP/SA-MDP
(d) MDPUC
(e) Confounded MDP
confounding
Ci
Ot
spunicus
backdoor path
81'
de
Gt
Tt
The probabilistic graphs of our formulation (SC-MDP) and other related formulations (specified in Appendix BB. due to the limited space). it means the first dimension of 8t. 8y is shorthand for 81+1- In SC-MDP, the orange line represents the backdoor path from state sp. to action ay opened by the confounder C₁, which makes the learned policy π rely on the value of cc.
Figure 2:
3.1 State-confounded MDPs (SC-MDPs)
We now present confoudonf MDPs (SC-MDPs), whose probabilistic graph is illustrated in Figure 2 with comparison to standard MDPs in Figure Besides the components standard T, rr.. introduce set of unobserved confounder C. {ct) ISST,, where E C denotes the confounder that is generated from some unknown but fixed distribution P. at time step i.e.,. C P. ACC).
MDPs .M A,
To characterize the causal effect of the confounder C₁ on the state dynamic, we resort to an SCM, where CN is the set of exogenous (unobserved) confounder and endogenous variables include all dimensions of states {s [s/hisis..serr: and actions {a, Jigegr Specifically, the structural function F is considered as {Pi 1ssss,,ia confounder C to each dimension of the next state 84 for all time steps, i.e., 8141 ~ P. St, dg, c.). Notably, the specified SCM does not confound the reward, i.e., [1( az) does not depend on the confounder ct.
the transition from the current state 8t, action de and the
Armed with the above SCM, denoting p = {PE), we can introduce state-confounded MDPs (SC-MDPs) represented by Msc = (S, A,T,r, (P:), pd} (Figure a)). a)). A policy is denoted as = {T where each Ny results in intervention (possibly stochastic) that sets a₂ ~ me(: (s) at time step regardless of the value of the confounder.
State-confounded value function and optimal policy. Given St, the causal effect of a, the next state 51+1 plays an important role in characterizing value function/Q-ffunction. To ensure the identifiability of the causal effect, the confounder are assumed obey the backdoor criterion [17 18), leading to the following confounde value function (SC-value function) and state- confounded O-function (SC-Q function) 19]
rhe (K( (x) 8t 8; Ph, 8441 Pa( 8k, aky (ls
V.H.P (8) Eer,
Pr (s, Er
Σ
The (5k. ax) s, az = a; Ck Ph Sk41 ~ P/ / axs cu)
(3)
Remark Note that the proposed SC-MDPs serve as a general formulation for broad family of RL problems that include standard MDPs as special case. Specifically, any standard MDP M AS P.T.ri r} can be equivalently represented by at least one SC-MDP My = A,,, C, (P)), as long as E. [P] St. Cr, ct) = [P se, at)], for all 1 < 1,,n,,s5
3.2 Robust state-confounded MDPs (RSC-MDPS)
In this work, consider robust ftae-coonooudeer MDPs (RSC-MDPs) variant of SC-MDPs promoting the robustness to the uncertainty of the unobserved confounder distribution P denoted by MMscrob (S,A,T,r, C, (Pi),
Here, the perturbed distribution of the unobserved confounder is assumed in uncertainty set t (P*) centered around the nominal distribution p

---

with radius measured by some "distance' function p: [C) x A(C) - R i.e. (P"))) @ u²((())
= (=
A(C) p (P. = o)
(4)
Consequently, the corresponding robust SC-value function and robust SC-Q function are defined
Vini (s) =
inf (PO)
Vs), Q, ==
inf
(5)
(PM("(P))
representing the worst-case cumulative rewards when the confounder distribution lies in the uncer- lainty set 21 (P").
Then natural question is: does there exist optimal policy that maximizes the robust SC-value function V. for any RSC-MDP so that we can target learn? To answer this, introduce the following theorem that ensures the existence of the optimal policy for all RSC-MDPs. The proof can be found in Appendix [.I]
Theorem 1 (Existence of an optimal policy). Ler II be the set of all non-stationary and stochastic policies. Consider any RSC-MDP, there exists at least one optimal policy xx, TT JssTT such that for all (8, a) E S x A and 1 S 1. S T, one hus
Vi (s) sup V. and Qi "is"" (s, FETI
(s. Q sup Qi ( (s)
=en
In addition, RSC-MDPs also possess benign properties similar to RMDPs such that for any policy and the robust optimal policy = , the corresponding robust SC Bellman consistency equation and robust SC Bellman optimality equation are also satisfied (specified in Appendix C.3.3).
Goal. Based on all the definitions and analysis above, this work aims to find an optimal policy for RSC-MDPs that maximizes the robust SC-value function in (5), yielding optimal performance in the worst case when the unobserved confounder distribution falls into an uncertainty set U (P*).
3.3 Advantages of RSC-MDPs over traditional RMDPs
The most relevant robust RL formulation to ours is RMDPs, which has been introduced in Section 2 Here, we provide a thorough comparison between RMDPs and our RSC-MDPs with theoretical justifications, and leave the comparisons and connections to other related formulations in Figure
and Appendix B. due space limits.
To begin, at each time step f, RMDPs explicitly introduce uncertainty to the transition probabil- ity kernels, while our RSC-MDPs add uncer- tainty to the transition kernels in latent (and hence more structured) manner via perturbing the unobserved confounder that partly deter- mines the transition kernels. As an example, imagining the true uncertainty set encountered in the real world is illustrated as the blue region in Figure 3, which could have a complicated structure. Since the uncertainty set in RMDPs is homogeneous (illustrated by the green circles), one often faces the dilemma of being either too conservative (when a is large) or too reckless (when is small). In contrast, the proposed RSC-MDPs - shown in Figure 3b)- - take advantage of the structured uncertainty set (illustrated by the orange region) enabled by the underlying SCM, which can potentially lead to much better estimation of the true uncertainty set. Specifically, the varying unobserved confounder induces diverse perturbation to different portions of the state through the structural causal function, enabling heterogeneous and structural uncertainty sets over the state space.
(a) RMDPs
(b) RSC-MDPs
Space
large
Valld
small IT
Nomhal pe
Nominal Pa
Space AiS)
Space A(S)
Figure 3: (a) RMDPs add homogeneous noise to states, while (b) RSC-MDPs perturb the con- founder influence states, resulting in subset of the valid space.
Theoretical guarantees of RSC-MDPs: advantages of structured uncertainty. To theoretically understand the advantages the proposed robust formulation of RSC-MDPs with comparison to prior works, especially RMDPs, the following theorem verifies that RSC-MDPs enable additional robustness against semantic attack besides small model perturbation or noise considered in RMDPs.
The proof is postponed to Appendix C.2
5

---

Theorem 2. Consider any T 2 2. Consider some standard MDPs M
equiv- alently represented as an SC-MDP M& = Ar, C, (Pi), PS) with C == 10.1). and total variation as the 'distance" function p to measure the uncertainty set (the admissible uncertainty level obeys o E [0, 1]). For the corresponding RMDP M with the uncertainty set n (P%), and the proposed RSC-MDP Msc.roh = (S,A,T,r, C, (Pi), (P=)). the optimal robust policy RRMD associated with Mrab and TRSC associated with Mac obey: given d2 E (2 1]. there exist RSC-MDPs with some initial state distribution & such that
T
.0 (0)
V2" MM
(0)
Yar [0.
(6)
In words, Theorem 2 reveals a fact about the proposed RSC-MDPs: RSC-MDPs could succeed in intense semantic atlacks while RMDPs fail. As shown by (6), when fierce semantic shifts appear between the training and test scenarios - perturbing the unobserved confounder in a large uncertainty set (P*)) solving RSC-MDPs with TRSČ succeeds in testing while RRMDP trained by solving RMDPs can fail catastrophically. The proof is achieved by constructing hard instances of RSC-MDPs that RMDPs could not cope with due to inherent limitations. Moreover, this advantage of RSC-MDPs is consistent with and verified by the empirical performance evaluation in Section 5.3 R1.
*001
An Empirical Algorithm to Solve RSC-MDPs: RSC-SAC
When addressing distributionally robust problems in RMDPs, the worst case is typically defined within prescribed uncertainty set in clear and implementation -feiinnd manner, allowing for it- erative or analytical solutions. However, solving RSC-MDPs could be challenging as the structured uncertainty set is induced by the causal effect of perturbing the confounder. The precise characteri- zation of this structured uncertainty set is difficult since neither the unobserved confounder nor the true causal graph of the observable variables is accessible, both of which are necessary for inter- vention counterfactual reasoning. Therefore, we choose to approximate the causal effect of per- turbing the confounder by learning from the data collected during training.
Algorithm 1: RSC-SAC Training
Input:
policy æ, data buffer D, transition
model Po, ratio of modified data /
for E [1,7] do
Sample action 0₂ (-)(s.)
(8++1. - Env( 8t, (t)
Add buffer D D U {st, At, 81+1,
for sample batch B E D do
Randomly select 33% data in B Modify 8; in selected data with (5)=1, ~ Po des 9o)
Replace data with (St, Gr 8 f2)
L= ||$ - 8 1/2 + ||, fell
Update @ and o with L + A|| p
Update π with SAC algorithm
In this section, we propose intuitive yer effective empirical approach named RSC-SAC for solving RSC-MDPs, which outlined in Algorithm [ We first estimate the effect of perturbing the distribution J of the confounder generate new states (Section 4.1). Then, we Iearn the structural causal model Pi to predict rewards and the next states given the perturbed states (Section (42)). By combining these two components, transitions (St, at Ft, 8t+1) from the structured uncertainty set. To learn the optimal policy, we construct the data buffer with a mixture of the original data and the generated data and then use the Soft Actor-Critic (SAC) algorithm 20 to optimize the policy.
construct data generator capable of simulating novel
4.1 Distribution of confounder
As we have no prior knowledge about the confounders, we choose to approximate the effect of per- turbing them without explicitly estimating the distribution Pe. We first randomly select a dimension i from the state s, to apply perturbation and then assign the dimension i of s; with a heuristic rule. We select the value from another sample sy that has the most different value from s; in dimension and the most similar value to 8; in the remaining dimensions. Formally, this process solves the following optimization problem to select sample ie from a batch of K samples:
||s 112
s arg max
k (1,..,KK K}
E-113 - "
where of and 8,'h means dimension i of 8; and other dimensions of 8, except for i, respectively. Intuitively, permuting the dimension of two samples breaks the spurious correlation and remains the

---

Door (composition)
Brighiness (distraction)
Day
High
Ego
Low
Close
Far
Figure 4: Two tasks used in experiments. Door is a composition task implemented in Robosuite with spurious correlation between the positions of the handle and the door. Brightness is a distraction task implemented in Carla with a spurious correlation between the number vehicles and day/night.
most semantic meaning of the state space, However, this permutation sometimes also breaks the true cause and effect between dimensions, leading to a performance drop. The trade-off between robustness and performance [2] is long-standing dilemma in the robust optimization framework, which we will leave to future work.
4.2 Learning of structural causal model
With the perturbed state 8t, then learn SCM to predict the next state and reward considering the effect of the action on the previous state. This model contains a causal graph to achieve better generalization unseen state-action pairs. Specifically, we simultaneously learn the model parameter and discover the underlying causal graph in a fully differentiable way with ($L+1. 14+1.f7 ft) ~ Po(se, de, Go), where 0 is the parameter of the neural network of the dynamic model and Φ E IR (((++d))x(++1) is the parameter to represent causal graph G between {st, de} and {1) This graph is represented by binary adjacency matrix G, where 1/0 means the existence/absence of an edge. Pa has an encoder-decoder structure with matrix G an intermediate linear transformation. The encoder takes state and action in and outputs features fe E (a++))xd for each dimension, where dy is the dimension of the feature. Then, the causal graph multiplied to generate the feature for the The decoder takes in Sa and outputs the next state and reward. The detailed architecture of this causal transition model can be found in Appendix D.I
decoder La = e G R22x(xx
The objective for training this model consists of two parts, one is the supervision signal from collected data ||st+1 - 1 + I[ - f and the other is penalty term > A||G||, IIp with weight to encourage the sparsity of the matrix G. The penalty is important to break the spurious correlation between dimensions of state since forces the model to eliminate unnecessary inputs for prediction.
5 Experiments and Evaluation
In this section, we first provide benchmark consisting of eight environments with spurious correla- tions, which may be of independent interest to robust RL. Then we evaluate the proposed algorithm RSC-SAC with comparisons to prior robust algorithms in RL.
5.1 Tasks with spurious correlation
To the best of our knowledge, no existing benchmark addresses the issues of spurious correlation in the state space of RL. To bridge the gap, we design benchmark consisting of eight novel tasks in self-driving and manipulation domains using the Carla [22] and Robosuite (23] platforms (shown in F) Tasks are designed to include spurious correlations terms of human common sense, which is ubiquitous in decision-making applications and could cause safety issues. We categorize the tasks into distraction correlation and composition correlation according to the type of spurious correlation. We specify these two types of correlation below and leave the full descriptions of the tasks in Appendix D.22
Distraction correlation between task-relevant and task-irrelevant portions of the state. The task-irreles irrelevant part could distract the policy model from learning important features and lead to a performance drop. A typical method to avoid distraction background augmentation 24 25 We design four tasks with this category of correlation, i.e., Lift, Wipe, Brightness, and CarType.
Composition correlation is between two task-relevant portions of the state. This correlation usually exists compositional generalization, where states are re-composed to form novel tasks
during testing. Typical examples multi-task RL [26] 27 and hierarchical RL (28 29 We design four tasks with this category of correlation, i.c., Stack, Door, Behavior, and Crossing.
7

---

Stack
Crossing
CarType
1.0
Nominal Environmntts
0.8
0.8
1.5
9.00
0.2
04
0.8
20
0.1
95
1.
25
55
15
1.
233
25
55
A
Shifted Ervironments
0.4
0.1
0.4
0.4
11
Toinng Stre
Traning
Tarting 5tp
Thaning 5hpp
Ours
MoCoDA
DBC
SAC
ATLA
2HDP-6
RMDP-
Figure 5:
The first row shows the testing reward on the nominal environments, while the second row shows the testing reward the shifted environments.
Table 1:
Testing reward on shifted environments. Bold font means the best reward.
Method
Brightness
Behavior
Crossing
CarType
Lift
Stack
Wipe
Door
SAC
0.56=0.13
0.134+0.03
0.81
0.63+0.14
0.58++0.33
0.26:10.12
0.16=0.20
0.08
RMDP-G
0.55
0.16110.0
0.47+0.13
0.53
0.31 4.008
0.
0.06
0.07-10.03
RMDP-U
0.54
0.13
0.6010.15
0.39
0.51 +0.17
0.23
0.06=0.17
0. 0010
MoCoDA
0.50
0.16+0.05
0.22+0.14
0.23+0.12
0.46+10.14
0.29:10.11
0.09 0.01
ATLA
0.48=0.11
0.14
0.61
0.52+0.14
0.61:t0.18
0.21:10.12
0.29=0.18
0.28=1019
DBC
0.52=0.18
0.16::0.03
0.68:10.12
0.45:10.10
0.12:t0.02
0.03:10002
0.19=0.35
0.01
Active
0.47=0.14 0.14
0.83
0.77:t0.14
0.35
0.24::0.12
0.17=0.17
0.05
RSC-SAC
0.99-0.11
1.02:+.09
1.04
1.03=0002
0.98++0.04
0.77=10.20
0.85=0.12
0.61
5.2 Baselines
Robustness RL has been explored in terms of diverse uncertainty set over state, action, or transition kernels. Regarding this, we use non-robust RL and four representative algorithms of robust RL as baselines, of which are implemented on top of the SAC (0) algorithm. Non-robust RI (SAC): This serves as a basic baseline without considering any robustness during training: Solving robust MDP: We generate the samples to cover the uncertainty set over the state space by adding perturbation around the nominal states that follows two distribution, i.e.. uniform distribution (RMDP-U) and Gaussian distribution (RMDP-G), Solving SA-MDP: We compare ATLA (6) strong algorithm that generates adversarial states using an optimal adversary within the uncertainty set. Invariant feature learning: We choose DBC 30 that learns invariant features using the bi-simulation metric [I] and 32 (Active) that actively sample uncertain transitions to reduce causal confusion. Counterfactual data augmentation: We select MoCoDA [33], which identifies local causality to switch components and generate counterfactual samples to cover the targeted uncertainty set. We adapt this algorithm using an approximate causal graph rather than the true causal graph.
5.3 Results and Analysis
To comprehensively evaluate the performance of the proposed method RSC-SAC, we conduct experiments to answer the following questions: Q1. Can RSC eliminate the harmful effect of spurious correlation in learned policy? Q2. Does the robustness RSC-SAC only come from the sparsity of model? Q3. How does RSC-SAC perform in the nominal environments compared to non-robust algorithms? Q4. Which module critical in our empirical algorithm? Q5. RSC-SAC robust to other types of uncertainty and model perturbation? Q6. How does RSC-SAC balance the tradeoff between performance and robustness? We analyze the results and answer these questions in the following.
8

---

Table 2: Testing reward on nominal environments. Underline means the reward is over 0.9.
Method
Brightness
Behavior
Crossing
CarType
Lift
Stack
Wipe
Door
SAC
0.00
1.00
0.00
1.00
1.00
1.00:10.09
1.00 1.00
RMDP-G
1.04
1.00
0.78
0.79
0.92=
0.86:10.14 14
0.994-0.13
0.9941.06
RMDP-U
1.02:10.09
1.04
0.90+
0.88
0.97=0005
0.92:10.12
0.97:10.14
0.88
MoCoDA
0.65=0.17
0.78
0.57
0,55
0.79::0.11
0.72:1:0.88
0.69=0.13
0.41
ATLA
0.99:10.11
0.98
0.89
0.88
0.94=0.08
0.88:10.10
0.96:10.12
0.97:10.05
DBC
0.75=0012
0.78:10.10
0.85
0.80
0.27:100004
0, 0..0080.08
0.31
0.01
Active
1.02++.10
1.08+0006
1.00=0.02
1.00=0.02
0.99=0.03
0.90= 12
0,93++0.20
0.99+0.05
RSC-SAC
0.9210.31
1.06:10.07
0.961..03
0.96110.03
0.96
1.04.10.08
0.92:10.14
0.98
R1. RSC-SAC is robust against spurious correlation. The testing results of our proposed method with comparisons to the baselines are presented in Table [ where the rewards are normalized by the episode reward of SAC in the nominal environment. The results reveal that RSC-SAC significantly outperforms other baselines in shifted test environments, exhibiting comparable performance to that of vanilla SAC on the nominal environment 5 out of tasks. An interesting and even surprising finding, as shown in Table M is that although RMDDP-G RMDP-U, and ATLA are trained desired to be robust against small perturbations, their performance tends to drop more than non-robust SAC in some tasks. This indicates that using the samples generated from the traditional robust algorithms could harm the policy performance when the test environment outside of the prescribed uncertainty set considered in the robust algorithms.
R2. Sparsity of the model is only rea- son for the robustness of RSC-SAC. As exist-
1.4
Ours
SAC-a-1
SAC-c:100
SAC
SAC-a-10
SAC-a-1000
ing literature shows (34), sparsity regularization benefits the elimination of spurious correlation and causal confusion. Therefore, we compare our method with a sparse version of SAC (SAC- Sparse): we add an additional penalty a|W]; 1 during the optimization, where W is the param- eter of the first linear layer of the policy and value networks and a is the weight. The results of both Distraction and Composition are shown in Figure 6 We have two important findings based on the results: (1) The sparsity improves the robustness of SAC in the setting of distrac- tion spurious correlation, which is consistent
1.0
0.8
0.4
0.0
Distraction (nominal)
Distraction (shiited)
Composition (nominal)
Composition (shifted)
Figure 6: Comparison between SparSe and our method. a is the weight of regularization.
with the findings in 344. (2) The sparsity does not help with the composition type of spurious correlation, which indicates that purely using sparsity regularization cannot explain the improvement of our RSC-SAC. In fact, the semantic perturbation in our method plays an important role in augmenting the composition generalization.
R3. RSC-SAC maintains great performance in the nominal environments. Previous literature [1] finds out that there usually exists a trade-off between the performance the nominal environment and the robustness against uncertainty. To evaluate the performance of RSC-SAC in the nominal environment, we conduct experiments and summarize results in Table which shows that RSC-SAC still performs well in the training environment. Additionally, the training curves are displayed in Figure 5 showing that RSC-SAC achieves similar rewards compared to non-robust SAC although converges slower than it.
R4. Both the distribution of confounder and the structural causal model are critical. To as- sess the impact of each module in our algorithm, we conduct three additional ablation studies (in Ta- w/o Po
Table 3: Influence of modules
Method
Lift
Behavior
Crossing
w/o Go
0.79
0.51
0.87-0.10
0.75-0.13
0.41+0.28
0.89+0.08
3), where we remove the causal graph Go. the transition model Pv. and the distribution of the con- founder pe respectively. The results demonstrate that the learnable causal graph G is critical for the performance that enhances the prediction of
w/o p
0.90
0.6640221
0.96+1004
Full model
0.98
1.02:
1.04:oozz

---

the next state and reward, thereby facilitating the generation of quslltualtty next states with current perturbed states. The transition model without Go may still retain numerous spurious correlations. resulting in performance drop similar to the one without Po. which does not alter the next state and reward. In the third row of Table 3 a crucial role in preserving semantic meaning and avoiding policy training distractions.
the performance drop indicates that the confounder Pe also plays
R5. RSC-SAC is also robust random pertur- bation. The final investigation aims to assess the generalizability of our method cope with random perturbation that is widely considered in robust RL (RMDPs). Towards this, we evaluate the proposed algorithm in the test environments added with ran- dom noise under the Gaussian distribution with two varying scales in the Lift environment. In Table Lift-0 indicates the nominal training environment, while Lift-0.01 and Lift-0. represent the environ- ments perturbed by the Gaussian noise with standard derivation 0.01 and 0.1, respectively. The results indicate that our RSC-SAC achieves comparable robustness compared to RMDP-0.01 in both large and small perturbation settings and outperforms RMDP methods in the nominal training environment.
Table 4: Random purterbuation
Method
Lift-0
Lift-0.01
Lift-0.1
SAC
1.00
0.77
0.46.10:23
RMDP-0.01
0.97
0.9610.06
1.51
RMDP-0.1
0.85
0.82=0.09
0.39+0.15
RSC-SAC
0.96##.05
0.94
0,44
R6. RSC-SAC keeps good performance and robustness for a wide range of B. As shown in Figure / the proposed RSC-SAC performs well in both nominal and shifted settings - keeping good performance in the nominal setting and achieving robustness, for a large range of (20%-70%). When the ratio of perturbed data is very small (1%), RSC-SAC almost achieves the same results as vanilla SAC in nominal settings and there is no robustness in shifted set- tings. As it increases (considering more robustness), the performance of RSC-SAC in the nominal setting gradually gets worse, while reversely gets better in the shifted settings (more robust). However, when the ratio is too large (>80%), the performances of RSC-SAC in both settings degrade lot, since the policy is too conservative so that fails in all environments.
1.1
1.0
0.9
Range of acceptable
0.8
Performance & Robustness
0.7
Carlype (nominal)
Cressing (nominal)
0.6
Car (shifted)
Crossing (shilted)
0.5
1%
10%
20%
30%
40%
50%
60%
70%
80%
90%
100%
Figure 7: cerformance-robustee tradeoff with differ- ent augmentation ratio B.
6 Conclusion and Limitation
This work focuses on robust reinforcement learning against spurious correlation in state space, which broadly exists in (sequential) decision-making tasks. We propose robust SC-MDPs a general framework to break spurious correlations by perturbing the value of unobserved confounders. We not only theoretically show the advantages of the framework compared to existing robust works in RL, but also design an empirical algorithm to solve robust SC-MDPs by approximating the causal effect of the confounder perturbation. The experimental results demonstrate that our algorithm is robust to spurious correlation - outperforms the baselines when the value of the confounder in the test environment derivates from the training one, It is important to note that the empirical algorithm we propose is evaluated only for low-dimensional states. However, the entire framework can be extended in the future to accommodate high-dimensional states by leveraging powerful generative models with disentanglement capabilities 355 and state abstraction techniques [36]
Acknowledgments and Disclosure of Funding
The work of W. Ding is supported by the Qualcomm Innovation Fellowship. The work of L. Shi and Y. Chi is supported in part by the grants ONR N00014-19-1-2400 NSF CCF-2106778, DMS-21 2134800 and CNS-2148212. L. Shi is also gratefully supported by the Leo Finzi Memorial Fellowship. Wei Shen and Xuehong Zhang Presidential Fellowship, and Liang Ji-Dian Graduate Fellowship at Carnegie Mellon University.
10

---

References
[1] OpenAI. Gpt-4 technical report, 2023.
[2] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Pannecrshelvam, Marc Lanc- tot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
[3] Mariusz Bojarski, Davide Del Testa, Daniel Dworakowski, Bernhard Firner, Beat Flepp, Prasoon Goyal, Lawrence D Jackel, Mathew Monfort, Urs Muller, Jiakai Zhang, et al. End to end learning for self-driving cars. arXiv preprint
2016.
[4] Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing goal-conditioned reinforce- ment learning with variational causal reasoning. arXiv preprint arXiv: 2022.
[5] A Rupam Mahmood, Dmytro Korenkevych, Gautham Vasan, William Ma, and James Bergstra. Benchmarking reinforcement learning algorithms on real-world robots. In Conference on robot learning, pages 561-591. PMLR, 2018.
[6] Huan Zhang, Hongge Chen, Duane Boning, and Cho-Jui Hsich. Robust reinforcement learning on state observations with learned optimal adversary. arXiv preprint arXiv:2101. 004522 2021.
[7] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
[8] Janosch Moos, Kay Hansel, Hany Abdulsamad, Svenja Stark, Debora Clever, and Jan Peters. Robust reinforcement learning: A review of foundations and recent advances. Machine Learning and Knowledge Extraction, 4(1):276-315, 2022.
[9] Huan Zhang. Hongge Chen, Chaowei Xiao, Bo Li, Mingyan Liu, Duane Boning, and Cho- Jui Hsieh. Robust deep reinforcement learning against adversarial perturbations on state observations. Advances in Neural Information Processing Systems, 33:21024-21037 2020. 2020.
[10] Songyang Han, Sanbao Su, Sihong He, Shuo Han, Haizhao Yang, and Fei Miao. What is the solution for state adversarial multi-agent reinforcement learning? arXiv preprint arXiv:22/2.02705, 2022.
[11] Chen Tessler, Yonathan Efroni, and Shie Mannor. Action robust reinforcement learning and applications in continuous control. In International Conference on Machine Learning, pages 6215-6224. PMLR, 2019.
[12] Kai Liang Tan, Yasaman Esfandiari, Xian Yeow Lee, and Soumik Sarkar. Robustifying reinforcement learning agents via action space adversarial training. In 2020 American control conference (ACC), pages 3959-3964. IEEE, 2020.
[13] Garud N Iyengar. Robust dynamic programming. Mathematics of Operations Research, 30(2):257-205 280, 2005.
[14] Wenhao Yang, Liangyu Zhang, and Zhihua Zhang. Toward thcoretical understandings of robust markov decision processes: Sample complexity and asymptotics. The Annals of Statistics, 50(6):3223-3244 3244888222--88, 2022.
[15] Laixi Shi and Yuejie Chi. Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity. arXiv preprint urXiv:2208.05766 2022.
[16] Wolfram Wiesemann, Daniel Kuhn, and Berç Rustem. Robust markov decision processes. Mathematics of Operations Research, 38(1):153- 183, 2013.
[17] Judea Pearl. Causality. Cambridge university press, 2009.
[18] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithmss. The MIT Press, 2017.

---

[19] Lingxiao Wang, Zhuoran Yang, and Zhaoran Wang. Provably efficient causal reinforcement learning with confounded observational data. Advances in Neural Information Processing Systems, 34:21164-21 1111114--117,, 175, 2021.
[20] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[21] Mengdi Xu, Peide Huang, Yaru Niu, Visak Kumar, Jielin Qiu, Chao Fang, Kuan-Hui Lee, Xuewei Qi, Henry Lam, Bo Li, et al. Group distributionally robust reinforcement learning with hierarchical latent variables. In International Conference on Artificial Intelligence and Statistics, pages 2677-2703. PMLR, 2023.
[22] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. In Conference on robot learning, pages 1-16. PMLR, 2017.
[23] Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martín, Abhishek Joshi, Soroush Nasiriany, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009.12293 2020.
[24] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised rep- resentations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5657 PMLR, 2020.
[25] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107,09644 2021.
[26] Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan. Vima: General robot manipulation with multimodal prompts. arXiv preprint arXiv:2210.0309, 2022.
[27] Yuchen Lu, Yikang Shen, Siyuan Zhou, Aaron Courville, Joshua Tenenbaum, and Chuang Gan. Learning task decomposition with ordered memory policy network. arXiv preprint arXiv:2103. 10972, 2021.
[28] Minjong Yoo, Sangwoo Cho, and Honguk Woo. Skills regularized task decomposition for multi-task offline reinforcement learning. Advances in Neural Information Processing Systems, 35:37432-3748 2022.
[29] Hoang Le, Nan Jiang, Alekh Agarwal, Miroslav Dudik, Yisong Yue, and Hal Daumé III. Hierarchical imitation and reinforcement learning. In International conference on machine learning, pages 2917-2926. PMLR, 2018
[30] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.
[31] Kim G Larsen and Ame Skou. Bisimulation through probabilistic testing (preliminary report). In Proceedings of the 16th ACM SIGPLAN-SIGACT symposium on Principles of programming languages, pages 344-352. 1989.
[32] Gunshi Gupta, Tim GJ Rudner, Rowan Thomas McAllister, Adrien Gaidon, and Yarin Gal. Can active sampling reduce causal confusion in offline reinforcement learning? In 2nd Conference Causal Learning and Reasoning, 2023.
[33] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. arXiv preprint arXiv:2210. /2287 2022.
[34] Jongjin Park, Younggyo Seo, Chang Liu, Li Zhao, Tao Qin, Jinwoo Shin, and Tie- Yan Liu. Object-aware regularization for addressing causal confusion in imitation learning. Advances in Neural Information Processing Systems, 34:3029-3042, 2021.
12

---

[35] Hyunjik Kim and Andriy Mnih. Disentangling by factorising, In International Conference on Machine Learning, pages 2649-2658. PMLR, 2018.
[36] David Abel. A theory of state abstraction for reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 9876-9877, 2019.
[37] Shirley Wu, Mert Yuksekgonul, Linjum Zhang, and James Zou. Discover and cure: Concept- aware mitigation of spurious correlation. arXiv preprint arXiv:2305.00650, 2023.
[38] Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes.
arXiv preprint arXiv: /50...2252229 2015.
[39] Huan Xu and Shie Mannor. Distributionally robust markov decision processes. Advances in Neural Information Processing Systems, 23, 2010.
[40] Eric M Wolff, Ufuk Topcu, and Richard M Murray. Robust control of uncertain markov decision processes with temporal logic specifications. In 2012 IEEE 51 IEEE Conference on Decision and Control (CDC), pages 3372-3379. IEEE, 2012.
[41] David I. Kaufman and Andrew J Schaefer. Robust modified policy iteration. INFORMS Journal on Computing, 25(3):396-411) 2013.
[42] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Fast bellman updates for robust mdps.
In International Conference on Machine Learning, pages 1979-1988, PMLR, 2018.
[43] Elena Smirnova, Elvis Dohmatob, and Jérémie Mary. Distributionally robust reinforcement learning. arXiv preprint arXiv:1902.08708, 2019.
[44] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Partial policy iteration for 11-robust markov decision processes. Journal of Machine Learning Research, 22(275): 2021.
[45] Vineet Goyal and Julien Grand-Clement. Robust markov decision processes: Beyond rectan- gularity. Mathematics of Operations Research, 2022.
[46] Esther Derman and Shie Mannor. Distributional robustness and regularization in reinforcement learning. arXiv preprint arXiv:2003, 02894, 2020.
[47] Aviv Tamar, Shie Mannor, and Huan Xu. Scaling up robust mdps using function approximation. In International conference on machine learning, pages 181-189. PMLR, 2014.
[48] Kishan Panaganti Badrinath and Dileep Kalathil. Robust reinforcement learning using least squares policy iteration with provable performance guarantees. In International Conference on Machine Learning, pages 51--520. PMLR, 2021.
[49] Zaiyan Xu, Kishan Panaganti, and Dileep Kalathil. Improved sample complexity bounds for distributionally robust reinforcement learning. arXiv preprint arXiv:2303.02778, 2023.
[50] Jing Dong, Jingwei Li, Baoxiang Wang, and Jingzhao Zhang. Online policy optimization for robust mdp. arXiv preprint 22X.::2209.13822 2022.
[51] Kishan Panaganti and Dileep Kalathil. Sample complexity of robust reinforcement learning with a generative model. In International Conference on Artificial Intelligence and Statistics, pages 9582-9602. PMLR, 2022.
[52] Zhengqing Zhou, Qinxun Bai, Zhengyuan Zhou, Linhai Qiu, Jose Blanchet, and Peter Glynn. Emite-sample regret bound for distributionally robust offline tabular reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pages 3331-3339, PMLR, 2021.
[53] Shengbo Wang, Nian Si, Jose Blanchet, and Zhengyuan Zhou. A finite sample complex- ity bound for distributionally robust Q-learning- In International Conference on Artificial Intelligence and Statistics, pages 3370-3398, PMLR, 2023.
13

---

[54] Jose Blanchet, Miao Lu, Tong Zhang, and Han Zhong. Double pessimism is provably efficient for distributionally robust offline reinforeement learning: Generic algorithm and robust partial coverage. arXiv preprint 0rX65:2305.09659, 2023.
[55] Zijian Liu, Qinxun Bai, Jose Blanchet, Perry Dong, Wei Xu, Zhengqing Zhou, and Zhengyuan Zhou. Distributionally robust g-learning. In International Conference on Machine Learning, pages I 13623-13643, PMLR, 2022.
[56] Shengbo Wang. Nian Si, Jose Blanchet, and Zhengyuan Zhou. Sample complexity of variance- reduced distributionally robust Q-learning. arXiv preprint arXiv:2305.18420, 2023.
[57] Zhipeng Liang. Xiaoteng Ma, Jose Blanchet, Jiheng Zhang, and Zhengyuan Zhou. Single- trajectory distributionally robust reinforcement learning. arXiv preprint arXiv:2301117211 2023.
[58] Yue Wang and Shaofeng Zou. Online robust reinforcement learning with model uncertainty. Advances in Neural Information Processing Systems, 34, 2021.
[59] Laixi Shi, Gen Li, Yuting Wei, Yuxin Chen, Matthieu Geist, and Yuejie Chi. The curious price of distributional robustness in reinforcement learning with a generative model. arXiv preprint arXiv:2305. 16589, 2023.
[60] Shyam Sundhar Ramesh, Pier Giuseppe Sessa, Yifan Hu, Andreas Krause, and Ilija Bogunovic. Distributionally robust model-based reinforcement learning with large state spaces. arXiv preprint 230x:22309.022366 2023.
[61] Kishan Panaganti, Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh. Robust reinforcement learning using offline data. Advances in neural information processing systems, 35:32211-32222 2022.
[62] Xiaoteng Ma, Zhipeng Liang, Jose Blanchet, Mingwen Liu, Li Xia, Jiheng Zhang, Qianchuan Zhao, and Zhengyuan Zhou. Distributionally robust offline reinforcement learning with linear function approximation. arXiv preprint arXiv:2209.0666 2022.
[63] Shengbo Wang, Nian Si, Jose Blanchet, and Zhengyuan Zhou, On the foundation of distribu- tionally robust reinforcement learning. arXiv preprint arXiv:2311.09018. 2023.
[64] Yuxin Pan, Yize Chen, and Fangzhen Lin. Adjustable robust reinforcement learning for online 3d bin packing. arXiv preprint arXiv:2310.04323, 2023.
[65] You Qiaoben, Xinning Zhou, Chengyang Ying, and Jun Zhu. Strategically-timed state- observation attacks on deep reinforcement learning agents. In ICML 2021 Workshop on Adversarial Machine Learning, 2021.
[66] Ke Sun, Yi Liu, Yingnan Zhao, Hengshuai Yao, Shangling Jui, and Linglong Kong. Exploring the training robustness of distributional reinforcement learning against noisy state observations. arXiv preprint arXiv:2 100.:0276077777 2021.
[67] Zikang Xiong, Joe Eappen, He Zhu, and Suresh Jagannathan. Defending observation attacks in deep reinforcement learning via detection and denoising. arXiv preprint arXiv:2206.07188, 2022.
[68] Zhihong Deng, Zuyue Fu, Lingxiao Wang, Zhuoran Yang, Chenjia Bai, Zhaoran Wang, and Jing Jiang. Score: Spurious correlation reduction for offline reinforcement learning. arXiv preprint arXiv:2110.1224 1246. 2021.
[69] Chenjia Bai, Lingxiao Wang, Lei Han, Animesh Garg, Jianye Hao, Peng Liu, and Zhaoran Wang. Dynamic bottleneck for robust self-supervised exploration. Advances in Neural Information Processing Systems, 34:17007-17020, 2021.
[70] Guy Tennenholtz, Assaf Hallak, Gal Dalal, Shie Mannor, Gal Chechik, and Uri Shalit. On covariate shift of latent confounders in imitation and reinforcement learning. arXiv preprint arXiv:2110.06599, 1021.
14

---

[71] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv:1907.02893 2019.
[72] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: 08731, 2019.
[73] Yonggang Zhang, Mingming Gong, Tongliang Liu, Gang Niu, Xinmei Tian, Bo Han, Bernhard Schölkopf, and Kun Zhang. Causaladv: Adversarial robustness through the lens of causality. arXiv preprint arXiv:2106.061
2021.
[74] Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Ré. Conrect-n-contrast: A contrastive approach for improving robustness spurious correlations. arXiv preprint arXiv:2203.01517,, 2022.
[75] Annie Xie, Shagun Sodhani, Chelsea Finn, Joelle Pineau, and Amy Zhang. Robust policy learning over multiple uncertainty sets, In International Conference on Machine Learning. pages 24414-24429. PMLR, 2022.
[76] Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in rein- forcement learning with mixture regularization. Advances in Neural Information Processing Systems, 13:7968-7978, 2020.
[77] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regulariz- ing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
[78] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep Q-learning with convnets and vision transformers under data augmentation. Advances in neural information processing systems,
2021.
[79] Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Auto- matic data augmentation for generalization in reinforcement learning. Advances in Neural Information Processing Systems, 34:5402-5415 -4450, 2021.
[80] Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 1 1361 1-13617. IEEE, 2021.
[81] Chaochao Lu, Biwei Huang. Ke Wang, José Miguel Hernández-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual bassd data augmentation. arXiv preprint arXiv:20/2.09092, 2020.
[82] Shubhankar Agarwal and Sandeep P Chinchali. model-based robotic control. In Conference on Robot Learning, pages 800-811. PMLR, 2023.
Synthesizing adversarial visual scenarios for
[83] Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33:3976-3990 2020.
[84] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference intelligent robots and systems (IROS), pages 23-30. IEEE, 2017.
[85] Nataniel Ruiz, Samuel Schulter, and Manmohan Chandraker. Learning to simulate. arXiv preprint 18700088..002113. 2018
[86] Bhairav Mehta, Manfrre Diaz, Florian Golemo, Christopher J Pal, and Liam Paull. Active domain randomization. In Conference on Robot Learning, pages I 1162-1176. PMLR, 2020.
[87] Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to- sim: Data-efficient robotic grasping via randomized-to-canoone adaptation networks. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition, pages 12627-12637, 2019.
15

---

[88] Yevgen Chebotar, Ankur Handa, Viktor Makoviychuk, Miles Macklin, Jan Issac, Nathan Ratliff, and Dieter Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In 20/9 International Conference on Robotics and Automation (ICRA), pages 8973-8979. IEEE, 2019.
[89] Sergey Zakharov, Wadim Kehl, and Slobodan Ilic. Deceptionnet: Network-driven domain randomization. In Proceedings the IEEE/CVF International Conference Computer Vision,
pages 532-541, 2019.
[90] Christopher Clark, Mark Yatskar, and Luke Zettlemoyer. Don't take the easy way out: Ensemble-based methods for avoiding known dataset biases. arXiv preprint VrXiv:/909.08: 2019.
[91] Divyansh Kaushik, Eduard Hovy, and Zachary C Lipton. Learning the difference that makes difference with counterfactually augmented data. arXiv preprint arXiv: 12434, 2019.
[92] Meike Nauta, Ricky Walsh, Adam Dubowski, and Christin Seifert. Uncovering and correcting shortcut learning in machine learning models for skin cancer diagnosis. Diagnostics, 12(1):40, 2021.
[93] Nimit Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. Advances in Neural Information Processing Systems, 33:19339-1993 2020.
[94] Seonguk Seo, Joon-Young Lee, and Bohyung Han. Unsupervised learning of debiased representations with pseudo-attributes. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16742-16751, 2022.
[95] Gregory Plumb, Marco Tulio Ribeiro, and Ameet Talwalkar. Finding and fixing spurious patterns with explanations. arXiv preprint arXiv:2106.02/12, 2021.
[96] Misgina Tsighe Hagos, Kathleen Curran, and Brian Mac Namee. Identifying spuri- correlations and correcting them with an explanation- based learning. arXiv preprint arXiv:2211.08225,, 2022.
[97] Abubakar Abid, Mert Yuksekgonul, and James Zou. Meaningfully debugging model mis- takes using conceptual counterfactual explanations. In International Conference on Machine Learning, pages 66-88. PMLR, 2022.
[98] Andrea Bontempelli, Stefano Teso, Fausto Giunchiglia, and Andrea Passerini. Concept-level debugging of part-prototype networks. arXiv preprint arXiv:2205. 2022.
[99] Mohammad Taha Bahadori and David E Heckerman. Debiasing concept-based explanations with causal analysis. arXiv preprint arXiv:2007. /5000, 2020.
[100] Alekh Agarwal, Nan Jiang, Sham M Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep, 32, 2019.
[101] Jiayi Weng, Huayu Chen, Dong Yan, Kaichao You, Alexis Dubureq, Minghao Zhang, Yi Su, Hang Su, and Jun Zhu. Tianshou: A highly modularized deep reinforcement learning library. Journal of Machine Learning Research, 23(267):1 1,,, 2022.
16

---

Appendix
Table of Contents
Broader Impact
17
Other Related Works
Related formulations
Related literature of robustmess in RL
B.3 Related literature of spurious correlation in RL
18
C Theoretical Analyses
19
C.I Proof of Theorem
19
Proof of Theorem
21
Proof auxiliary results
24
D)
Experiment Details
27
D
Architecture the structural causal model
................................
27
D.2
Details of Tasks
................................
28
D.3
Example of Generated Data by Perturbations
29
D.4
Computation Resources
30
D.5
Hyperparameters
30
D.6
Discovered Causal Graph in SCM
30
A Broader Impact
Incorporating causality into reinforcement learning methods increases the interpretability of artificial intelligence, which helps humans understand the underlying mechanism of algorithms and check the source of failures. However, the learned causal transition model may contain auman-readable private information about the environment, which could raise privacy issues. To mitigate this potential negative societal impact, the causal transition model needs be encrypted and only accessible to algorithms and trustworthy users.
B Other Related Works
In this section, besides the most related formulation, robust RL introduced in Sec 3.3 we also introduce some other related RL problem formulations partially shown in Figure 3 Then, we limit our discussion to mainly two lines of work that are related ours: (1) promoting robustness RL; (2) concerning the spurious correlation issues in RL.
B.1 Related RL formulations
Robustness noisy state: POMDPs and SA-MDPs. State-noisy MDPs refer to the RL problem that the agent can only access and choose the action based on noisy observation rather than the true state each step, including two existing types of problems: Partially observable MDPs (POMDPs) and state-adversarial MDPs (SA-MDPs), shown in Figure b(... In particular, at each step in POMDPs, the observation or is generated by a fixed probability transition O(- st) (we refer the case that only depends on the state but not action); for anversaavraa MDPs, the observation an adversary v(s)) against and thus determined by the conducted policy, leading to the worst performance by perturbing the state in small set around itself. To defend the state perturbation, both POMDPs, and SA-MDPs are indeed robust to the noisy observation, or called agent-observed state, but not the real state that transitions to the environment and next steps. In contrast, our RSC-MDPs
17

---

propose the robustness to the real state shift that will directly transition to the next state in the environment, involving additional challenges induced by the appearance of -uttt--distriuution states.
Robustness to unobserved confounder: MDPUC and confounded MDPs. To address the mislead- ing spurious correlations hidden in components of RI, people formulate RL problems MDPs with some additional components unobserved confounders. In particular, the Markov decision process with unobserved confounders (MDPUC) [37 serves as a general framework to concern all types of possible spurious correlations in RL problems - at each step, the state, action, and reward are all possibly influenced by some unobserved confounder, shown in Figure dd); confounded MDPs [9 mainly concerns the misleading correlation between the current action and the next state, illustrated in Figure The proposed confonoodf MDPs (SC-MDPs) can be seen as a specified type of MDPUC that focuses on breaking the spurious correlation between different parts of the state space itself (different from confounded MDPs which consider the correlation between action and next state), motivated by various real-world applications in self-driving and control tasks. In addition, the proposed formulation is more flexible and can work in both online and offline RI. settings.
Contexual MDPs (CDMPs). A contextual MDP (CMDP) [38] is basically a set of standard MDPs sharing the same state and action space but specified by different contexts within a context space. In particular, the transition kemnll, reward, and action of CMDP are all determined by a (possibly unknown) fixed context. The proposed robust confounded MDPs (RSC-MDPs) are similar to CMDPs if we cast the unobserved confounder as the context in CMDPs, while different in two aspects: (1) In a CMDP, the context fixed throughout an episode, while the unobserved confounder in RSC-MDPs can vary as {ce IsgsT; (2) In the online setting, the goal of CMDP is to beat the optimal policy depending on the context, while RSC-MDPs seek to learn the optimal policy that does not depend on the confounder (chnggg.
B.2 Related literature of robustness in RL
Robust RL (robust MDPs). Concerning the robust issues in RL, a large portion of works focus on robust RL with explicit uncertainty the transition kernel, which is well-posed and a natural way to consider the uncertainty of the environment [13] 33H488. However, to define the uncertainty set for the environment, most existing works use task atructure-agnostic and heuristic 'distance'
such as R-contamination, KI. divergence, x² and total variation 49 50 14
51
52
15
49
53
58 49
60 to measure the shift between the training and test transition kernel, leading
to homogeneous (almost structure-free) uncertainty set around the state space. In contrast, we consider a more general uncertainty set that enables the robustness to a task desk-dependent heterogeneous uncertainty set shaped by unobserved confounder and causal structure, in order to break the spurious correlation hidden in different parts of the state space.
Robustness in RL. Despite the remarkable success that standard RL has achieved, current RL algorithms are still limited since the agent is vulnerable if the deployed environment is subject to uncertainty and even structural changes. To address these challenges, a recent line of RL works begins to concern robustness to the uncertainty or changes over different components of MDPs state, action, reward, and transition kernel, where review [8] can be referred to. Besides robust RL framework concerning the shift of the transition kernel and reward, to promote robustness in RL, there exist various works [1 12 that consider the robustness to action uncertainty, i.e., the deployed action in the environment is distorted by an adversarial agent smoothly or circumstantially; some works [9] 6 I0 65-67 investigate the robustness to the state uncertainty including but not limited to the introduced POMDPs and SA-MDPs in Appendix [B.T] where the agent chooses the action based observation the perturbed state determined by some restricted noise adversarial attack. The proposed RSC-MDPs can be regarded as addressing the state uncertainty since the shift of the unobserved confounder leads to state perturbation. In contrast, RSC-MDPs consider the out-of-distributtin of the real state that will directly influence the subsequent transition in the environment, but not the observation in POMDPs and SA-MDPs that will not directly influence the environment.
B.3 Related literature of spurious correlation in RL
Confounder in RL. These works mainly focus on the confounder between action (treatment) and state (effect), which is long-standing problem that exists in the causal inference area. However,
18

---

we find that the confounder may cause problems from another perspective, where the confounder is built upon different dimensions of the state variable. Some people focus on the confounder between action and state, which is common in offline settings since the dataset is fixed and intervention is not allowed. But in the online setting, actions are controlled by an agent and intervention is available to eliminate spurious correlation. [68] reduces the spurious correlation between action and state in the offline setting. [69] deal with enviromment-irrelevan white noise; possible shift + causal [70]. The confounder problem is usually easy to solve since agents can interact with the environment to do interventions. However, different from most existing settings, we find that even with the capability of intervention, the confounding between dimensions in states cannot be fully eliminated. Then the learned policy is heavily influenced if these confounders change during testing.
Invariant Feature learning. The problem of spurious correlation has attracted attention in the supervised learning area for a long time and many solutions are proposed learn invariant features to eliminate spurious correlations. A general framework to remedy the ignorance of spurious correlation empirical risk minimization (ERM) is invariant risk minimization (IRM) 71. Other works tackle this problem with group distributional robustness 772. adversarial robustness 773,, and contrastive
learning [74]. These methods are also adapted to sequential settings. The idea of increasing the
robustness of RL agents by training agents on multiple environments has been shown in previous works 775 30 30 However, a shared assumption among these methods is that multiple environments with different values of confounder accessible, which is not always true in the real world.
Counterfactual Data Augmentation in RL. One way to simulate multiple environments is data augmentation. However, most data augmentation works [24] 76 25 77-80 apply image transfor- mation raw inputs, which requires strong domain knowledge for image manipulation and cannot be applied to other types of inputs. In RL, the dynamic model and reward model follow certain causal structures, which allow counterfactual generation new transitions based on the collected samples. This line of work, named counterfactual data augmentation, is very close to this work. Deep generative models [8]] and adversarial examples
[82] are considered for the generation to improve sample efficiency in model-based RL. CoDA [83 and MocoDA (33) leverage the concept of locally factored dynamics to randomly stitch components from different trajectories. However, the assumption of local causality may be limited.
Domain Randomization. we are allowed control the data generation process, e-g., the underlying mechanism of the simulator, we can apply the golden rule in causality Randomized Controlled Trial (RCT). The well-known technique, domain randomization [84]. exactly follows the idea of RCT, which randomly perturbs the internal state of the experiment in simulators. Later literature follows this direction and develops variants including randomization guided by downstream tasks in the target domain 88 [6), randomization to match real-works distributions [87 88), and randomization to minimize data divergence [89]. However, is usually impossible to randomly manipulate internal states in most situations in the real world. In addition, determining which variables to randomize is even harder given so many factors in complex systems.
Discovering Spurious Correlations Detecting spurious correlations helps models remove features that are harmful to generalization. Usually, domain knowledge is required to find such correlations (90) [2]. However, when prior knowledge is accessible, techniques such as clustering can also be used to reveal spurious attributes [37
93 94] When human inspection is available, recent works [95 also use explainability techniques to find spurious correlations. Another area for discovery is concept- level and interactive debugging 988 99], which leverage concepts or human feedback to perform debugging.
C Theoretical Analyses
C.1 Proof of Theorem
1
In this section, we verify the existence of an optimal policy of the proposed RSC-MDPs, involving
additional components - confounder C, and the infimum optimization problems with comparisons to standard MDPs [100).
19

---

To begin with. we recall that the goal is to find a policy क fassesr E II such that for all (s,a) EEEXX A:
(s) (s)
sup Tell
V. ( (s) and (s,a) = (s,a)
sup well
"
(s, (s,a), (8)
which we called an optimal policy. Towards this, we start from the first claim in (8)
Step 1: Introducing additional notations. Before proceeding, we let {Ss,, Rr, Ct} denote the random variables state, action, reward, and confounder, at time step t for all S - T. Then invoking the Markov properties, we know that conditioned on current state the future state, action, and reward are all independent from the previous $1, d1, 71, C1, 8t-1, at-1, 1t-1, Ct-1. In addition, we represent P. E (C)) as some distribution of confounder at time step t, for all 1 S t S T. For convenience, we introduce the following notation that is defined over time step S T:
ME P- 8ickeTPa and U (P++) = BiSKSTUL (PE),
which represent some collections of variables from time step to the end of the episode. In addition, recall that the transition kernel from time step / + is denoted as 8/41 Pi( 88, (t) for i E 1,2... n. With slight abuse of notation, we denote 8141 ~ P.(- 8t, at (t) and abbreviate
E. Pe (e)
[+ whenever it is clear.
Step Establishing recursive relationship. Recall that the nominal distribution of the confounder
E Pf at time step L. We choose से {Te) which obeys: for all < 1
}
(s) arg max Ex, (T,,,r((,, + inf N,EA(A)
Ex,
Banrs
B
(10)
PiEu" (PT)
Armed with these definitions and notations, for any (1,8) E f1,2,,... ,T} x S, one has
V.
T
sup inf sell (Pc)
P.. (s)
sup inf sell E. (P+.)
Er,
Σ
Erass
(ii) sup
inf
Er, az)
well Pur (Pq.)
T
Basti
Σ k=t+1
T& (sk, ax) n, P+(++), (St, A,, Re, S:41) (8, Ct 8141)
sup En (r)( =ell
+
inf Pue (Pz.)
Ex, Beww
T
Esc
Σ
Tk(s,, ax) m, P+(+-1) (St, Ar, Rrs Ci, Se41) Te, 841)
where (i) holds by the definitions in ( (ii) is due to 3 and that V.T.P (8) only depends on P++ by the Markov property, (iii) follows from expressing the term of interest by moving one step ahead and Es, is taken with respect a, -T(- 1 s).
To continue, we observe that the V+ (s) can be further controlled as follows:
(s))
sup E,, (r)) xell
inf P.reur (rş.)
Ex. EeywP
E
Σ b=t+1
ra(s,, ax) m. P+(--)) (St, Ae, Re, Ct, S441) (8,01, (s,00,
20

---

sup En, (r((s, Fen
inf (P()
Eus Pe Bott1
inf
Th (a( (x) n,, P+(++1) (St, Ar, Re, Ct,St+1) = (8, 04s rts Ct, 8141)
Eu (Mo))) k=1+1
sup En, [r((s, + xell
inf (PP)
En. Ee
sup x'ell P (Ur (P())))
inf
(k(s ax)l P+(++1); (St, Ar, Rr, Cr, St+1) = (s, 0e, Tt: Cc,
k=1+1
(1)
sup Em, (rr(s, at) nell
T
inf (PE)
Exc
Bai. Eai1
sup r'ell EL
inf
Ex, Ex P+cet))
Σ k=
(sk: ax)
Exi BenPs Boem [Vti Viti(seti)]]
sup xell
Ex,, +
inf PiEle (PC)
Bsg Bear [Becs [Vti(eeti]]]
sup *<EA(AA
Ex, [r, +
inf (PE)
inf (PE)
E re(s, a.) Ee
[Ti la (
(11)
P Prom (Pian)) is independent from πT conditioned on fixed distribution of St+11 (ii) arises from the Markov property such that the rewards {ru(Sk ag) h++155T conditioned on (St,A, Rr, S241) or S441 are the same, and the last equality follows from the
where (i) holds by the operator inf
definition of के in 10
Step Completing the proof by applying recursion.
Applying
11 recursively for + 1,.... T, we arrive at (r((s,a)) ac) + Ee Eaiis
V
<
inf Pell"
E
Vin 441 (s0+1)
|oe (+))
inf (PF) Pareu (PT))
inf
E rr( EerrP [Eatt.]
Te41(8:11, a441) + BE Bou4z V/ ((22()))2)
(141)= (c), (F)(s)), (
5...g inf
Ex,P.
Σ k=d
(x)
1/2.0())
(12)
Observing from 12 that
VsES: V.r < F7.0 (s) I sup 50.0 (8)
-0)))(((())) (s),
(13)
for all S. The second assertion in ( can be achieved analogously. Until now, we verify that there exists at least policy ãे that obeys 8 which we refer as an optimal policy since its value is equal to larger than any other non-stationary and stochastic policies over all states E S.
which directly verifies the first assertion in 8 V.a. (s) = V.
C.2 Proof of Theorem
We establish the proof by separating it into several key steps.
21

---

0
a
0
1-4
[0.0]
[0.1]
r=
r=1 1
[0.0]
[0.1]
r=
0
=
a=
[1,0]
[1,1]
r=1
r=0
[1.0]
[1,1]
r=1
(a) Standard MDP kemel Pt
SC-MDP kemel Pı
Figure 8: The illustration of the transition kernels of the standard MDP M and the proposed SC-MDP Mac at the first time step t = i.. Pf and P1 respectively.
Step 1: Constructing hard instance .M of standard MDP. In this section, we consider the
following standard MDP instance . = {S, .. p ,T.r) where = {[0, 0]. [0. 1], [1. 0j. [1, 1]] is the state space consisting four elements in dimension n 2, and A= (0. with only two options. The transition kernel p {PP) at different time steps < t < 7 is defined as
the action space
1(s' = [0. (a) 1(s = 10. (a) if 00.0) 1(s' ==)
(P(5'18,00
(14)
otherwise
which is illustrated in Fig. a), and
( 8, 1(s' s), 8(,,, E (2,3, x
(15)
Note that this transition kernel pD ensures that the next state transitioned from the state [0, 0] is either [0, 0] or [0, 1j. The reward function is specified as follows: for all time steps < S.
{
if [0, 0J [1,
as
(16)
0 otherwise
Step 2: The equivalence between M and one SC-MDP. Then, we shall show that the constructed standard MDP M can be equivalently represented by one SC-MDP Mac {S,A, (S,,,T,r, T,,, C, (P;), Pd)
with C := {0, 1}. The equivalence is defined the sequential observations {st, de, IS)STT induced by any policy and any initial state distribution in two Markov processes are identical. To specify, S, are kept the same as M. Here, {Pi} shall be specified in while, which determines the transition to each dimension of the next state conditioned on the current state, action, and confounder distribution for all time steps, i.e., 8/41 Re [Pi(- 8t, at at)c for any i-th dimension of the state (i 2 and all time step S 1 T. For convenience, we denote = [P] P?] A(8) as the transition kernel towards the next state, namely, 81+1 ~ Ea, [P 8t, az cz)].
Then we simply set the nominal distribution of the confounder as follows:
(et) = 1(c =), Vi EC
(17)
In addition, before introducing the transition kernel (Pi} of the SC-MDP Mse, we introduce auxiliary transition kernel p = (P.C) } as follows:
{ 1 = 11,0))((a 0) + 1 = (a,1))((a= =1) 1) (s, = (10,0),0) 1 (ss 8)
(s' a,
otherwise
(18)
and
(r' s, 1(s'- s), V(t,s,a) (2, ... .T) x S A.
(19)
It can be observed that pse is similar to pe except for the transition in the state [0, 0j.
22

---

Armed with this transition kernel pmr, the {Pt) of the SC-MDP Mse is set to obey
(,) - { Pi(s' 8 (1,s'1s,,,,
(1- 11-a)P P (s,0) + P = 00
(20)
1(s'=
otherwise
which is illustrated in Fig. 8)b. and
Pa(s' (, = 1(' V(t,s,,,a) E 12,3, T) X
(21)
With them in mind, we are ready to verify that the marginalized transition from the current state and action to the next state in the SC-MDP Msc is identical to the one in MDP M: for all (f, St, at, 8641) E [1,2, ... 7} x S x A S:
P((8+1 St, (x) = PF [Pr( 8t, Ct, a)] = (r 8r, oz 0) = p (8c+1 8t, az)
(22)
where the second equality holds by the definition of pa (1). and the last equality holds by the definitions of P (see (20) and (2T).
In summary, verified that the standard MDP M {8, A, Po, Tr is equal to the above specified SC-MDP MMs.
Step 3: Defining corresponding RMDP and RSC-MDP. Equipped with the equivalent standard
MDP M and SC-MDP Mae. we consider the robust variants of them respectively - a RMDPPMMM. Mrob =
{S, AS (Pº), Jrrr with some uncertainty level 01, and the proposed RSC-MDP Msc-rob (8,,, C, (Pi))22=(PP
with some uncertainty level 02.
In this section, without loss of generality, we consider total variation as the 'distance' function p for the uncertainty sets of both RMDP Mrob and RSC-MDP Mse-rob, i.e., for any probability vectors P.P E ACC (or P E A(S)), (P', = |P' -Pl Consequently, for any uncertainty level a E [0, 1], the uncertainty set 4/"1 (Po) of the RMDP (see (T) and 4 (Po) of the RSC-MDP Moc.rob (see
are defined as follow, respectively:
(P(P))== =0 60 1 U (aa))
1 u (Pt = PPis E A(S)
=fPiso
2 PP Po Phnalll so
4PPP)) ( [PA(C)) (pE)))(EEC)
1 11P : Pfll s so
=P(P))=
(23)
Step 4: Comparing between the performance of the optimal policy of RMDP Mrab (
RMMD 4:00
and that of RSC-MDP Macrob ). To continue, we specify the robust optimal policy RRMDP associated with Mrab and TRSC associated with MMsc.rob and then compare their performance on RSC-MDP with some initial state distribution.
4.02
To begin, we introduce the following lemma about the robust optimal policy R associated with the RMDP Mrob.
Lemma 1. For any 0i E (0, 1], the robust optimal policy of Mrob obeys
Vs S:
x [TMDPP
(0] = 1.
(24a)
In addition, characterize the robust SC-value functions of the RSC-MDP Msc-rob associated with any policy, combined with the optimal policy and its optimal robust SC-value functions, shown in the following lemma.
Lemma 2. Consider any 02 ( 1] and the RSC-MDP Mscrob (S,,,, C, (P)), For any policy n, the corresponding robust SC-value functions satisfy
V..022 (0,00) = 11 (T
Edrp T (0000000) (1)(1)))))))--) m PEU"
inf
(25a)
In addition, the optimal robust SC- Salve function and the robust optimal policy R of the RMDP M rob obeys:
2,02
- (1.00) = T-1
V RSc d2 (10, 0])
(26)
2
23

---

Armed with above lemmas, applying Lemma 2 with policy π in Lemma
RRMDP obeying [=
[RMDDP] [0| () =
[ one has
V FAMMP 12 ([0,0)) 11 (T -
inf (PT)
BeirP
1 c-
[1
3
T 1
2 (T -
(27)
4
where the inequality holds by the fact that the probability distribution P obeying Pi(0) = and
Pi(1) = 3 is inside the uncertainty set 4 (Pf) (recall that 02 (g) 1] and Pf(0) = 1).
Finally, combining (27 and (26) together. we complete the proof by showing that with the initial state distribution @ defined as @[ 0)) we arrive
T-1 T
Vi TRUBO 22 (10, 0]) V. (Ø) Vi RRMOO 22 010,0)= 2
(28)
4
8
where the last inequality holds by T> 2.
C.3 Proof of auxiliary results
C.3.1 Proof of Lemma i
Step I: specifying the minimum of the robust value functions over states, For any uncertainty set
σ E (0, 1). we first characterize the robust value function of any policy over different states. To start, we denote the minimum of the robust value function over states each time step & as below:
V
min Va 0,
(29)
SES
where the last inequality holds that the reward function defined in (66) is always non-negative. Obviously, there exists at least one state smin, that satisfies V. (smin.)
min.t"
With this in mind, we shall verify that for any policy n,
VI S & : Vil.a ((0, 1J) = V. (I,00) =
(30)
To achieve this, we use recursive argument. First, the base case can be verified since when + 1 = T + 1, the value functions are all zeros at T + step, i.c., V741 (s) = 0 for all s E S, Then, the goal is to verify the following fact
V. ([0, 1])
V. ([1,0)
(31)
with the assumption that V441 (s) for any state { 1), [1, 01}. It easily observed that for any policy т, the robust value function when state {10, 1], [1, at any time step t obeys
0 S (s) Eawt () re(ssa)+ +
inf PEU
PVT.
0+ (1 0. 1 (5) + σ V=.0
0 a VT.01
min.t+1
< 0 1 17 (s) =0
(32)
where (i) holds by rx(8.0) a) 0 for all {(0, 1]. [1, 0J). the fact (
for E S (see
(14) and (15),, and the definition of the uncertainty set (1: ( (PD) 23 Here (ii) follows from the
recursive assumption ++1 (s) = 0 for any state {10, 1). [1, 0.. and the last equality holds by
Vo.. "in,/+1 I VVT 1)) (see 29 ). Until now, we complete the proof for (31 and then verify (30).
Note that 30 direcly leads to
LA VI
V min.t 0.
(33)
Step 2: Considering the robust value function at state [0, 0). Armed with the above facts, we are now ready to derive the robust value function for the state [0, 0].
24

---

When 2 T. one has
V, 0]) am m ((0.0))) Fr([0,0).a)
inf PPu.o
PV41
1 (1 (1--o)\V 44 ((0,)) + Vinnt min.c+1
1 + (1 on V441 ([),0)
(34)
where (i) holds by ().
for all E {0, 1) and the definition of p (see 15 ), and the last
equality arises from 33
Applying 34 recursively for L, 1 1, ,T yields that
Vi v. = (1-00))~1 1.
(35)
When = 1, the robust value function obeys:
V/T:*(((,,0) = Eaam (: (1000))
(r((0,0).g +
inf
PV2
PEur:(Ps,
@ 1 1 [0] (P
inf
PV +T +1| PP
inf
PV
1 (0) 10, 0)) [1 - (1-0;)13 01 " ((0,0) 0]) + V
+ (1 [0, 0)) (1 (1)' " ((0, 1)) min,2
= =1+ n (0) 00 0) -1)1 O))
(36)
where (i) holds by rı ( (0,0).a) a) = for all E f0. 1), (ii) follows from the definition of po (see (4p), and the last equality arises from 30 and (33)
Step 3: the optimal policy TRMDP- Observing that Vi ([0, 0J) is increasing monotically as , (0) [0, 0]) is larger, we directly have that +.00 RR [0, 0)) = 1.
Considering that the action does not influence the state transition for 2, 3, T and all other states s + [0, OJ, without loss of generality, we choose the robust optimal policy as
Vs S: RMDPP (0|(s) s)
(37)
C.3.2 Proof of Lemma
To begin with, for any uncertainty level 02 E (1) 1] and any policy क = {me) J. we consider the robust SC-value function V,T.00 of the RSC-MDP Mse-rob-
Step 1: deriving t. for st T. Towards this, for any 2 < t 1 T and S, one has
V (s) (
i (e)x
re(ssa)+ +
inf (PU)(
Baar Piisac az 6+1
(((s,a) +
BanP Propa Vin PEU
inf
= r, a) + (s),
(38)
where (i) follows from the confoudedd Bellman consistency equation in (47). (ii) holds by that
the reward function r, and P, all independent from the action (see (16) and (1 and the last
inequality holds by (s s, a, c.) = ((s' = s) is independent from c (see 21
Applying the above fact recursively for 1, ..... . .T leads to that for any E S,
: =r((sa))
V (s) re(s, a) + (s,1 a(41) (s)
25

---

T
0t) +
Σ
Tk(Kk, ax);
(39)
k=t+1
which directly yields (see reward r in 16
V2 (10,0) 1/2 ([1, 1J) T -1 and V2 ((0, 11) " ([1,0) 0.
(40)
Step 2: characterizing V/T T. ((0, 0)) for any policy n. this section, we consider the value of V 2 on the state [0, 0j. To proceed, one has
V (10, 0]) Bawni ((0,01)
@1 (10,0,
E Ba (10,01)
ri (r((0, 0].
Eaop (P1.0000,a.cc v2 (EU((PF)
inf
iii
+
inf
Eai-P
(T] P,,00,00 + +T (0, (.0)..i.)) 12
PEt
(i))
1+ inf PEl (PT)
BanP
T1 [o, 0]) (11 (11 (1--)) 1,1000.0 + ci P 1.0.0..) ii
Ti (1)(00))0 (11 0i P)) 1,000011 + t.
I((0.0).1 ) 7
(iv)
inf (PT)
TT (000000,0 (11 (1-c)) (1- ((0,0) EaNP
+ (0.)
(,,22(111000) ([1,0))
+ +1)0100000 (11 (11- ((0 [J) + a V. n. ( [1)
=1+ -T-1) inf
EapP क (0] (0) ai + M [0, 0)ecr
Petl" (PF)
+ (T - 1)= (0. 0)) (T-
inf (PF)
E P (1 2z, [0. 01)
(41)
where (i) holds by robust state-confounded Bellman consistency equation in 47 (ii) follows from Ti (10, 0], = for all E {0, 1} which independent from c₁. (iii) arises from the definition of P in (2),, (iv) can be verified by plugging in the definitions from (14) and (18). and the penultimate equality holds by (40).
Step 3: characterizing the optimal robust SC-value functions. Before proceeding, we recall the fact that 41 (Pf) (P E (C) |P pill S 02).
Observing from 41 that for any fixed क (0 [o, 0)). e (1 (1-2x, (0 [o. 0)) is monotonously increasing
with ej when 2 (0 [o, 0)) 2 0 and decreasing with e₁ otherwise, is easily verified that the maximum of the following function
(0 0)) -T inf Eg~P [c - 2m (0 [o, 0))]
(42)
Pell(
obeys
0
T ( [0,0) 2
max (mx ()) = {(r ( 1)00 2m (0 (0)0))
(43)
otherwise
Then, note that the value _.. "(0," only depends on क (: (-1 0]) which can be represent by Ti (0 [o, 0]). Plugging in (43) into (41) arrives when (0 [0. 0)) 2 }
V V ((..0)
max
x [01
26

---

maX [0 [0.0))
1 (T () (0 [0, 0]) + (T 10)0 2 (0) [0.0]]
1 (7 - 1)02 (T -1)
(1 20 Ii (0 [0,0] 0))
11 [0,01]2
1 -T- + (T 1)(1 2022
--1 =1+
(44)
2
2
where the penultimate equality holds by 02 and letting (0) [0, 0J) (1 [0, 0])
Similarly, when
2"
T-1
maaX (0.0))
((,00
max <0.0))
1+ (1 <1 2
(45)
Consequently, combining and 445 we conclude that
T
1 RSC ((0,0) Vi ((0,0) max V 20 ([0,0) =1+
(46)
2
C.3.3 Auxiliary results of RSC-MDPS
easily verified that for any RSC-MDP Msc-rob (8, A, C. (2) any policy and optimal policy n" satisfy the corresponding robust state-confoudded Bellman consistency equation and Bellman optimality equation shown below, respectively:
(s,a) =)) inf (PE)
EaeP (Praaa Vin 1+1
(s,0) = a)) + inf (PU)
Eanp [Proaa V.oo
(47)
= where Precae E Rxxs such that (s)) = Pi(s'|s,a,co for S. and Vir
SUPAEA(A)) (TT(I1(s, (z)) inf (P Er, BaND (Pisaa Veti (8441)
D Experiment Details
D.I Architecture of the structural causal model
position embedding
position
Causal Graph
embedding
Encoder
feature
feature
G
State
Decoder
+
Encoder
Next State
Encoder
matrix multiplication
Decoder
Action
Decoder
Reward
Encoder
13 [x++))
Figure 9: Model architecture of the structural causal model. Encoder, Decoder, position embedding, and Causal Graph are learnable during the training stage.
We plot the architecture of the structural causal model we used in our method in Figure In normal neural networks, the input treated as whole to pass through linear layers or convolution layers. However, this structure blends all information in the input, making the causal graph useless to separate cause and effect. Thus, our model, we design an encoder that shared across all dimensions the input. Since different dimensions could have exactly the same values, add learnable position
27

---

embedding to the input of the encoder. In summary, the input dimension of the encoder is 1 + dpoc. where dpos the dimension of the position embedding.
After the encoder, we obtain set of independent features for each dimension of the input. We now multiply the features with learnable binary causal graph G. The element (i, of the graph is sampled from a Gumbel-Softmax distribution with parameter Øi, to ensure the loss function is differentiable w.r.t é.
The multiplication of the causal graph and the input feature creates a linear combination of the input feature with respect the causal graph. The obtained features then passed through decoder to predict the next state and reward. Again, the decoder is shared across all dimensions to avoid information leaking between dimensions. Position embedding is included in the input to the decoder and the output dimension of the decoder is 1.
D.2 Details of Tasks
We design four self-driving tasks in the Carla simulator 22 and four manipulation tasks in the Robosuite platform [23]. All of these realistic tasks contain strong spurious correlations that are explicit to humans. We provide detailed descriptions of all these environments in the following.
Brightness
Behavior
Crossing
CarType
Day
Wait
Red light
Slow
Ego
go
Ege
Overtake
uummnn uniuiii
Night
Cross
Fast
Leading
Green light
Ego
Ego
Figure 10: Illustration of tasks in the Carla simulator.
Brightness. The nominal environments are shown in the 1 column of Figure 10 where the brightness and the traflic density are correlated. When the ego vehicle drives in the daytime, there are many surrounding vehicles (first row). When the ego vehicle drives the evening, there is no surrounding vehicle (second row). The shifted environment swaps the brightness and traffic density in the nominal environment, i.e., many surrounding vehicles in the evening and no surrounding vehicles in the daytime.
Behavior: The nominal environments shown in the 2nd column of Figure 10 where the other vehicle has aggressive driving behavior. When the ego vehicle is in front of the other vehicle, the other vehicle always accelerates and overtakes the ego vehicle in the left lane. When the ego vehicle is behind the other vehicle, the other vehicle will always accelerate. In the shifted environment, the behavior of the other vehicle is conservative, i.c., the other vehicle always decelerates to block the ego vehicle.
Crossing. The nominal environments shown in the 3rd column of Figure 10 where the pedestrian follows the traffic rule and only crosses the road when the traffic light is green. In the shifted environment, the pedestrian disobeys the traffic rules and crosses the road when the traffic light is red.
CarType. The nominal environments are shown in the 4th column of Figure where the type of vehicle and the speed of the vehicle are correlated. When the vehicle is truck, the speed is low and
28

---

when the vehicle is motoreycle, the speed is high. In the shifted environment, the truck drives very fast and the motoreycle drives very slow.
Lift
Stack
Wipe
Door
Right
Low
Left
Righh
Green
Right
Diagonal
Right
High
Left
Lueft
Anti- diagonal
Right
Left
Left
Figure 11: Illustration of tasks in the Robosuite simulator.
Lift. The nominal environments are shown the 1hh column Figure where the position of the cube and the color of the cube are correlated. When the cube is in the left part of the table, the color of the cube is green, when the cube is in the right part of the table, the color of the cube red. The shifted environment swaps the color and position of the cube in the nominal environment, i.e., the cube is green when it is in the right part and the cube red when it is in the left part.
Stack. The nominal environments are shown in the 2dd column of Figure 11 where the position of the red cube and green plate are correlated. When the cube in the left part of the table, the plate also in the left part; when the cube is in the right part of the table, the plate is also in the right part. In the shifted environment, the relative position of the cube and the plate changes, i.c., When the cube is
in the left part of the table, the plate is in the right part; when the cube is in the right part of the table, the plate is the left part.
Wipe. The nominal environments are shown in the 3rd column of Figure where the shape of the dirty region correlated the position of the cube. When the dirty region diagonal, the cube is on the right-hand side of the robot arm. When the dirty region is anti-diagonal the cube is on the left-hand side of the robot arm. the shifted environment, the correlation changes, i.e., when the dirty region is diagonal, the cube is on the left-hand side of the robot arm; when the dirty region is anti anti-diagonal, the cube on the right-hand side of the robot arm.
The nominal environments are shown in the 4th column of Figure 11 where the height of the handle and the position of the door are correlated. When the door is closed the robot arm, the handle is in a low position. When the door is far from the robot arm, the handle is in a high position. In the shifted environment, the correlation changes, i.e., when the door is closed to the robot arm, the handle is in a high position; when the door is far from the robot arm, the handle is in low position.
Door.
D.3 Example of Generated Data by Perturbations
We show an example of generated trajectories in the Lift task to demonstrate the reason why our method obtains robustness against spurious correlation. In Figure 12 (a), we show the collected trajectories from the data buffer. Since the green block is always generated on the left side of the
table, the trajectories of the green block mainly appear on the left side of the table. In Figure 12 (b). we generate new trajectories from trained transition model and we observe that the distribution of trajectories follows the collected data. In Figure 12 (c), we directly perturbed the dimensions of the state and used the same transition model to generate new trajectories. We find that the generated trajectories blend the color but fail to maintain the spatial distribution of the original data. In Figure 12
29

---

(d), we use the causal- based transition model to generate new trajectories and we find that the results only follow the spatial distribution but also blend the color.
The results shown in Figure T2 illustrate that the data generated by our method eliminates the spurious correlation between the color and position of the block, therefore, enabling the policy model to generalize to the shifted environment.
Nominal (tainig) envionment
(woding) anviooment
Togectay harm
Gereradac tranddry we paraatalioo
wht cauee graph of pertetarion
forw cound grech parkatation
Ours
The generated transition data from different perturbation methods. (a) Trajectories collected from the policy interact with the nominal environment. (b) Generated trajectories without any perturbation. (c) Generated trajectories with perturbation but without the causal graph. (d) Generated trajectories with perturbation and with the causal graph.
Figure 12:
D.4 Computation Resources
Our algorithm is implemented on top of the Tianshou [101 package. All of our experiments are conducted on machine with an Intel i9-9900K CPU@3 (16 core) CPU, an NVIDIA GeForce GTX 1080Ti GPU, and 64GB memory.
D.5
Hyperparameters
We summarize all hyper- sppee-parameters used in the Carla experiments (Table 5) and Robosuite experi- ments (Table
D.6
Discovered Causal Graph in SCM
To show the performance learned SCM, we plot the estimated causal graphs of all experiments in Figure 13
Figure
[4 Figure 15 Figure 16 and Figure 17
30

---

Table in Carla experiments
Environment
Parameters
Notation
Brightness
Behavior
Crossing
CarType
Horizon steps
T
100
100
100
100
State dimension
24
12
12
12
Action dimension
d
2
2
2
2
Max training steps
1x10
1x10
5x10
5x10
Weight of ||G|],
0.1
norm of || G
0.1
Actor learning rate
3 10
Critic learning rate
1 x 10
Batch size
256
Discount factor
7 SAC
0.99
Soft update weight
SAC
0.005
Weight of entropy
in SAC
0.1
Hidden layers
[256, 256, 256]
Returns estimation step
4
Buffer size
1 x 10
Steps per update
10
Table 6: Hyper-parameters in Robosuite experiments
Environment
Parameters
Notation
Lift
Stack
Door
Wipe
Horizon steps
T
300
300
300
500
Control frequency (Hz)
20
20
20
20
State dimension
50
110
22
30
Action dimension
d,
4
Controller type
OSC position
OSC position
Joint velocity
Joint velocity
Max training steps
1x10"
5x10
1x
1x10
Weight of |G]|,
A
0.01
norm of | lp
P
0.1
Actor learning rate
x
Critic learning rate
x 10
Batch size
128
Discount factor
in SAC
0.99
Soft update weight
in SAC
0.005
alpha learning rate
Ira in SAC
x 10-4
Hidden layers
[256, 256, 256]
Returns estimation step
Buffer size
x 10
Steps per update
10
31

---

Brightness
Behavior
Crossing
CarType
Figure 13: Estimated Causal Graphs of four tasks Carla.
32

---

add pos3
ade
d'oppricate pog2)
piopercate (o41)
cond5
aule cosel
part sun out
10000000
joet bou csd AnO
par pos eagg
jor pos
jorn poe ends
jart pee ech
per witt
pirt wit
jest wiTl
ouf peati
set aaga
qpect
stope spo
grope eped
grpper oold
pippe
greper owe
groper owell
guper
con
Figure 14: Estimated Causal Graphs of the Lift task in Robosuite.
33

---

ode A
dde AsM
cde
cass quatt
cass quart
plowe soa
plice
por odr
serpa ung
por wel
relowt
selpoii
of turt
ed sart
orger und
usper poe
grppe tovi
prope qall
grow qate
glow qal
usee pid
gape R
rita o que
D cast
mnn and
per pak cat
terten witll
gat paripe unti
odd cask
oncer mosk
ont hape sore
ond sastt
coo ()
pust poa
yor
Figure 15: Estimated Causal Graphs of the Stack task in Robosuite.
34

---

door posSt
door
honcle poct
huecle pen2)
Eaga spoc
pelss
14 pes wc
perpes wed
Jed will
pär w
with
por will
qest
pow g
grow qee
prrow p
oteer owe
ateer owt
pen will
3000 80000030
ves with
vet with
ge with
Figure 16: Estimated Causal Graphs of the Door task in Robosuite.
35

---

Figure 17: Estimated Causal Graphs of the Wipe task in Robosuite.
36


Reviews:
[{"review_id": "7H8QMYemv7", "reviewer": "Reviewer_3Qte", "summary": "The paper studies how to develop more robust reinforcement learning algorithms when spurious correlation exists in the observation space. The paper studies the problem from a causal perspective and present robust state confounded MDP as the problem formulation. The paper proposes an algorithm to solve this problem formulation by learning the structural causal model. The paper demonstrate experimentally that the proposed method is more robust to shift in the environment compared to baselines.\n", "questions": "Please see my questions about related work", "limitations": "Please see my questions about related work.\n\nI would have like to see the limitations that the method currently only work for state-based policies more prominently displayed and discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is mostly extremely well written and easy to follow. The structure of the paper makes sense, which first discusses the formulation of the problem in the form of robust state confounded MDP and comparison to existing formulation. The algorithm the paper proposes is intuitive and the improvements over the baseline in the experiment section are clear.\n", "weaknesses": "The discussion around related work focuses on a few most relevant papers and does not address the broader literature if I am not mistaken. The problem of spurious correlation is known within the community and various prior works [1, 2]. Please see [2] for example, which also learns the dynamics model to alleviate the issue of spurious correlation. These prior works tackle this problem even though they did not propose the same robust state confounded MDP that the paper proposes. I would have liked to see a more thorough comparison with the broader literature in a related work section.\n\nOn page 3, the paper writes “people usually prescribe the uncertainty set of the transition kernel using a  heuristic and simple function rho with a relatively small sigma”. However, because of the lack of a related work session, I do not see evidence for this statement.\n\n[1] https://ben-eysenbach.github.io/rpc/\n[2] https://arxiv.org/abs/1909.11373\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude to the reviewer for their insightful feedback. We are glad to know that the reviewer recognizes the novelty of our contributions, the clarity of our problem formulation, and the empirical algorithm that sufficiently shows the advantages compared to baselines. We provide our response to the questions below.\n\n\n### **Q1: Comparisons with broader literature such as [1][2].**\nThanks for the important question and providing the important related works [1][2].\n*  We have included a thorough related work section in Appendix B in the original manuscript with comparisons to literature about 1) other related RL formulations, 2) robustness in RL investigations, and 3) spurious correlation in RL. Within Appendix B.3, we discussed different spurious correlation types that have been considered in the RL community.\n* We will definitely add [2] as a broader reference as the reviewer suggested since it also addresses spurious correlation in RL. We want to note that [2] is kind of far from our topic since 1) [2] sought to solve multi-task RL, while we focus on single-task RL; 2) the spurious correlation that [2] deals with is largely different from the one considered in this work: [2] considers the spurious correlation between the dataset distribution and the dataset identity to design a better task inference module, while the spurious correlation considered in this work is between different states.\n*  After carefully reading [1], we found that it is not quite related to spurious correlation but is implicitly related to robustness in RL. So we will add it to the related section in Appendix B.2. [1] sought to learn a 'simple' policy by minimizing the used information to seek better performance in standard RL, which turns out to have some robustness benefits.\n\n\n### **Q2: Related works about robustness in RL.**\nFor the claim mentioned by the reviewer \"people usually prescribe the uncertainty set of the transition kernel using a heuristic and simple function rho with a relatively small sigma\" on page 3, we add references [3][4] to the end of the sentence to support. A more detailed review of the existing investigated uncertainty set can be referred to in Appendix B.2, where we thoroughly summarize that most existing works use task structure-agnostic and heuristic 'distance' such as KL divergence and total variation.\n\n\n### **Q3: More discussion about the potential of the proposed method which currently only work for state-based policies.**\nWe add more discussion in the last section:\n> The current method is of great potential to be applied to more complicated problems. In particular, the proposed method requires swapping the dimensions of states to break spurious correlation, where it implicitly assumes that each dimension has semantic meaning. In more complicated problems such as high-dimensional image states, each dimension becomes a pixel without semantic meaning anymore and the dimension size may explode. To extend our method to such cases, we can leverage existing state abstraction techniques [5,6] to project images to a low-dimensional latent space, where each dimension represents a semantic feature. Then, our method can be straightforwardly applied to the low-dimensional latent space.\n\n    [1] Robust predictable control. Eysenbach, Ben, Russ R. Salakhutdinov, and Sergey Levine. Neurips 2021\n    [2] Multi-task batch reinforcement learning with metric learning. Li, Jiachen, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Henrik Christensen, and Hao Su. Neurips 2021\n    [3] Toward theoretical understandings of robust Markov decision processes: Sample complexity and asymptotics. Yang, Wenhao, Liangyu Zhang, and Zhihua Zhang. The Annals of Statistics 50.6 (2022): 3223-3248.\n    [4] Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity. Shi, Laixi, and Yuejie Chi. arXiv preprint arXiv:2208.05767 (2022).\n    [5] Disentangling by factorising. Hyunjik Kim and Andriy Mnih. In International Conference on Machine Learning, pages 2649–2658. PMLR, 2018.\n    [6] A theory of state abstraction for reinforcement learning. David Abel. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 9876–9877, 2019.\n"}, {"review_id": "ZuO4hfTM6m", "reviewer": "Reviewer_uQEX", "summary": "This paper aims to address the spurious correlation challenge that arises in RL. Such correlation is typically useless to decision-making but may be learned by agents, which leads to failure in applying to unknown test cases. To this end, the authors proposed a novel RSC-MDP framework that models the spurious correlation challenge with an unobserved confounder. The authors then justify that previous robust MDP methods may fail on the RSC-MDP framework due to inaccurate uncertainty modeling. The authors proposed a heuristic method to handle the RSC-MDP challenge authors then conducted extensive experiments to compare the proposed method and SOTA baselines for robust RL. Experiments justified that the proposed method has better performance for the proposed spurious correlation attacks.", "questions": "See Weaknesses.", "limitations": "N.A.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is very well written, with a clear explanation of motivation and well organized problem setup and experimental design.\n2. The authors provide theoretical justification that the proposed robustness is milder and more accurate than the traditional robust MDP.\n3. The authors provide novel experimental design to assess the effect of spurious correlation in the state space of RL. The proposed baseline has the potential to be adopted by future study.", "weaknesses": "1. In the proposed driving example, the confounder issue spans over multiple steps of interactions, whereas the latent variable c introduced in RSC-MDP only affects one step, and each step is affected independently (i.e., $c_t$ is independently generated for all $t$). In reality, the brightness and traffic density do not normally get perturbed per each step of the transition. Do the proposed RSC-MDPs cover the targeted spurious correlation problem?\n\n2. Seems that the confounding within stat is not affecting the distribution of transition; that said, the observation still follows the distribution of trajectories after marginalizing the confounder $c$. To me, the spurious correlation challenge of RSC-MDPs is more of a robustness challenge than a confounding challenge to me, where the offline trajectory distribution is typically different from the interventional transition dynamics. How does the spurious correlation within each state affect the learning of transition P(s_{t+1} | s_t, a_t), hence affecting regular RL algorithms?\n\n3. Could authors provide more reasoning on the proposed heuristic perturbation in approximating the true underlying perturbation of latent variables? Intuitively it seems to align well with the motivative driving example, but it seems that the perturbation proposed in equation (6) introduces extra correlations within states. In particular, sample $s^i_k$ selected to replace $s^i_t$ is based on a rule that explicitly involves $s^j_t$ for $j\\neq i$. Do we worry that such introduced correlation ruins training for similar reasons of the targeted spurious correlation? \n\nMisc:\n\nline 81: discounted infinite-horizon -> finite horizon?\n\nline 220: $\\sigma_2 \\in (3/4, 1]$ according to the proof.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for valuable suggestions and the praise of our systematic formulation and experiments. In what follows, we provide our response to the reviewer's comments.\n\n\n### **Q1: Does the proposed RSC-MDP formulation cover the targeted spurious correlation problems in the real world and our experiments?**\nThanks for pointing out this important question! **Our RSC-MDP is a general formulation designed to cover most types of unobserved confounders.**  Note that the proposed RSC-MDP allows but not requires the confounder variable to vary across different time steps, i.e., $c_t$ can have different values for different $t$ but not forced to. So it includes the case that the confounder variable keeps the same in the entire horizon, i.e., $c_1 = c_2 = \\cdots c_T$. In addition, RSC-MDP can deal with other scenarios when the confounder variable $c_t$ changes over time, e.g., the car is driving through sudden weather changes.\n\n### **Q2: How does the spurious correlation within each state affect the learning of transition $P(s_{t+1} |s_t, a_t)$, hence affecting regular RL algorithms?**\nThanks for raising this insightful question. Recall the transition kernel for the proposed state-confounded MDP (SC-MDP) -- $s_{t+1} \\sim \\mathcal{P}\\_t(\\cdot | s_t, a_t, c_{t})$. The spurious correlation -- represented by an unknown confounder distribution $P_t^c$ (i.e., the confounder $c_t \\sim P_t^c$) will determine the expected transition kernel $\\sum_{c_t} \\mathcal{P}\\_t(\\cdot | s_t, a_t, c_{t})  P_t^c(c_t)$. So the expected transition kernel will change if $P_t^c$ changes. The reviewer is correct that if the distribution $P_t^c$ is fixed, the influence of confounder on the transition kernel will not exist after marginalizing w.r.t. $c_t$. However, in this work, we desire to learn a policy that can address possible perturbation of the confounder distribution $P_t^c$ --- robust SC-MDP. In this case, the standard RL algorithm may learn the transition kernel based on one confounder distribution $P_t^c$ (spurious correlation) and fail catastrophically when $P_t^c$ varies in the testing environment.\n\n### **Q3: More reasoning on the proposed heuristic perturbation for approximating the true underlying perturbation of the latent variables?**\nThanks for raising this valuable question. We would like to answer this question from two aspects:\n* **Why our perturbation within the state approximates changing latent variables.** We provide more explanation by showing a concrete example, which visualizes the original state trajectories and the generated trajectories by our perturbation algorithm. Please refer to the figure in **General Response (3)**. In the nominal environment, the green cube is always initialized on the left part of the table and the red cube is initialized on the right part. We assume the latent variable $z$ is discrete and has 4 values: $z=0$ (generate green-left samples), $z=1$ (generate red-right samples), $z=2$ (generate green-right samples), $z=3$ (generate red-left samples). The nominal environment only contains cases with $z=0,1$, which has a strong spurious correlation between color and position (Figure (a)). Without explicitly setting $z=2,3$ (we can't do this during training to unobserved $z$), we directly perturb the states to mimic the effect of controlling $z$, which gives us the samples in Figure (c) that contains the cases with $z=2,3$ (green-right and red-left samples).\n* **Why perturbation within state works and doesn't involve additional spurious correlation.** The correlation between two dimensions in a state could be either causation or spurious correlation. To avoid the neural network overfits to harmful spurious correlation, we randomly perturb a small portion of the data ($\\beta\\%$) by swapping some dimensions of states to regularize the model. Since this is a random perturbation, it does not introduce additional spurious correlation. As we only perturb a small portion of the data, the model can still learn meaningful features from the remaining large portion of the data.\n\n### **Q4: Other minor comments.**\nThanks for improving the writing of our work. We have revised and polished the paper accordingly to the reviewer's comments.\n"}, {"review_id": "TKaGNY7W9i", "reviewer": "Reviewer_EdL3", "summary": "The paper proposes a state-confounded (SC-)  and a robust-state confounded (RSC-) MDP formulation to account for setups where a confounder satisfying the backdoor criterion confounds the states. The RSC-MDP setup assumes that the confounder lies in an uncertainty set which is part of the MDP parameters, and the aim is to learn a value function that is the lower bound of the different value functions arising out of different instances of the RSC-MDP at different values of the confounder within its uncertainty set. They propose an empirical algorithm  to approximate the effect of perturbing the confounder by perturbing single dimensions of the states and setting them to values present in ‘otherwise nearby’ states (defined by distance between the two states in dimensions other than the dimension being edited). The paper claims that this indirectly estimates the effects of perturbing the confounder. They learn a  graph-based model that is supposed to be a causal model, to generate new transitions (given a state and action s_t, a_t, generate s_t+1, r_t) based on the perturbed states and (original) action pairs. They augment SAC by randomly replacing part of the sampled batches with the generated transitions and learn as usual, and show that this recipe is robust to the spurious correlation in the nominal envs they define. All experiments are in the state space (non-vision inputs).\n", "questions": "1. How is the experiment to answer R3 designed? It's unclear how to test the w/o P^{c} case since the confounder isn't known anyway? \n\n2. As stated previously, it would help to get more insights into how novel the generated transitions from the causal model really are.\n\n3. Have you observed or quantified the robustness-performance tradeoff mentioned in line 263, which will lead to a performance drop in some envs? I expect the perturbation procedure of states to only work for specific kinds of envs, so it would be good to show the failure cases as well.\n\n", "limitations": "Given how heuristic the technique of generating perturbed states is (swapping single dimensions between otherwise nearby states), the paper should properly discuss the many challenges of scaling this approach to a high-dimensional state space like that in visual domains.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper studies an important problem and proposes a set of  benchmarks with handcrafted spurious correlations in a CARLA environment and in a robosuite environment which could be useful for more future works to study spurious correlations in RL. \n2. The RSC-MDP formulation could also be useful, although I’m not sure how much the notation and proofs scale to an arbitrary number of confounding variables.\n", "weaknesses": "1. The paper structure is somewhat confusing to me: The MDP formulations and the empirical algorithm don’t seem to have much of a connection and seems like both were developed independently - please correct me if I’m missing something here. It would be great to further clarify exactly how the algorithm is helping to solve the robust SC-MDP.\n\n2. The notion of semantic uncertainty is not very precise or clear (I don’t think there’s any references either), and could be removed in my opinion since the RSC-MDP formulation can simply use the “uncertainty set” phrasing.\n\n3. The appendix suggests that the training environments have perfect spurious correlations: the correlation isn’t broken in even a few instances of the training environments, so for eg. there would never be a (non-generated) transition in the SAC buffer that would have the brightness value low in its state space value, in the case that the training envs are following the correlation regime where brightness and traffic are both high during the day. It’s unclear to me then, why the causal model would ever generate a state value that doesn’t also have the brightness value set to high (same logic for the perturbation procedure). Am I missing something else here? Or are the training envs set up such that an agent will train on two kinds of envs at once - night time with less traffic and day time with more traffic, and so it s possible to see samples with the brightness value set to both low and high across otherwise nearby states?\nIt would be great to clarify what exactly is the train and test env distribution and what is expected of the graphical causal model - is it to generate counterfactual transitions (counterfactual to what?) and conclusively show that it is actually doing that by visualizing or plotting some property of the samples it is generating. It's possible that I have misinterpreted something here, and I'm willing to raise my score if the authors' response clarifies some of these questions. \n\n4. I would also suggest mentioning relevant related work explicitly aimed at resolving causal confusion in online [1] and offline RL [2] (and comparing to the closest version if applicable).\n\n5. It would be great to include the training curves for the results in Table 1 as well, since that gives insight into whether other methods simply converge later to a similar highest reward, or whether they are entirely limited in their ability to converge to a similarly high reward as the best performing method. \n\n[1] Resolving Causal Confusion in Reinforcement Learning via Robust Exploration . Clare Lyle, Amy Zhang, Minqi Jiang, Joelle Pineau, Yarin Gal. ICLR Self-Supervised RL Workshop 2021.\n\n[2] Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning? Gunshi Gupta, Tim G. J. Rudner, Rowan Thomas McAllister, Adrien Gaidon, Yarin Gal, NeurIPS Offline RL Workshop 2022, CleaR 2023\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We gratefully thank the reviewer for recognizing our contributions to problem formulation and the creation of a useful benchmark! We provide our response below:\n\n### **Q1: The connection between the proposed problem formulation (robust SC-MDP) and the empirical algorithm.**\nThe empirical algorithm implicitly constructs and manipulates the unknown uncertainty set proposed in the robust SC-MDP formulation. To solve the RSC-MDP problem, we need to optimize over an unknown transition kernel uncertainty set -- structured ball around some nominal confounder value. Due to lacking the information on both the structure and nominal value, the empirical algorithm approximates the uncertainty set by constructing samples within it -- perturbing the states to mimic different confounder values around the nominal value.\n\n### **Q2: The notion of semantic uncertainty.**\nThanks for the valuable suggestion. We agree that the term \"semantic uncertainty\" may cause confusion without further explanation and references. We replace it with \"structured uncertainty\" since the uncertainty set is determined by some underlying task-specific causal structure.\n\n### **Q3: What exactly is the train and test env distribution? Show that our approach is actually generating desired samples by visualizing.**\nWe use Brightness to illustrate the train and test envs. We assume the latent confounder has 4 values: $z=0$ (generate day-heavy samples), $z=1$ (generate night-light samples), $z=2$ (generate day-light samples), and $z=3$ (generate night-heavy samples). The training environment is generated with $z=0,1$ and the testing environment is generated with $z=2,3$. Therefore, the shift between training and testing comes from different compositions of brightness and traffic. **Since we cannot explicitly set $z=2,3$ in training, our perturbation method simulates the effect of setting $z=2,3$ by perturbing the training data (with $z=0,1$).**\n\n**We visualize the generated trajectories of Lift in General Response (3) (includes figures and detailed explanation) with comparisons to original trajectories in training env, showing that our perturbation method can generate counterfactual examples to the unobserved confounder.** \n\n### **Q4: Mention the related works [1][2] explicitly.**\nThanks for providing important related works [1][2]. We add the following discussion to our related work:\n> [1] designs an exploration algorithm to conduct intervention and improve state-action coverage to avoid biased data collection. However, our work deal with a more general setting, where the testing environment could vary from the training one, leaving the exploration strategy inapplicable.\n\n> [2] proposes an uncertainty-based acquisition function to resample from the data buffer, learning more from the samples that do not have spurious correlations. This method does not explicitly handle spuriousness and could fail when spurious correlations exist in most samples.\n\n**We add [2] as a baseline and show the results in General Response (1)**. \n\n    [1] Resolving Causal Confusion in Reinforcement Learning via Robust Exploration. Lyle et.al. ICLR Workshop 2021\n    [2] Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning? Gupta et.al. CleaR 2023\n\n### **Q5: The robustness-performance tradeoff.**\nThe robustness-performance tradeoff is mainly determined by $\\beta\\%$, the ratio of the perturbed data.\nTo investigate more on this tradeoff, we add an ablation study on $\\beta\\%$ and show the results in **General Response (2)**. Two important messages:\n* **The existence of the tradeoff.** When the ratio of perturbed data $\\beta\\%$ is very small (1%), RSC-SAC achieves similar results as vanilla SAC in nominal settings and shows no robustness in shifted settings. As $\\beta\\%$ increases (considering more robustness), the performance of RSC-SAC in the nominal setting gets worse, while reversely the performance gets better in the shifted settings (more robust). However, when the ratio is too large (>80%), the performances of RSC-SAC in both settings degrade a lot, since the policy is too conservative.\n* **Our RSC-SAC maintains both robustness and good performance.** Although the tradeoff exists, the proposed RSC-SAC can perform well in both nominal and shifted settings for a large range of $\\beta\\%$ (10%-80%) -- keeping good performance in the nominal setting and achieving robustness.\n\n### **Q6: How are the experiments designed to answer R3? How to test the w/o $P^{c}$ case since the confounder is unknown?**\n* **How to design 3 ablation studies in R3.** All 3 ablation methods are modified on RSC-SAC. For \"w/o $\\textbf{G}\\_\\phi$\", we use a full graph to replace the learnable causal graph during training. For \"w/o $P_{\\theta}$\", we replace the entire causal model with a fully-connected NN. For \"w/o $P^c$\", we don't do the perturbation of state introduced in Section 4.1.\n* **The experiments of w/o $P^{c}$.** We denote the unknown confounder distribution as $P^{c}$. \"w/o $P^{c}$\" means that we don't permute the dimensions of states in Equation 6 in the data generation process. So the generated data will still be in the same distribution as the nominal training environment. We change it to \"w/o $P^{c}$ perturbation\" in Table 4 to avoid confusion.\n\n### **Q7: Discussion about scaling our approach to high-dimensional state spaces such as visual domains.**\nWe add more discussion in the conclusion section. Please check **General Response (4)**.\n\n### **Q8: Adding the training curves for the results in Table 1.**\nThe training curve is shown in the bottom row of Figure 5, where we display the testing reward on the shifted testing environment with the training step increases. After a long training process, the proposed method RSC-SAC still outperforms all baselines. To verify that all methods are converged (enter a flat area), we also plot the testing reward on the nominal environment during training in the top row of Figure 5.\n"}, {"review_id": "ycxKgrcLqn", "reviewer": "Reviewer_27Kb", "summary": "This work aims to make RL robust to distribution shifts by limiting their reliance on spurious correlations between state features.  This is achieved through a new RL algorithm designed in a variant of Robust MDPs extended to include a more structured uncertainty representation.", "questions": "How sensitive is the approach to changes in B% and K?", "limitations": "It is important to discuss somewhere how hyper-parameters (such as k and B%) were chosen, and could be chosen in a new domain.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* This paper addresses a serious underlying flaw in the standard approaches for Robust MDPs, in that the uncertainty set is usually artificial an no alined to the ways we need our policies to be robust.", "weaknesses": "The language is often confusing, for instance the term \"spurious correlations\" is often used to refer specifically to spurious correlations between state variables.  This is very confusing as it overloads the term, and  makes it hard to refer to other spurious correlations (such as a spurious correlation between the agents actions and the reward or another agents actions).  Better to say \"Spurious state correlations\".\n\nOn a similar note \"semantic uncertainty\" is a confusing term that wasn't explained until deep into the paper.  It sounds like uncertainty **about**  semantics (which is how it has been used in the literature), but it is meant to be uncertainty that is not just a norm-ball perturbation over the transition function.  A better term would be \"structured\" uncertainty.\n\n* Other  more minor confusing wordings make the paper more difficult to read than it needs to be:\n\t* I don't know what it means for portions of the state to \"not have causality\"\n\t* I don't know what it means for an uncertainty set to be \"shaped by the unobserved confounder and sequential structure of RL\"\n\t* I don't know what it means for some approach to be \"superior in breaking spurious correlations\"\n\t* There is a typo in : \"Despite various types of uncertainty have been investigated in RL\"\n\n\nIn addition it's unclear how sensitive the approach is to the degree of robustness.  It's a famous issue of RMDPs that the degree of robustness has to be carefully selected, too large and you have  an overly conservative policy, and too small you get no robustness.  Giving that part of the argument for this approach is alleviating that problem (as shown in Figure 3), a sensitivity analysis here would greatly improve the paper, as a main draw away from RMDPs is that they are very sensitive to the robustness parameter.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We gratefully thank the reviewer for valuable suggestions and the recognition of our proposed problem setting. In what follows, we provide our response to the reviewer's comments.\n\n### **Q1. Change the term \"spurious correlation\" to \"spurious state correlation\".**\nThanks for raising this point. We agree that \"spurious correlation\" is generally used between different variables of interest and may lead to confusion without further specification. So we shall explicitly introduce that this work focuses on the spurious correlation between state and change \"spurious correlation\" to \"spurious state correlation\" when it is better.\n\n\n### **Q2: Change the term \"semantic uncertainty\" to \"structured uncertainty\".**\nThanks for the valuable suggestion. We agree that the terminology \"structured uncertainty\" is better to describe our setting, where the uncertainty set is not a norm ball using some divergence function but a possible heterogeneous ball with a causal structure. \n\n\n### **Q3: Study of the uncertainty level/degree sensitivity ($\\beta\\%$ and K).**\nThanks for raising this question. Different from the traditional RMDPs that use a radius parameter $\\sigma$ to control the size of the uncertainty ball, our RSC-MDP characterizes the uncertainty set by constructing samples inside it by perturbing the confounder over state space. To control the degree of the perturbation and hence implicitly control the uncertainty size, we introduce two hyper-parameters $\\beta \\%$ and $K$. We agree with the reviewer that the sensitivity analysis of these two parameters is critical so we conduct additional experiments and show the results in **General Response (2)** (two tables). The results provide three important findings:\n* **Our RSC-SAC is not sensitive to $\\beta$.** Shown in the first table, the proposed RSC-SAC can perform well in both nominal and shifted settings --- keeping good performance in the nominal setting and achieving robustness, for a large range of $\\beta\\%$ (10%-80%). It verifies that RSC-SAC is not sensitive to hyperparameter choices.\n* **Our RSC-SAC is not sensitive to $K$.** We evaluate the proposed RSC-SAC using different $K = [32,64,\\cdots, 1024]$ and achieve similar results, shown in the second table. It shows that RSC-SAC is not sensitive to the size $K$ of candidate samples for permutation.\n* **Performance-robustness tradeoff.** From the first table, when the ratio of perturbed data $\\beta\\%$ is very small (1%), RSC-SAC almost achieves the same results as vanilla SAC in nominal settings and there is no robustness in shifted settings. As $\\beta\\%$ increases (considering more robustness), the performance of RSC-SAC in the nominal setting gradually gets worse, while reversely gets better in the shifted settings (more robust). However, when the ratio is too large (>80%), the performances of RSC-SAC in both settings degrade a lot, since the policy is too conservative so that fails in all environments.\n\n### **Q4: Other comments about writing.**\nThanks for the careful reading and valuable suggestions about the writing. We address them below:\n* \"...different portions of the state that do not have causality...\" -> \"different portions of the state that don't have correlations induced by unobserved confounder.\"\n* \"an uncertainty set shaped by the unobserved confounder and sequential structure of RL.\" -> \"an uncertainty set determined by some causal structure and unobserved confounder\"\n* \"an approach that is superiority in breaking spurious correlations...\" -> \"an approach that can achieve robust performance by avoiding learning useless spurious correlations\"\n* \"Despite various types of uncertainty have been investigated in RL,...\" -> \"Despite various types of uncertainty that have been investigated in RL,...\"\n"}]

Author rebuttal:
We thank the reviewers for their careful reading of the paper and their insightful and valuable feedback. We provide new experimental results and discussions to answer some common questions raised by reviewers.

### **(1) Add a new baseline [1], which also tackles spurious correlation in RL.**

|Env|Brightness|Behavior|Crossing|CarType|Lift|Stack|Wipe|Door|
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|[1] (nominal)|**1.07±0.10**|1.00±0.02|**1.08±0.06**|**1.00±0.02**|**0.99±0.03**|0.90±0.12|**0.93±0.20**|**0.99±0.05**|
|Ours(nominal)|0.92±0.31|**1.06±0.07**|0.96±0.03|0.96±0.03|0.96±0.05|**1.04±0.08**|0.92±0.14|0.98±0.05|
|[1] (shifted)|0.47±0.14|0.83±0.09|0.14±0.03|0.77±0.14|0.35±0.09|0.24±0.12|0.17±0.17|0.05±0.02|
|Ours(shifted)|**0.99±0.11**|**1.02±0.09**|**1.04±0.02**|**1.03±0.02**|**0.98±0.04**|**0.77±0.20**|**0.85±0.12**|**0.61±0.17**|

The results indicate that [1] has very limited robustness in the shifted testing environment compared to our method, especially in Crossing, Stack, Wipe, and Door tasks.


    [1] Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning? Gunshi Gupta, Tim G. J. Rudner, Rowan Thomas McAllister, Adrien Gaidon, Yarin Gal, NeurIPS Offline RL Workshop 2022, CleaR 2023

### **(2) Add sensitivity analysis of the ratio of perturbed data $\beta\%$ and the number of perturbation candidates $K$.**

|$\beta\%$|1%|10%|20%|30%|40%|50%|60%|70%|80%|90%|100%|
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|CarType (nominal)|0.998 ±0.018|0.989 ±0.028|0.987 ±0.020|0.979 ±0.030|0.974 ±0.032|0.984 ±0.027|0.966 ±0.029|0.965 ±0.032|0.952 ±0.035|0.914 ±0.053|0.854 ±0.096|
|CarType (shifted)|0.654 ±0.210|0.826 ±0.127|0.977 ±0.051|1.003 ±0.044|0.995 ±0.042|1.012 ±0.035|1.014 ±0.028|1.017 ±0.028|1.001 ±0.039|0.905 ±0.145|0.825 ±0.172|
|Crossing (nominal)|1.002 ±0.018|0.995 ±0.026|0.993 ±0.024|0.988 ±0.022|0.975 ±0.028|0.968 ±0.031|0.964 ±0.029|0.952 ±0.039|0.909 ±0.041|0.869 ±0.135|0.818 ±0.166|
|Crossing (shifted)|0.675 ±0.120|0.990 ±0.051|1.012 ±0.043|1.031 ±0.028|1.029 ±0.032|1.019 ±0.018|1.025 ±0.028|1.012 ±0.027|0.977 ±0.039|0.915 ±0.140|0.859 ±0.147|

|K|32|64|128|256|512|1024|
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|CarType (nominal)|0.972±0.022|0.967±0.029|0.978±0.022|0.975±0.023|0.978±0.026|0.967±0.027|
|CarType (shifted)|1.009±0.037|1.009±0.032|1.020±0.031|1.014±0.033|1.005±0.027|1.009±0.034|
|Crossing (nominal)|0.971±0.030|0.974±0.032|0.987±0.023|0.974±0.030|0.983±0.034|0.980±0.024|
|Crossing (shifted)|1.036±0.022|1.040±0.021|1.041±0.021|1.039±0.017|1.050±0.027|1.039±0.019|

The results demonstrate three important messages:
* **Our RSC-SAC is not sensitive to $\beta$.** As shown in the first table, the proposed RSC-SAC performs well in both nominal and shifted settings --- keeping good performance in the nominal setting and achieving robustness, for a large range of $\beta\%$ (10%-80%). It verifies that RSC-SAC is not sensitive to hyperparameter choices.
* **Our RSC-SAC is not sensitive to $K$.** As shown in the second table, we evaluate the proposed RSC-SAC using different $K = [32,64,\cdots, 1024]$ and achieve similar results. It shows that RSC-SAC is not sensitive to the size $K$ of candidate samples for permutation.
* **Performance-robustness tradeoff.** In the first table, when the ratio of perturbed data $\beta\%$ is very small (1%), RSC-SAC almost achieves the same results as vanilla SAC in nominal settings and there is no robustness in shifted settings. As $\beta\%$ increases (considering more robustness), the performance of RSC-SAC in the nominal setting gradually gets worse, while reversely gets better in the shifted settings (more robust). However, when the ratio is too large (>80%), the performances of RSC-SAC in both settings degrade a lot, since the policy is too conservative so that fails in all environments.


### **(3) Add a visualization of the generated trajectory in Lift environment by our perturbation algorithm.**
The figure is in the uploaded **PDF**.

In the nominal (training) environment, the green cube is always initialized on the left part of the table and the red cube is initialized on the right part. In the shifted (testing) environment, the green cube is always initialized on the right part of the table and the red cube is initialized on the left part. 

Figure (a) shows the trajectory of the state in the nominal environment. If we don't do any perturbation of the state, the generated trajectories will still have a spurious correlation (Figure (b)). However, with our perturbation within the state, we generate trajectories that break the spurious correlation and blend the color (Figure (c)). In the shifted (testing) environment, we will have the green cube on the right part and the red cube on the left part. With the counterfactual data (to unobserved confounder) generated by our algorithm (Figure (c)), we can prevent the RL model from overfitting to the spurious correlation between the color and the position of the cube.


### **(4) Discussion about the limitation and potential improvement of our method for high-dimensional tasks.**

> The current method is of great potential to be applied to more complicated problems. In particular, the proposed method requires swapping the dimensions of states to break spurious correlation, where it implicitly assumes that each dimension has semantic meaning. In more complicated problems such as high-dimensional image states, each dimension becomes a pixel without semantic meaning anymore and the dimension size may explode. To extend our method to such cases, we can leverage existing state abstraction techniques [5,6] to project images to a low-dimensional latent space, where each dimension represents a semantic feature. Then, our method can be straightforwardly applied to the low-dimensional latent space.


Meta-review: Paper proposes Robust State-Confounded MDPs as a way to mitigate against effects of correlations induced by unobserved confounders between states. Reviewers' main concerns were mainly the many related works not address or included, but the authors have since included many of these included running additional comparisons, which has largely satisfied most of the reviewers' concerns.
Decision: Accept (poster)