SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The optimization of the world model involves joint optimization of its four components to maximize the variational lower bound [24] or, more generally, the variational information bottleneck [25, 26]. The bound encompasses reconstruction terms for both observations and rewards, along with a KL regularizer:
$$
\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} \| p _ {\gamma}\right). \tag {5}
$$
Further, the KL regularizer  $\mathcal{I}_{\mathrm{D}}^{t}$  can be decomposed into four components based on our factorization of the state variables. We introduce additional hyperparameters to regulate the amount of information contained within each category of variables:
$$
\mathcal {J} _ {\mathrm {D}} ^ {t} = - \beta_ {1} \cdot \mathrm {K L} \left(q _ {\phi_ {1}} \| p _ {\gamma_ {1}}\right) - \beta_ {2} \cdot \mathrm {K L} \left(q _ {\phi_ {2}} \| p _ {\gamma_ {2}}\right) - \beta_ {3} \cdot \mathrm {K L} \left(q _ {\phi_ {3}} \| p _ {\gamma_ {3}}\right) - \beta_ {4} \cdot \mathrm {K L} \left(q _ {\phi_ {4}} \| p _ {\gamma_ {4}}\right). \tag {6}
$$
Additionally, we introduce two supplementary objectives to explicitly capture the distinctive characteristics of the four distinct representation categories. Specifically, we characterize the reward-
(a) Learning world model from experiences.
(b) Learning policy by imagination.
Figure 2: (a) The agent learns the disentangled latent dynamics from prior experiences. The yellow arrow represents a one-one mapping from  $h_t^*$  to  $s_t^*$  with the same superscript. (b) Within the latent space, state values and actions are forecasted to maximize future value predictions by backpropagating gradients through imagined trajectories. Only  $s_t^r$  (reward-relevant) are used for policy optimization.
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})$ . 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}^{\bar{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 {7}
$$
The conditional mutual information can be expressed as the disparity between two mutual information.
$$
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),
$$
$$
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). \tag {8}
$$
After removing the common term, we leverage mutual information neural estimation [27] to approximate the value of mutual information. Thus, we reframe the objective in Eq.7 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 {9}
$$
We employ additional neural networks, parameterized by  $\alpha$ , to estimate the mutual information. To incorporate the conditions from the original objective, we apply the stop-gradient operation  $\mathbf{sg}$  to the variable  $s_{t-1}^r$ . Similarly, to ensure that the representations  $s_t^a$  are directly controllable by actions, while  $s_t^{\tilde{a}}$  are not, we maximize the following objective:
$$
\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\}. \tag {10}
$$
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^a$  cannot. The total objective for learning the world model is summarized as:
$$
\mathcal {J} _ {\text {T O T A L}} = \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}. \tag {11}
$$
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. A thorough derivation and discussion of the objective function are given in Appendix B.
It is important to note that our approach to learning world models with identifiable factorization is independent of the specific policy optimization algorithm employed. In this work, we choose to build upon the state-of-the-art method Dreamer [6, 7], which iteratively performs exploration, model-fitting, and policy optimization. As illustrated in Figure 2, the learning process for the dynamics involves the joint training of the four categories of latent variables. However, only the variables of  $s_t^r$  are utilized for policy learning. We provide the pseudocode for the IFactor algorithm in Appendix C.
In this section, we begin by evaluating the identifiability of our method using simulated datasets. Subsequently, we visualize the learned representations of the four categories in a cartpole environment that includes distractors. Further, we assess the advantages of our factored world model in policy learning by conducting experiments on variants of Robodesk and DeepMind Control Suite. The reported results are aggregated over 5 runs for the Robodesk experiments and over 3 runs for others. A detailed introduction to each environment is provided in Appendix D, while experiment details and hyperparameters can be found in Appendix E. The source code is available at https://github.com/AlexLiuyuren/IFactor
5.1 Latent State Identification Evaluation
5.1.1 Synthetic environments
Evaluation Metrics and Baselines We generate synthetic datasets that satisfy the identifiability conditions outlined in the theorems (see Appendix D.1). To evaluate the block-wise identifiability of the four categories of representations, we follow the experimental methodology of [28] and compute the coefficient of determination  $R^2$  from  $s_t^*$  to  $\hat{s}_t^*$ , as well as from  $\hat{s}_t^*$  to  $s_t^*$ , for  $* \in \{ar, \bar{a}r, a\bar{r}, \bar{a}\bar{r}\}$ .  $R^2$  can be viewed as the identifiability score.  $R^2 = 1$  in both directions suggests a one-to-one mapping between the true latent variables and the recovered ones. We compare our latent representation learning method with Denoised MDP [9]. We also include baselines along the line of nonlinear ICA: BetaVAE [29] and FactorVAE [30] which do not consider temporal dependencies; SlowVAE [31] and PCL[32] which leverage temporal constraints but assume independent sources; and TDRL [19] which incorporates temporal constraints and non-stationary noise but does not utilize the action and reward information in the Markov Decision Process.
Results Figure 9 demonstrates the effectiveness of our method in accurately recovering the four categories of latent state variables, as evidenced by high  $R^2$  values ( $>0.9$ ). In contrast, the baseline methods exhibit distortions in the identification results. Particularly, Denoised MDP assumes independent latent process for  $x_{t}, y_{t}$ , and the existence of instantaneous causal effect from  $x_{t}, y_{t}$  to  $z_{t}$  (see Figure 1(d)), leading to unidentifiable latent variables without further intervention. BetaVAE, FactorVAE, SlowVAE, and PCL, which assume independent sources, show subpar performance. Although TDRL allows for causally-related processes under conditional independence assumptions, it fails to explicitly leverage the distinct transition structures of the four variable categories. More results on the identifiability scores of baselines are provided in Appendix E.1.
(c)
Figure 3: Simulation results: (a) The coefficient of determination  $(R^2)$  obtained by using kernel ridge regression to regress estimated latents on true latents.(b) The  $R^2$  obtained by using kernel ridge regression [33] to regress true latents on estimated latents. (c) The average  $R^2$  over four types of representations during training (True latents  $\rightarrow$  Estimated latents).(d) The average  $R^2$  over four types of representations during training (Estimated latents  $\rightarrow$  True latents).
5.1.2 Modified Cartpole with distractors
We have introduced a variant of the original Cartpole environment by incorporating two distractors. The first distractor is an uncontrollable Cartpole located in the upper portion of the image, which is irrelevant to the rewards. The second distractor is a controllable but reward-irrelevant green light positioned below the reward-relevant Cartpole in the lower part of the image.
After estimating the representation model, we utilize latent traversal to visualize the four categories of representations in the modified Cartpole environment (see Figure 4). Specifically, the recovered  $s_t^{a\bar{r}}$  corresponds to the position of the cart, while  $s_t^{\bar{a} r}$  corresponds to the angle of the pole. This demonstrates the ability of our method to automatically factorize reward-relevant representations based on (one-step) controllability: the cart experiences a force at the current time step, whereas the pole angle is influenced in the subsequent time step. Note that during the latent traversal of  $s_t^{\bar{a}\bar{r}}$ , the position of the cart remains fixed, even though the pole originally moves with the cart in the video. Additionally,  $s_t^{a\bar{r}}$  corresponds to greenness of the light, and  $s_t^{\bar{a}\bar{r}}$  corresponds to the uncontrollable and reward-irrelevant cartpole. These findings highlight the capability of our representation learning method to disentangle and accurately recover the ground-true latent variables from videos. The identifiability scores of four categories of representations are provided in Appendix E.2.
Figure 4: Latent traversal for four types of representations in the modified Cartpole environment.
5.2 Policy Optimization Evaluation
To assess the effectiveness of our method in enhancing policy learning by utilizing minimal yet sufficient representations for control, we conducted experiments on more complex control tasks. One of these tasks is a variant of Robodesk [34], which includes realistic noise elements such as flickering lights, shaky cameras, and a dynamic video background. In this task, the objective for the agent is to change the hue of a TV screen to green using a button press. We also consider variants of DeepMind Control Suite (DMC) [35, 9], where distractors such as a dynamic video background, noisy sensor readings, and a jittering camera are introduced to the original DMC environment. These additional elements aim to challenge the agent's ability to focus on the relevant aspects of the task while filtering out irrelevant distractions. Baseline results except for DreamerPro [36] are derived from the Denoised MDP paper. We have omitted the standard deviation of their performance for clarity.
Evaluation Metrics and Baselines We evaluate the performance of the policy at intervals of every 10,000 environment steps and compare our method with both model-based and model-free approaches. Among the model-based methods, we include Denoised MDP [9], which is the state-of-the-art method for variants of Robodesk and DMC. We also include DreamerPro [36], TIA [11] and Dreamer [6] as additional model-based baselines. For the model-free methods, we include DBC [37], CURL[38], and PI-SAC [39]. To ensure a fair comparison, we have aligned all common hyperparameters and neural network structures with those used in the Denoised MDP [9].
5.2.1 RoboDesk with Various Noise Distractors
In Figure 5, the left image demonstrates that our model IFactor achieves comparable performance to Denoised MDP while outperforming other baselines. The results of baselines except for DreamerPro are directly copied from the paper of Denoised MDP [9] (its error bar cannot be directly copied), and the replication of Denoised MDP results is given in Appendix E.3. Furthermore, to investigate whether  $s_t^r$  serves as minimal and sufficient representations for policy learning, we retrain policies using the Soft Actor-Critic algorithm [40] with different combinations of the four learned categories as input. Remarkably, the policy trained using  $s_t^r$  exhibits the best performance, while the performance of other policies degrades due to insufficient information (e.g.,  $s_t^{ar}$  and  $s_t^{\bar{r}}$ ) or redundant information
Figure 5: Results on Robodesk. We have omitted the standard deviation of the performance of baselines for clarity. The left image shows that our method IFactor achieves comparable performance with Denoised Mdp, outperforming other baselines. The middle image shows the policy optimization process by SAC [40] using representations learned by IFactor, where policies that use  $s_t^r$  as input achieve the best performance. The right image shows the latent traversal of the representations.
(e.g.,  $s_t^r + s_t^{a\bar{r}}$  and  $\mathbf{s_t}$ ). Moreover, we conduct latent traversal on the learned representations to elucidate their original meaning in the observation. Interestingly, our model identifies the greenness of the screen as  $s_t^{a\bar{r}}$ , which may initially appear counter-intuitive. However, this phenomenon arises from the robot arm pressing the green button, resulting in the screen turning green. Consequently, the screen turning green becomes independent of rewards conditioned on the robot arm pressing the green button, aligning with the definition of  $s_t^{a\bar{r}}$ . This empirically confirms the minimality of  $s_t^r$  for policy optimization.
5.2.2 DeepMind Control Suite (DMC)
Figure 6 demonstrates the consistent superiority of our method over baselines in the Deepmind Control Suite (DMC) across various distractor scenarios. Our approach exhibits robust performance in both noiseless and noisy environments, demonstrating the effectiveness of our factored world model in eliminating distractors while preserving essential information for effective and efficient control in policy learning. The ablation study of IFactor's objective function terms is presented in Appendix B.1 and Appendix E.3, which shows that the inclusion of mutual information constraints is essential to promote disentanglement and enhance policy performance. The policy performance that includes the standard error is provided in Appendix E.4. Visualization of the learned representations is also provided in Appendix E.5.
World Model Learning in RL. Image reconstruction [41, 42] and contrastive learning [43, 44, 45, 46, 47] are wildly used to learn representations in RL. Dreamer and its subsequent extensions [5, 6, 7, 8] adopt a world model-based learning approach, where the policy is learned solely through dreaming in the world model. However, the agents trained by these techniques tend to underperform in noisy environments [9]. To address this issue, numerous approaches have been introduced to enhance robustness to distractions. Task Informed Abstractions (TIA) [11] explicitly partition the latent state space into reward-relevant and reward-irrelevant features. Denoised MDP [9] takes a step further and decomposes the reward-relevant states into controllable and uncontrollable components. InfoPower [48] prioritizes information that is correlated with action based on mutual information. Iso-Dream [12] learns controllable and noncontrollable sources of spatiotemporal changes on isolated state transition branches. Our method extends these work by accommodating a more general factorization with block-wise identifiability. Recent research also explores reconstruction-free representation learning methods [36, 49, 50, 51]. DreamerPro [36] combines Dreamer with prototypes, which distills temporal structures from past observations and actions. Temporal Predictive Coding [50] encodes elements in the environment that can be predicted across time. Contrastively-trained Structured World Models [51] utilize graph neural networks and a contrastive approach for representation learning in environments with compositional structure. Latent states learned by these methods are not identifiable due to the reconstruction-free property, where the mixing function is not invertible anymore. A detailed comparison between IFactor and related work is also given in Appendix F.
Identifiability in Causal Representation Learning. Temporal structure and nonstationarities were recently used to achieve identifiability in causal representation learning [52]. Methods such as TCL
Figure 6: Policy optimization on variants of DMC with various distractors.
[53] and PCL [32] leverage nonstationarity in temporal data through contrastive learning to identify independent sources. CITRIS [54] and iCITRIS [55] takes advantage of observed intervention targets to identify causal factors. LEAP [18] introduces identifiability conditions for both parametric and nonparametric latent processes. TDRL [19] explores identifiability of latent processes under stationary environments and distribution shifts. ASRs [13] establish the identifiability of the representations in the linear-gaussian case in the RL setting. Our work extends the theoretical results in ASRs to enable block-wise identifiability of four categories of latent variables in general nonlinear cases.
In this paper, we present a general framework to model four distinct categories of latent state variables within the RL system, based on their interactions with actions and rewards. We establish the blockwise identifiability of these latent categories in general nonlinear cases, under weak and realistic assumptions. Accordingly, we propose IFactor to extract four types of representations from raw observations and use reward-relevant representations for policy optimization. Experiments verify our theoretical results and show that our method achieves state-of-the-art performance in variants of the DeepMind Control Suite and RoboDesk. The basic limitation of this work is that the underlying latent processes are assumed to have no instantaneous causal relations but only time-delayed influences. This assumption does not hold true if the resolution of the time series is much lower than the causal frequency. We leave the identifiability of the latent dynamics in the presence of instantaneous causal and the extension of our method to heterogeneous environments to future work.
This work is supported by National Key Research and Development Program of China (2020AAA0107200), the Major Key Project of PCL (PCL2021A12), China Scholarship Council, NSF Grant 2229881, the National Institutes of Health (NIH) under Contract R01HL159805, a grant from Apple Inc., a grant from KDDI Research Inc., and generous gifts from Salesforce Inc., Microsoft Research, and Amazon Research. We thank Guangyi Chen, Weiran Yao and the anonymous reviewers for their support and helpful discussions on improving the paper.
[49] Arnav Kumar Jain, Shivakanth Sujit, Shruti Joshi, Vincent Michalski, Danijar Hafner, and Samira Ebrahimi Kahou. Learning robust dynamics through variational sparse gating. In Advances in neural information processing systems, volume 35, pages 1612-1626, New Orleans, Louisiana, 2022.
[50] Tung D. Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 8130-8139, Virtual Event, 2021.
[51] Thomas N. Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In Proceedings of the 8th International Conference on Learning Representations, Addis Ababa, Ethiopia, 2020.
[52] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proc. IEEE, 109(5):612-634, 2021.
[53] Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ICA. In Advances in Neural Information Processing Systems 29, pages 3765–3773, Barcelona, Spain, 2016.
[54] Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M. Asano, Taco Cohen, and Stratis Gavves. CITRIS: causal identifiability from temporal intervened sequences. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pages 13557-13603, Baltimore, Maryland, 2022.
[55] Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. Causal representation learning for instantaneous and temporal effects in interactive systems. In Proceedings of the 11th International Conference on Learning Representations, 2023.
[56] P. Spirtes, C. Glymour, and R. Scheines. Causation, Prediction, and Search. Spring-Verlag Lectures in Statistics, 1993.
[57] J. Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, Cambridge, 2000.
[58] Julius von Kugelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Scholkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems, pages 16451-16467, Virtual Event, 2021.
[59] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22:268:1-268:8, 2021.
A Theoretical Proofs
Proposition 1 shows that  $s_t^r$ , which has directed paths to  $r_{t + \tau}$  (for  $\tau \geq 0$ ), is minimally sufficient for policy learning that aims to maximize the future reward and can be characterized by conditional dependence with the cumulative reward variable  $R_{t}$ .
Proposition 1. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data,  $s_t^r \subseteq \mathbf{s}_t$  is a minimal subset of state dimensions that are sufficient for policy learning, and  $s_{i,t} \in s_t^r$  if and only if  $s_{i,t} \not\perp R_t|a_{t-1:t}, s_{t-1}^r$ .
We first give the definitions of the Markov condition and the faithfulness assumption [56, 57], which will be used in the proof.
Definition 1 (Global Markov Condition). The distribution  $p$  over a set of variables  $\mathbf{V}$  satisfies the global Markov property on graph  $G$  if for any partition  $(A, B, C)$  such that if  $B$ $d$ -separates  $A$  from  $C$ , then  $p(A, C|B) = p(A|B)p(C|B)$ .
Definition 2 (Faithfulness Assumption). There are no independencies between variables that are not entailed by the Markov Condition in the graph.
Below, we give the proof of Proposition 1.
Proof. The proof contains the following three steps.
- In step 1, we show that a state dimension  $s_{i,t}$  is in  $s_t^r$ , that is, it has a directed path to  $r_{t + \tau}$ , if and only if  $s_{i,t} \not\perp R_t|a_{t - 1:t}, \mathbf{s}_{t - 1}$ .
- In step 2, we show that for  $s_{i,t}$  with  $s_{i,t} \nsubseteq R_t | a_{t-1:t}, s_{t-1}$ , if and only if  $s_{i,t} \nsubseteq R_t | a_{t-1:t}, s_{t-1}^r$ .
- In step 3, we show that  $s_t^r$  are minimally sufficient for policy learning.