SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In words, Theorem 2 reveals a fact about the proposed RSC-MDPs: RSC-MDPs could succeed in intense semantic attacks 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  $\mathcal{U}^{\sigma_2}(P^c)$ , solving RSC-MDPs with  $\pi_{\mathrm{RSC}}^{\star,\sigma_2}$  succeeds in testing while  $\pi_{\mathrm{RMDP}}^{\star,\sigma_1}$  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.
4 An Empirical Algorithm to Solve RSC-MDPs: RSC-SAC
When addressing distributionally robust problems in RMDPs, the worst case is typically defined within a prescribed uncertainty set in a clear and implementation-friendly manner, allowing for iterative 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 characterization 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 intervention or counterfactual reasoning. Therefore, we choose to approximate the causal effect of perturbing the confounder by learning from the data collected during training.
Algorithm 1: RSC-SAC Training
Input: policy π, data buffer D, transition model Pθ, ratio of modified data β
for t ∈ [1,T] do
Sample action att ~ π(·|st)
(st+1,rt) ← Env(st, at)
Add buffer D = D ∪ {st, at, st+1, rt}
for sample batch B ∈ D do
Randomly select β% data in B
Modify st in selected data with (7)
(ˆt+1,ˆt) ~ Pθ(st, at, Gφ)
Replace data with (st, at,ˆt+1,ˆt)
L = ||st+1 -ˆt+1||2 + ||rt -ˆrt||2
Update θ and φ with L + λ||G||p
Update π with SAC algorithm
In this section, we propose an intuitive yet effective empirical approach named RSC-SAC for solving RSC-MDPs, which is outlined in Algorithm 1. We first estimate the effect of perturbing the distribution  $P^c$  of the confounder to generate new states (Section 4.1). Then, we learn the structural causal model  $\mathcal{P}_t^i$  to predict rewards and the next states given the perturbed states (Section 4.2). By combining these two components, we construct a data generator capable of simulating novel transitions  $(s_t, a_t, r_t, s_{t+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.
4.1 Distribution of confounder
As we have no prior knowledge about the confounders, we choose to approximate the effect of perturbing them without explicitly estimating the distribution  $P^c$ . We first randomly select a dimension  $i$  from the state  $s_t$  to apply perturbation and then assign the dimension  $i$  of  $s_t$  with a heuristic rule. We select the value from another sample  $s_k$  that has the most different value from  $s_t$  in dimension  $i$  and the most similar value to  $s_t$  in the remaining dimensions. Formally, this process solves the following optimization problem to select sample  $k$  from a batch of  $K$  samples:
$$
s _ {t} ^ {i} \leftarrow s _ {k} ^ {i}, k = \arg \max  \frac {\| s _ {t} ^ {i} - s _ {k} ^ {i} \| _ {2} ^ {2}}{\sum_ {\neg i} \| s _ {t} ^ {\neg i} - s _ {k} ^ {\neg i} \| _ {2} ^ {2}}, k \in \{1, \dots , K \} \tag {7}
$$
where  $s_t^i$  and  $s_t^{-i}$  means dimension  $i$  of  $s_t$  and other dimensions of  $s_t$  except for  $i$ , respectively. Intuitively, permuting the dimension of two samples breaks the spurious correlation and remains the
Figure 4: Two tasks used in experiments. Door is a composition task implemented in Robosuite with a 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 [21] is a 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  $s_t$ , we then learn an 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 to unseen state-action pairs. Specifically, we simultaneously learn the model parameter and discover the underlying causal graph in a fully differentiable way with  $(\hat{s}_{t+1}, \hat{r}_t) \sim P_\theta(s_t, a_t, \mathcal{G}_\phi)$ , where  $\theta$  is the parameter of the neural network of the dynamic model and  $\phi \in \mathbb{R}^{(n + d_A) \times (n + 1)}$  is the parameter to represent causal graph  $\mathcal{G}$  between  $\{s_t, a_t\}$  and  $\{s_{t+1}, r_t\}$ . This graph is represented by a binary adjacency matrix  $G$ , where  $1/0$  means the existence/absence of an edge.  $P_\theta$  has an encoder-decoder structure with matrix  $G$  as an intermediate linear transformation. The encoder takes state and action in and outputs features  $f_e \in \mathbb{R}^{(n + d_A) \times d_f}$  for each dimension, where  $d_f$  is the dimension of the feature. Then, the causal graph is multiplied to generate the feature for the decoder  $f_d = f_e^T G \in \mathbb{R}^{d_f \times (n + 1)}$ . The decoder takes in  $f_d$  and outputs the next state and reward. The detailed architecture of this causal transition model can be found in Appendix D.1.
The objective for training this model consists of two parts, one is the supervision signal from collected data  $\| s_{t + 1} - \hat{s}_{t + 1}\| _2^2 +\| r_t - \hat{r}_t\| _2^2$ , and the other is a penalty term  $\lambda \| G\| _p$  with weight  $\lambda$  to encourage the sparsity of the matrix  $G$ . The penalty is important to break the spurious correlation between dimensions of state since it forces the model to eliminate unnecessary inputs for prediction.
In this section, we first provide a benchmark consisting of eight environments with spurious correlations, 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 a benchmark consisting of eight novel tasks in self-driving and manipulation domains using the Carla [22] and Robosuite [23] platforms (shown in Figure 4). Tasks are designed to include spurious correlations in 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.2.
- Distraction correlation is between task-relevant and task-irrelevant portions of the state. The task-irrelevant part could distract the policy model from learning important features and lead to a performance drop. A typical method to avoid distraction is 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 in compositional generalization, where states are re-composed to form novel tasks during testing. Typical examples are multi-task RL [26, 27] and hierarchical RL [28, 29]. We design four tasks with this category of correlation, i.e., Stack, Door, Behavior, and Crossing.
Figure 5: The first row shows the testing reward on the nominal environments, while the second row shows the testing reward on 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.13±0.03	0.81±0.13	0.63±0.14	0.58±0.13	0.26±0.12	0.16±0.20	0.08±0.07
RMDP-G	0.55±0.15	0.16±0.04	0.47±0.13	0.53±0.16	0.31±0.08	0.33±0.15	0.06±0.17	0.07±0.03
RMDP-U	0.54±0.19	0.13±0.05	0.60±0.15	0.39±0.13	0.51±0.17	0.23±0.11	0.06±0.17	0.10±0.13
MoCoDA	0.50±0.14	0.16±0.05	0.22±0.14	0.23±0.12	0.46±0.14	0.29±0.11	0.01±0.24	0.09±0.14
ATLA	0.48±0.11	0.14±0.03	0.61±0.14	0.52±0.14	0.61±0.18	0.21±0.12	0.29±0.18	0.28±0.19
DBC	0.52±0.18	0.16±0.03	0.68±0.12	0.45±0.10	0.12±0.02	0.03±0.02	0.19±0.35	0.01±0.01
Active	0.47±0.14	0.14±0.03	0.83±0.09	0.77±0.14	0.35±0.09	0.24±0.12	0.17±0.17	0.05±0.02
RSC-SAC	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
Robustness in RL has been explored in terms of diverse uncertainty set over state, action, or transition kernels. Regarding this, we use a non-robust RL and four representative algorithms of robust RL as baselines, all of which are implemented on top of the SAC [20] algorithm. Non-robust RL (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], a 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 [31] 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.
To comprehensively evaluate the performance of the proposed method RSC-SAC, we conduct experiments to answer the following questions: Q1. Can RSC-SAC eliminate the harmful effect of spurious correlation in learned policy? Q2. Does the robustness of 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 is critical in our empirical algorithm? Q5. Is 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.
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	1.00±0.09	1.00±0.08	1.00±0.02	1.00±0.03	1.00±0.03	1.00±0.09	1.00±0.12	1.00±0.03
RMDP-G	1.04±0.09	1.00±0.11	0.78±0.05	0.79±0.05	0.92±0.07	0.86±0.14	0.99±0.13	0.99±0.06
RMDP-U	1.02±0.09	1.04±0.07	0.90±0.03	0.88±0.03	0.97±0.05	0.92±0.12	0.97±0.14	0.88±0.31
MoCoDA	0.65±0.17	0.78±0.15	0.57±0.07	0.55±0.13	0.79±0.11	0.72±0.08	0.69±0.13	0.41±0.22
ATLA	0.99±0.11	0.98±0.11	0.89±0.05	0.88±0.04	0.94±0.08	0.88±0.10	0.96±0.12	0.97±0.05
DBC	0.75±0.12	0.78±0.10	0.85±0.08	0.86±0.06	0.27±0.04	0.12±0.08	0.31±0.21	0.01±0.01
Active	1.02±0.10	1.08±0.06	1.00±0.02	1.00±0.02	0.99±0.03	0.90±0.12	0.93±0.20	0.99±0.05
RSC-SAC	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
R1. RSC-SAC is robust against spurious correlation. The testing results of our proposed method with comparisons to the baselines are presented in Table 1, 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 in 5 out of 8 tasks. An interesting and even surprising finding, as shown in Table 1, is that although RMDP-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 is outside of the prescribed uncertainty set considered in the robust algorithms.
R2. Sparsity of the model is only one reason for the robustness of RSC-SAC. As existing 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 (SACSparse): we add an additional penalty  $\alpha |W|_{1}$  during the optimization, where  $W$  is the parameter of the first linear layer of the policy and value networks and  $\alpha$  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 distraction spurious correlation, which is consistent with the findings in [34]. (2) The sparsity does
Figure 6: Comparison between SAC-Sparse and our method.  $\alpha$  is the weight of regularization.
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 [21] finds out that there usually exists a trade-off between the performance in 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 2, 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 assess the impact of each module in our algorithm, we conduct three additional ablation studies (in Table 3), where we remove the causal graph  $G_{\phi}$ , the transition model  $P_{\theta}$ , and the distribution of the confounder  $P^{c}$  respectively. The results demonstrate
Table 3: Influence of modules
Method	Lift	Behavior	Crossing
w/o Gφ	0.79±0.15	0.51±0.24	0.87±0.10
w/o Pθ	0.75±0.13	0.41±0.28	0.89±0.08
w/o Pc	0.90±0.09	0.66±0.21	0.96±0.04
Full model	0.98±0.04	1.02±0.09	1.04±0.02
that the learnable causal graph  $G_{\phi}$  is critical for the performance that enhances the prediction of
the next state and reward, thereby facilitating the generation of high-quality next states with current perturbed states. The transition model without  $G_{\phi}$  may still retain numerous spurious correlations, resulting in a performance drop similar to the one without  $P_{\theta}$ , which does not alter the next state and reward. In the third row of Table 3, the performance drop indicates that the confounder  $P^{c}$  also plays a crucial role in preserving semantic meaning and avoiding policy training distractions.
R5. RSC-SAC is also robust to random perturbation. The final investigation aims to assess the generalizability of our method to 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 random noise under the Gaussian distribution with two varying scales in the Lift environment. In Table 4,
Table 4: Random purterbuation
Method	Lift-0	Lift-0.01	Lift-0.1
SAC	1.00±0.03	0.77±0.13	0.46±0.23
RMDP-0.01	0.97±0.05	0.96±0.06	0.51±0.21
RMDP-0.1	0.85±0.12	0.82±0.09	0.39±0.15
RSC-SAC	0.96±0.05	0.94±0.06	0.44±0.18
Lift-0 indicates the nominal training environment, while Lift-0.01 and Lift-0.1 represent the environments 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.
R6. RSC-SAC keeps good performance and robustness for a wide range of  $\beta$ . As shown in Figure 7, 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-
Figure 7: Performance-robustness tradeoff with different augmentation ratio  $\beta$ .
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 a lot, since the policy is too conservative so that fails in all environments.
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 as 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 [35] and state abstraction techniques [36].
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-2404, NSF CCF-2106778, DMS-2134080, 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.
[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 2019 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 of the IEEE/CVF International Conference on 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 arXiv:1909.03683, 2019.
[91] Divyansh Kaushik, Eduard Hovy, and Zachary C Lipton. Learning the difference that makes a difference with counterfactually-augmented data. arXiv preprint arXiv:1909.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-19352, 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.02112, 2021.
[96] Misgina Tsighe Hagos, Kathleen M Curran, and Brian Mac Namee. Identifying spurious correlations and correcting them with an explanation-based learning. arXiv preprint arXiv:2211.08285, 2022.
[97] Abubakar Abid, Mert Yuksekgonul, and James Zou. Meaningfully debugging model mistakes 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.15769, 2022.
[99] Mohammad Taha Bahadori and David E Heckerman. Debiasing concept-based explanations with causal analysis. arXiv preprint arXiv:2007.11500, 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 Duburcq, 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-6, 2022.
A Broader Impact 17
B Other Related Works 17
B.1 Related RL formulations 17
B.2 Related literature of robustness in RL 18
B.3 Related literature of spurious correlation in RL 18
C Theoretical Analyses 19
C.1 Proof of Theorem 1 19
C.2 Proof of Theorem 2 21
C.3 Proof of auxiliary results 24
DExperiment Details 27
D.1 Architecture of 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
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 human-readable private information about the environment, which could raise privacy issues. To mitigate this potential negative societal impact, the causal transition model needs to 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 to ours: (1) promoting robustness in RL; (2) concerning the spurious correlation issues in RL.
B.1 Related RL formulations