SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\mathcal{P}_{t,s,a,c_t}\in \mathbb{R}^{1\times S}$  such that  $\mathcal{P}_{t,s,a,c_t}(s')\coloneqq \mathcal{P}_t(s'|s,a,c_t)$  for  $s'\in S$ , and  $\widetilde{V}_t^{\star,\sigma}(s) = \sup_{\pi_t\in \Delta (\mathcal{A})}\left\{\mathbb{E}_{\pi_t}[r_t(s,a_t)] + \inf_{P_t\in \mathcal{U}^\sigma (P_t^c)}\mathbb{E}_{\pi_t}\left[\mathbb{E}_{c_t\sim P_t}\left[\mathcal{P}_{t,s,a,c_t}\widetilde{V}_{t + 1}^{\star,\sigma}(s_{t + 1})\right]\right]\right\}$ .
D.1 Architecture of the structural causal model
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 9. In normal neural networks, the input is treated as a 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, in our model, we design an encoder that is shared across all dimensions of the input. Since different dimensions could have exactly the same values, we add a learnable position
embedding to the input of the encoder. In summary, the input dimension of the encoder is  $1 + d_{pos}$ , where  $d_{pos}$  is the dimension of the position embedding.
After the encoder, we obtain a set of independent features for each dimension of the input. We now multiply the features with a learnable binary causal graph  $G$ . The element  $(i,j)$  of the graph is sampled from a Gumbel-Softmax distribution with parameter  $\phi_{i,j}$  to ensure the loss function is differentiable w.r.t  $\phi$ .
The multiplication of the causal graph and the input feature creates a linear combination of the input feature with respect to the causal graph. The obtained features are then passed through a 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.
Figure 10: Illustration of tasks in the Carla simulator.
Brightness. The nominal environments are shown in the  $1^{\text{th}}$  column of Figure 10, where the brightness and the traffic density are correlated. When the ego vehicle drives in the daytime, there are many surrounding vehicles (first row). When the ego vehicle drives in 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 are shown in the  $2^{\mathrm{nd}}$  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.e., the other vehicle always decelerates to block the ego vehicle.
Crossing. The nominal environments are shown in the  $3^{\mathrm{rd}}$  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  $4^{\text{th}}$  column of Figure 10, where the type of vehicle and the speed of the vehicle are correlated. When the vehicle is a truck, the speed is low and
when the vehicle is a motorcycle, the speed is high. In the shifted environment, the truck drives very fast and the motorcycle drives very slow.
Figure 11: Illustration of tasks in the Robosuite simulator.
Lift. The nominal environments are shown in the  $1^{\text{th}}$  column of Figure 11, 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 is 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 is red when it is in the left part.
Stack. The nominal environments are shown in the  $2^{\mathrm{nd}}$  column of Figure 11, where the position of the red cube and green plate are correlated. When the cube is in the left part of the table, the plate is 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.e., 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 in the left part.
Wipe. The nominal environments are shown in the  $3^{\mathrm{rd}}$  column of Figure 11, where the shape of the dirty region is correlated to the position of the cube. When the dirty region is 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. In 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-diagonal, the cube is on the right-hand side of the robot arm.
Door. The nominal environments are shown in the  $4^{\text{th}}$  column of Figure 11, where the height of the handle and the position of the door are correlated. When the door is closed to 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 a low position.
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 a 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
(d), we use the causal-based transition model to generate new trajectories and we find that the results not only follow the spatial distribution but also blend the color.
The results shown in Figure 12 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.
Figure 12: 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.
D.4 Computation Resources
Our algorithm is implemented on top of the Tianshou [101] package. All of our experiments are conducted on a machine with an Intel i9-9900K CPU@3.60GHz (16 core) CPU, an NVIDIA GeForce GTX 1080Ti GPU, and 64GB memory.
D.5 Hyperparameters
We summarize all hyper-parameters used in the Carla experiments (Table 5) and Robosuite experiments (Table 6).
D.6 Discovered Causal Graph in SCM
To show the performance of our learned SCM, we plot the estimated causal graphs of all experiments in Figure 13, Figure 14, Figure 15, Figure 16, and Figure 17.
Table 5: Hyper-parameters in Carla experiments
Parameters	Notation	Environment
Brightness	Behavior	Crossing	CarType
Horizon steps	T	100	100	100	100
State dimension	n	24	12	12	12
Action dimension	dA	2	2	2	2
Max training steps	1×105	1×105	5×105	5×105
Weight of ||G||p	λ	0.1	-	-	-
norm of ||G||p	p	0.1	-	-	-
Actor learning rate	3×10-4	-	-	-
Critic learning rate	1×10-3	-	-	-
Batch size	256	-	-	-
Discount factor	γ in SAC	0.99	-	-	-
Soft update weight	τ in SAC	0.005	-	-	-
Weight of entropy	α in SAC	0.1	-	-	-
Hidden layers	[256, 256, 256]	-	-	-
Returns estimation step	4	-	-	-
Buffer size	1×105	-	-	-
Steps per update	10	-	-	-
Table 6: Hyper-parameters in Robosuite experiments
Parameters	Notation	Environment
Lift	Stack	Door	Wipe
Horizon steps	T	300	300	300	500
Control frequency (Hz)	20	20	20	20
State dimension	n	50	110	22	30
Action dimension	dA	4	4	8	7
Controller type	OSC position	OSC position	Joint velocity	Joint velocity
Max training steps	1×10^6	5×10^6	1×10^6	1×10^6
Weight of ||G||p	λ	0.01	-	-	-
norm of ||G||p	p	0.1	-	-	-
Actor learning rate	3 × 10^-4	-	-	-
Critic learning rate	1 × 10^-3	-	-	-
Batch size	128	-	-	-
Discount factor	γ in SAC	0.99	-	-	-
Soft update weight	τ in SAC	0.005	-	-	-
alpha learning rate	lrα in SAC	3 × 10^-4	-	-	-
Hidden layers	[256, 256, 256]	-	-	-
Returns estimation step	4	-	-	-
Buffer size	1 × 10^6	-	-	-
Steps per update	10	-	-	-
Figure 13: Estimated Causal Graphs of four tasks in Carla.
cuiiie 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 1011 101
Figure 14: Estimated Causal Graphs of the Lift task in Robosuite.
a 10000000000000000000000000000000000000000000000000000000000000000000000000000
Figure 15: Estimated Causal Graphs of the Stack task in Robosuite.
Figure 16: Estimated Causal Graphs of the Door task in Robosuite.
Figure 17: Estimated Causal Graphs of the Wipe task in Robosuite.