SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} I \left(\mathbf {s} _ {1: T}; i _ {1: T} \mid a _ {1: T}\right) \\ = \mathrm {E} _ {q (o _ {1: T}, r _ {1: T}, \mathbf {s} _ {\mathbf {t}: \mathbf {T}}, a _ {1: T}, i _ {1: T})} \left(\sum_ {t} \ln q (\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}, i _ {t}) - \ln p (\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1})\right) \\ = \mathrm {E} \left(\sum_ {t} \ln q _ {\phi} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}, o _ {t}\right) - \ln p \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right)\right) \tag {30} \\ \leq \mathrm {E} \left(\sum_ {t} \ln q _ {\phi} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}, o _ {t}\right) - \ln p _ {\gamma} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right)\right) \\ = \mathrm {E} \left(\sum_ {t} \mathrm {K L} \left(q _ {\phi} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}, o _ {t}\right) \| p _ {\gamma} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right)\right)\right). \\ \end{array}
$$
According to equation 26, we have  $p_{\gamma} = p_{\gamma_1} \cdot p_{\gamma_2} \cdot p_{\gamma_3} \cdot p_{\gamma_4}$  and  $q_{\phi} = q_{\phi_1} \cdot q_{\phi_2} \cdot q_{\phi_3} \cdot q_{\phi_4}$ .
$$
\begin{array}{l} \mathrm {K L} \left(q _ {\phi} \| p _ {\gamma}\right) = \mathrm {K L} \left(q _ {\phi_ {1}} \cdot q _ {\phi_ {2}} \cdot q _ {\phi_ {3}} \cdot q _ {\phi_ {4}} \| p _ {\gamma_ {1}} \cdot p _ {\gamma_ {2}} \cdot p _ {\gamma_ {3}} \cdot p _ {\gamma_ {4}}\right) \\ = E _ {q _ {\phi}} \left(\ln \frac {q _ {\phi_ {1}}}{p _ {\gamma_ {1}}} + \ln \frac {q _ {\phi_ {2}}}{p _ {\gamma_ {2}}} + \ln \frac {q _ {\phi_ {3}}}{p _ {\gamma_ {3}}} + \ln \frac {q _ {\phi_ {4}}}{p _ {\gamma_ {4}}}\right) \tag {31} \\ = \mathrm {K L} \left(q _ {\phi_ {1}} \| p _ {\gamma_ {1}}\right) + \mathrm {K L} \left(q _ {\phi_ {2}} \| p _ {\gamma_ {2}}\right) + \mathrm {K L} \left(q _ {\phi_ {3}} \| p _ {\gamma_ {3}}\right) + \mathrm {K L} \left(q _ {\phi_ {4}} \| p _ {\gamma_ {4}}\right) \\ \end{array}
$$
We introduce additional hyperparameters to regulate the amount of information contained within each category of variables:
$$
\mathcal {J} _ {\mathrm {D}} ^ {t} = - \beta_ {1} \cdot \operatorname {K L} \left(q _ {\phi_ {1}} \| p _ {\gamma_ {1}}\right) - \beta_ {2} \cdot \operatorname {K L} \left(q _ {\phi_ {2}} \| p _ {\gamma_ {2}}\right) - \beta_ {3} \cdot \operatorname {K L} \left(q _ {\phi_ {3}} \| p _ {\gamma_ {3}}\right) - \beta_ {4} \cdot \operatorname {K L} \left(q _ {\phi_ {4}} \| p _ {\gamma_ {4}}\right). \tag {32}
$$
Additionally, we introduce two supplementary objectives to explicitly capture the distinctive characteristics of the four distinct representation categories. Specifically, we characterize the reward-relevant representations by measuring the dependence between  $s_t^r$  and  $R_t$ , given  $a_{t-1:t}$  and  $s_{t-1}^r$ , that is  $I(s_t^r, R_t \mid a_{t-1:t}, s_{t-1}^r)$  (see Figure 15(a). Note that if the action  $a_t$  is not dependent on  $s_t^r$ , such as when it is randomly chosed,  $a_t$  does not need to be conditioned. In this case, the mutual information turns into  $I(s_t^r, R_t \mid a_{t-1}, s_{t-1}^r)$ . To ensure that  $s_t^r$  are minimally sufficient for policy training, we maximize  $I(s_t^r, R_t \mid a_{t-1:t}, s_{t-1}^r)$  while minimizing  $I(s_t^r, R_t \mid a_{t-1:t}, s_{t-1}^r)$  to discourage the inclusion of redundant information in  $s_t^r$  concerning the rewards:
$$
I \left(s _ {t} ^ {r}; R _ {t} \mid a _ {t - 1: t}, s _ {t - 1} ^ {r}\right) - I \left(s _ {t} ^ {\bar {r}}; R _ {t} \mid a _ {t - 1: t}, s _ {t - 1} ^ {r}\right). \tag {33}
$$
The conditional mutual information can be expressed as the disparity between two mutual information values.
$$
\begin{array}{l} I \left(s _ {t} ^ {r}; R _ {t} \mid a _ {t - 1: t}, s _ {t - 1} ^ {r}\right) = I \left(R _ {t}; s _ {t} ^ {r}, a _ {t - 1: t}, s _ {t - 1} ^ {r}\right) - I \left(R _ {t}; a _ {t - 1: t}, s _ {t - 1} ^ {r}\right), \end{array} \tag {34}
$$
$$
I \left(s _ {t} ^ {\bar {r}}; R _ {t} \mid a _ {t - 1: t}, s _ {t - 1} ^ {r}\right) = I \left(R _ {t}; s _ {t} ^ {\bar {r}}, a _ {t - 1: t}, s _ {t - 1} ^ {r}\right) - I \left(R _ {t}; a _ {t - 1: t}, s _ {t - 1} ^ {r}\right).
$$
Combining the above two equations, we eliminated the identical terms, ultimately yielding the following formula
$$
I \left(R _ {t}; s _ {t} ^ {r}, a _ {t - 1: t}, s _ {t - 1} ^ {r}\right) - I \left(R _ {t}; s _ {t} ^ {\bar {r}}, a _ {t - 1: t}, s _ {t - 1} ^ {r}\right). \tag {35}
$$
We use the Donsker-Varadhan representation to express mutual information as a supremum over functions.
$$
\begin{array}{l} I (X; Y) = D _ {K L} (p (x, y) \parallel p (x) p (y)) \\ = \sup  _ {T \in \mathcal {T}} \mathbb {E} _ {p (x, y)} [ T (x, y) ] - \log \mathbb {E} _ {p (x) p (y)} [ e ^ {T (x, y)} ]. \tag {36} \\ \end{array}
$$
We employ mutual information neural estimation [27] to approximate the mutual information value. We represent the function  $T$  using a neural network that accepts variables  $(x,y)$  as inputs and is parameterized by  $\alpha$ . The neural network is optimized through stochastic gradient ascent to find the supremum. Substituting  $x$  and  $y$  with variables defined in Equation 35, our objective is reformulated as follows:
$$
\mathcal {J} _ {\mathrm {R S}} ^ {t} = \lambda_ {1} \cdot \left\{I _ {\alpha_ {1}} \left(R _ {t}; s _ {t} ^ {r}, a _ {t - 1: t}, \mathbf {s g} \left(s _ {t - 1} ^ {r}\right)\right) - I _ {\alpha_ {2}} \left(R _ {t}; s _ {t} ^ {\bar {r}}, a _ {t - 1: t}, \mathbf {s g} \left(s _ {t - 1} ^ {r}\right)\right) \right\}. \tag {37}
$$
To incorporate the conditions from the original objective, we apply the stop-gradient operation to the variable  $s_{t-1}^r$ . Similarly, to ensure that the representations  $s_t^a$  are directly controllable by actions, while  $s_t^{\bar{a}}$  are not, we maximize the following objective:
$$
I \left(s _ {t} ^ {a}; a _ {t - 1} \mid \mathbf {s} _ {\mathbf {t} - \mathbf {1}}\right) - I \left(s _ {t} ^ {\bar {a}}, a _ {t - 1} \mid \mathbf {s} _ {\mathbf {t} - \mathbf {1}}\right), \tag {38}
$$
By splitting the conditional mutual information and eliminating identical terms, we obtain the following objective function:
$$
\mathcal {J} _ {\mathrm {A S}} ^ {t} = \lambda_ {2} \cdot \left\{I _ {\alpha_ {3}} \left(a _ {t - 1}; s _ {t} ^ {a}, \mathbf {s g} \left(\mathbf {s} _ {\mathbf {t} - 1}\right)\right) - I _ {\alpha_ {4}} \left(a _ {t - 1}; s _ {t} ^ {\tilde {a}}, \mathbf {s g} \left(\mathbf {s} _ {\mathbf {t} - 1}\right)\right) \right\}. \tag {39}
$$
where  $\alpha_{1},\alpha_{2},\alpha_{3},\alpha_{4}$  can be obtained by maximizing Equation 36. Intuitively, these two objective functions ensure that  $s_t^r$  is predictive of the reward, while  $s_t^{\bar{r}}$  is not; similarly,  $s_t^a$  can be predicted by the action, whereas  $s_t^{\bar{a}}$  cannot.
Combine the equation 29, equation 32, equation 37 and equation 39, the total objective function is:
$$
\begin{array}{l} \mathcal {J} _ {\text {T O T A L}} = \max  _ {\phi , \theta , \gamma , \alpha_ {1}, \alpha_ {3}} \min  _ {\alpha_ {2} \alpha_ {4}} \mathrm {E} _ {q _ {\phi}} \left(\sum_ {t} \left(\mathcal {J} _ {\mathrm {O}} ^ {t} + \mathcal {J} _ {\mathrm {R}} ^ {t} + \mathcal {J} _ {\mathrm {D}} ^ {t} + \mathcal {J} _ {\mathrm {R S}} ^ {t} + \mathcal {J} _ {\mathrm {A S}} ^ {t}\right)\right) + \text {c o n s t} \\ = \max  _ {\phi , \theta , \gamma , \alpha_ {1}, \alpha_ {3}} \min  _ {\alpha_ {2}, \alpha_ {4}} \mathbb {E} _ {q _ {\phi}} \left\{\log p _ {\theta} \left(o _ {t} \mid \mathbf {s} _ {\mathbf {t}}\right) + \log p _ {\theta} \left(r _ {t} \mid s _ {t} ^ {r}\right) \right. \tag {40} \\ - \sum_ {i = 1} ^ {4} \beta_ {i} \cdot \mathrm {K L} \left(q _ {\phi_ {i}} \| p _ {\gamma_ {i}}\right) + \lambda_ {1} \cdot \left(I _ {\alpha_ {1}} - I _ {\alpha_ {2}}\right) + \lambda_ {2} \cdot \left(I _ {\alpha_ {3}} - I _ {\alpha_ {4}}\right) \} + \text {c o n s t}. \\ \end{array}
$$
The expectation is computed over the dataset and the representation model. Throughout the model learning process, the objectives for estimating mutual information and learning the world model are alternately optimized.
B.1 Discussions
In this subsection, we examine the mutual information constraints in equation 37 and equation 39 and their relationship with other objectives. Our findings reveal that while other objectives partially fulfill the desired functionality of the mutual information constraints, incorporating both mutual information objectives is essential for certain environments.
The objective functions can be summarized as follows:
$$
\begin{array}{l} \mathcal {J} _ {\mathrm {O}} ^ {t} = \ln p _ {\theta} \left(o _ {t} \mid \mathbf {s} _ {\mathbf {t}}\right), \quad \mathcal {J} _ {\mathrm {R}} ^ {t} = \ln p _ {\theta} \left(r _ {t} \mid s _ {t} ^ {r}\right), \quad \mathcal {J} _ {\mathrm {D}} ^ {t} = - \mathrm {K L} \left(q _ {\phi} \parallel p _ {\gamma}\right), \\ \mathcal {J} _ {\mathrm {R S}} ^ {t} = \lambda_ {1} \cdot \left\{I _ {\alpha_ {1}} \left(R _ {t}; s _ {t} ^ {r}, a _ {t - 1: t}, \mathbf {s g} \left(s _ {t - 1} ^ {r}\right)\right) - I _ {\alpha_ {2}} \left(R _ {t}; s _ {t} ^ {\bar {r}}, a _ {t - 1: t}, \mathbf {s g} \left(s _ {t - 1} ^ {r}\right)\right) \right\}, \tag {41} \\ \mathcal {J} _ {\mathrm {A S}} ^ {t} = \lambda_ {2} \cdot \left\{I _ {\alpha_ {3}} \left(a _ {t - 1}; s _ {t} ^ {a}, \mathbf {s g} \left(\mathbf {s} _ {\mathbf {t} - 1}\right)\right) - I _ {\alpha_ {4}} \left(a _ {t - 1}; s _ {t} ^ {\bar {a}}, \mathbf {s g} \left(\mathbf {s} _ {\mathbf {t} - 1}\right)\right) \right\}, \\ \end{array}
$$
and the KL divergence term can be further decomposed into 4 components:
$$
\begin{array}{l} \mathcal {J} _ {D _ {1}} ^ {t} = - \beta_ {1} \cdot \mathrm {K L} \left(q _ {\phi_ {1}} \left(s _ {t} ^ {a r} \mid o _ {t}, s _ {t - 1} ^ {r}, a _ {t - 1}\right) \| p _ {\gamma_ {1}} \left(s _ {t} ^ {a r} \mid s _ {t - 1} ^ {r}, a _ {t - 1}\right)\right) \\ \mathcal {J} _ {D _ {2}} ^ {t} = - \beta_ {2} \cdot \mathrm {K L} \left(q _ {\phi_ {2}} \left(s _ {t} ^ {\bar {a} r} \mid o _ {t}, s _ {t - 1} ^ {r}\right) \| p _ {\gamma_ {2}} \left(s _ {t} ^ {\bar {a} r} \mid s _ {t - 1} ^ {r}\right)\right) \tag {42} \\ \mathcal {J} _ {D _ {3}} ^ {t} = - \beta_ {3} \cdot \mathrm {K L} \left(q _ {\phi_ {3}} \left(s _ {t} ^ {a \bar {r}} \mid o _ {t}, \mathbf {s} _ {\mathbf {t} - \mathbf {1}}, a _ {t - 1}\right) \| p _ {\gamma_ {3}} \left(s _ {t} ^ {a \bar {r}} \mid \mathbf {s} _ {\mathbf {t} - \mathbf {1}}, a _ {t - 1}\right)\right) \\ \mathcal {J} _ {D _ {4}} ^ {t} = - \beta_ {4} \cdot \mathrm {K L} (q _ {\phi_ {4}} \left(s _ {t} ^ {\bar {a} \bar {r}} \mid o _ {t}, \mathbf {s _ {t - 1}}\right) \parallel p _ {\gamma_ {4}} \left(s _ {t} ^ {\bar {a} \bar {r}} \mid \mathbf {s _ {t - 1}}\right)). \\ \end{array}
$$
Specifically, maximizing  $I_{\alpha_1}$  in  $\mathcal{J}_{\mathrm{RS}}^t$  enhances the predictability of  $R_t$  based on the current state  $s_{t-1}^r$  conditioning on  $(s_{t-1}^r, a_{t-1:t})$ . However, notice that this objective can be partially accomplished by optimizing  $\mathcal{J}_{\mathrm{R}}^t$ . When learning the world model, both the transition function and the reward function are trained: the reward function predicts the current reward  $r_t$  using  $s_t^r$ , while the transition model predicts the next state. These combined predictions contribute to the overall prediction of  $R_t$ .
Minimizing  $I_{\alpha_2}$  in  $\mathcal{J}_{\mathrm{RS}}^t$  eliminates extraneous reward-related information present in  $s_t^{\bar{r}}$ . According to our formulation,  $s_t^{\bar{r}}$  can still be predictive of  $R_t$  as long as it does not introduce additional predictability beyond what is already captured by  $(s_{t - 1}^r,a_{t - 1:t})$ . This is because we only assume that  $s_t^{\bar{r}}$  is conditionally independent from  $R_t$  when conditioning on  $(s_{t - 1}^r,a_{t - 1:t})$ . If we don't condition on  $(s_{t - 1}^r,a_{t - 1:t})$ , it introduces  $s_{t - 1}^{r}$  as the confounding factor between  $s_t^{\bar{r}}$  and  $R_t$ , establishing association between  $s_t^{\bar{r}}$  and  $R_t$  (refer to Figure 15(a)). Note that the KL divergence constraints govern the information amount within each state category. By amplifying the weight of the KL constraints on  $s_t^{\bar{r}}$ , the value of  $I_{\alpha_2}$  can indirectly be diminished.
By maximizing  $I_{\alpha_3}$  and minimizing  $I_{\alpha_4}$  in  $\mathcal{J}_{\mathrm{AS}}^t$ , we ensure that  $s_t^a$  can be predicted based on  $(a_{t - 1},\mathbf{s}_{t - 1})$  while  $s_t^{\bar{a}}$  cannot. The KL constraints on  $s_t^{ar}$  and  $s_t^{a\bar{r}}$  incorporate the action  $a_{t - 1}$  into the prior and posterior, implicitly requiring that  $s_t^a$  should be predictable given  $a_{t - 1}$ . Conversely, the KL constraints on  $s_t^{\bar{a} r}$  and  $s_t^{\bar{a}\bar{r}}$  do not include the action  $a_{t - 1}$  in the prior and posterior, implicitly requiring that  $s_t^a$  should not be predictable based on  $a_{t - 1}$ . However, relying solely on indirect constraints can sometimes be ineffective, as it may lead to entangled representations that negatively impact policy performance (see Figure 13).
Ablation of the mutual information constraints. The inclusion of both  $\mathcal{I}_{\mathrm{RS}}^t$  and  $\mathcal{I}_{\mathrm{AS}}^t$  is essential in certain environments to promote disentanglement and enhance policy performance, despite sharing
Figure 7: Ablation of the mutual information constraints in the Reacher environment with video background and jittering camera. The dashed brown line illustrates the policy performance of Denoised MDP after 1 million environment steps.
some common objectives. We have observed improved training stability in the variant of Robodesk environment (see Figure 12) and significant performance gains in the Reacher environment with video background and camera jittering (see Figure 7). When two mutual information objectives are removed, we notice that entangled representations emerged in these environments, as depicted in Figure 13. We assign values of 0.1 to  $\lambda_{1}$  and  $\lambda_{2}$  in the environment of modified Cartpole, variant of Robodesk and Reacher with video background and jittering camera. Empirically, a value of 0.1 has been found to be preferable for both  $\lambda_{1}$  and  $\lambda_{2}$ . Using a higher value for regularization might negatively impact the learning of representation and transition model. In other DMC environments, the ELBO loss alone has proven effective due to the inherent structure of our disentangled latent dynamics. The choice of hyperparameters  $(\beta_{1},\beta_{2},\beta_{3},\beta_{4})$  depends on the specific goals of representation learning and the extent of noise interference in the task. If the objective is to accurately recover the true latent variables for understanding the environment, it is often effective to assign equal weights to the four KL divergence terms (for experiments on synthetic data and modified cartpole). When the aim is to enhance policy training stability by mitigating noise, it is recommended to set the values of  $\beta_{1}$  and  $\beta_{2}$  higher than  $\beta_{3}$  and  $\beta_{4}$  (for experiments on variants of Robodesk and DMC). Moreover, in environments with higher levels of noise, it is advisable to increase the discrepancy in values between the hyperparameters.
Algorithm 1: IFactor
Input: Representation model:  $q_{\phi}(s_t \mid s_{t-1}, a_{t-1}, o_t) = q_{\phi_1} \cdot q_{\phi_2} \cdot q_{\phi_3} \cdot q_{\phi_4}$ ; Transition model:  $p_{\gamma}(s_t \mid s_{t-1}, a_{t-1}) = p_{\gamma_1} \cdot p_{\gamma_2} \cdot p_{\gamma_3} \cdot p_{\gamma_4}$ ; Observation Model:  $p_{\theta}(o_t \mid s_t)$ ; Reward model:  $p_{\theta}(r_t \mid s_t^r)$ ; Policy Model:  $\pi_{\psi}(a_t \mid s_t^r)$ ; Value model:  $v_{\psi}(s_t^r)$ ; Mutual information estimator:  $I_{\alpha_1}, I_{\alpha_2}, I_{\alpha_3}, I_{\alpha_4}$ ; Policy optimization algorithm Pr-Opt, which is in default the same as that used in Dreamer.
Output: Transition Model  $p_{\gamma}$ . Representation Model  $q_{\phi}$ . Policy Model  $\pi_{\psi}$
1: while training do
2: // Exploration
3: Collect new trajectories  $D' = \{(o_t, a_t, r_t)\}_t$  with  $\pi$  acting on  $q_{\phi}$  encoded outputs
4: Add experience to the replay buffer  $D = D \cup D'$
5: // Model learning
6: Sample a batch of data sequences  $\{(o_t, a_t, r_t)\}_t$  from the reply buffer  $D$
7: Obtain  $s_t$  by the representation model and estimate mutual information terms  $I_{\alpha_1}, I_{\alpha_2}, I_{\alpha_3}, I_{\alpha_4}$
8: for  $i = 1$  to  $n$  do
9: Train mutual information neural estimators by maximizing  $I_{\alpha_1}, I_{\alpha_2}, I_{\alpha_3}, I_{\alpha_4}$
10: end for
11: Freeze the parameters in  $I_{\alpha_1}, I_{\alpha_2}, I_{\alpha_3}, I_{\alpha_4}$  and Train  $q_{\phi}, p_{\gamma}$  and  $p_{\theta}$  with Equation 11.
12: // Policy learning by dreaming
13: Imagine trajectories of  $s_t^r$  using the learned world model.
14: Train  $\pi_{\psi}$  and  $\nu_{\psi}$  by running P1-OPT.
15: end while
D Environment Descriptions
D.1 Synthetic Data
For the sake of simplicity, we consider one lag for the latent processes in Section 4. Our identifiability proof can actually be applied for arbitrary lags directly because the identifiability does not rely on the number of previous states. We extend the latent dynamics of the synthetic environment to incorporate a general time-delayed causal effect with  $\tau \geq 1$  in the synthetic environment. When  $\tau = 1$ , it reduces to a common MDP. The ground-truth generative model of the environment is as follows:
$$
\left\{ \begin{array}{l l} \text {O b s e r v a t i o n M o d e l :} & p _ {\theta} \left(o _ {t} \mid \mathbf {s} _ {\mathbf {t}}\right) \\ \text {R e w a r d M o d e l :} & p _ {\theta} \left(r _ {t} \mid s _ {t} ^ {r}\right) \\ \text {T r a n s i t i o n M o d e l :} & p _ {\gamma} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - \tau : \mathbf {t} - 1}, a _ {t - \tau : t - 1}\right) \end{array} \quad \text {T r a n s i t i o n}: \left\{ \begin{array}{l} p _ {\gamma_ {1}} \left(s _ {t} ^ {a r} \mid s _ {t - \tau : t - 1} ^ {r}, a _ {t - \tau : t - 1}\right) \\ p _ {\gamma_ {2}} \left(s _ {t} ^ {\bar {a} r} \mid s _ {t - \tau : t - 1} ^ {r}\right) \\ p _ {\gamma_ {3}} \left(s _ {t} ^ {\bar {a} r} \mid \mathbf {s} _ {\mathbf {t} - \tau : \mathbf {t} - 1}, a _ {t - \tau : t - 1}\right) \\ p _ {\gamma_ {4}} \left(s _ {t} ^ {\bar {a} r} \mid \mathbf {s} _ {\mathbf {t} - \tau : \mathbf {t} - 1}\right) \end{array} \right. \right. \tag {43}
$$
Data Generation We generate synthetic datasets with 100,000 data points according to the generating process in Equation 43, which satisfies the identifiability conditions stated in Theorem 1. The latent variables  $\mathbf{s_t}$  have 8 dimensions, where  $s_t^{ar} = s_t^{\bar{a}r} = s_t^{\bar{a}\bar{r}} = s_t^{\bar{a}\bar{r}} = 2$ . At each timestep, a one-hot action of dimension 5, denoted as  $a_t$ , is taken. The lag number of the process is set to  $\tau = 2$ . The observation model  $p_{\theta}(o_t,|\mathbf{s_t})$  is implemented using a random three-layer MLPs with LeakyReLU units. The reward model  $p_{\theta}(r_t,|s_t^r)$  is represented by a random one-layer MLP. It's worth noting that the reward model is not invertible due to the scalar nature of  $r_t$ . Four distinct transition functions, namely  $p_{\gamma_1}, p_{\gamma_2}, p_{\gamma_3}$ , and  $p_{\gamma_4}$ , are employed and modeled using random one-layer MLP with LeakyReLU units. The process noise is sampled from an i.i.d. Gaussian distribution with a standard deviation of  $\sigma = 0.1$ . To simulate nonstationary noise for various latent variables in RL, the process noise terms are coupled with the historical information by multiplying them with the average value of all the time-lagged latent variables, as suggested in [19].
D.2 Modified Cartpole
We have modified the original Cartpole environment by introducing two distractors. The first distractor is an uncontrollable Cartpole located in the upper portion of the image, which does not affect the rewards. The second distractor is a controllable green light positioned below the reward-relevant Cartpole in the lower part of the image, but it is not associated with any rewards. The task-irrelevant cartpole undergoes random actions at each time step and stops moving when its angle exceeds 45 degrees or goes beyond the screen boundaries. The action space consists of three independent degrees of freedom: direction (left or right), force magnitude (10N or 20N), and green light intensity (lighter or darker). This results in an 8-dimensional one-hot vector. The objective of this variant is to maintain balance for the reward-relevant cartpole by applying suitable forces.
D.3 Variant of Robodesk
The RoboDesk environment with noise distractors [9] is a control task designed to simulate realistic sources of noise, such as flickering lights and shaky cameras. Within the environment, there is a large TV that displays natural RGB videos. On the desk, there is a green button that controls both the hue of the TV and a light. The agent's objective is to manipulate this button in order to change the TV's hue to green. The agent's reward is determined based on the greenness of the TV image. In this environment, all four types of information are present (see Table 1).
Figure 8: Visualization of the environments used in our experiments.
Ctrl + Rew	Ctrl + Rew	Ctrl + Rew	Ctrl + Rew
Modified Cartpole	Agent	Agent	Green Light	Distractor cartpole
Robodesk	Agent, Button, Light on desk	TV content, Button sensor noise	Blocks on desk, Handle on desk, Other movable objects (Green hue of TV)	Jittering and flickering environment lighting, Jittering camera
Noiseless	Agent	(Agent)	—	—
Video Background	Agent	(Agent)	—	Background
DMC	Video Background + Noisy Sensor	Agent	(Agent) Background	—
Video Background + Camera Jittering	Agent	(Agent)	—
Background Jittering camera
Table 1: Categorization of various types of information in the environments we evaluated. We use the red color to emphasize the categorization difference between IFactor and Denoised MDP. Unlike Denoised MDP that assumes independent latent processes, IFactor allows for causally-related processes. Therefore, in this paper, the term "controllable" refers specifically to one-step controllability, while "reward-relevant" is characterized by the conditional dependence between  $s_t^*$  and the cumulative reward variable  $R_t$  when conditioning on  $(s_{t-1}, a_{t-1:t})$ . Following this categorization, certain agent information can be classified as (one-step) uncontrollable (including indirectly controllable and uncontrollable factors) but reward-relevant factors, such as some position information determined by the position and velocity in the previous time-step rather than the action. On the other hand, the green hue of TV in Robodesk is classified as controllable but reward-irrelevant factors, as they are independent of the reward given the state of the robot arm and green button, aligning with the definition of  $s_t^{a\tilde{r}}$ .
D.4 Variants of DeepMind Control Suite
Four variants [9] are introduced for each DMC task:
- Noiseless: Original environment without distractors.
Video Background: Replacing noiseless background with natural videos [37] (Ctrl + Rew).
- Video Background + Sensor Noise: Imperfect sensors sensitive to intensity of a background patch ( $\overline{\mathbf{Ctrl}} + \mathbf{Rew}$ ).
- Video Background + Camera Jittering: Shifting the observation by a smooth random walk  $(\overline{\mathbf{Ctrl}} + \overline{\mathbf{Rew}})$ .