Learning World Models with Identifiable Factorization
Authors: Yu-Ren Liu
Biwei Huang
Zhengmao Zhu
Honglong Tian
Mingming Gong
Yang Yu
Kun Zhang
Venue: NeurIPS 2023 poster (2023)
Keywords: Model-based Reinforcement Learning; Causal Representation Learning;
TL;DR: We propose a general framework to model four distinct categories of latent state variables within the RL system with block-wise identifiability.

Abstract: Extracting a stable and compact representation of the environment is crucial for efficient reinforcement learning in high-dimensional, noisy, and non-stationary environments.  Different categories of information coexist in such environments -- how to effectively extract and disentangle the information remains a challenging problem. In this paper, we propose IFactor, a general framework to model four distinct categories of latent state variables that capture various aspects of information within the RL system, based on their interactions with actions and rewards. Our analysis establishes block-wise identifiability of these latent variables, which not only provides a stable and compact representation but also discloses that all reward-relevant factors are significant for policy learning. We further present a practical approach to learning the world model with identifiable blocks, ensuring the removal of redundancies but retaining minimal and sufficient information for policy optimization. Experiments in synthetic worlds demonstrate that our method accurately identifies the ground-truth latent variables, substantiating our theoretical findings. Moreover, experiments in variants of the DeepMind Control Suite and RoboDesk showcase the superior performance of our approach over baselines.

Learning World Models with Identifiable Factorization
Yu-Ren Liu Biwei Huang +2 Zhengmao Zhu! Honglong Tian',
Mingming Gong Yang Yu..167, Kun Zhang
National Key Laboratory for Novel Software Technology. Nanjing University, China
University of California San Diego, USA
Carnegie Mellon University, USA Mohamed bin Zayed University of Artificial Intelligence, UAE
University of Melbourne, Australia
Polixir.ii, China
7 Peng Cheng Laboratory, China
(liuyr, zhuzm, tianhl, yuy}@lamda nju. edu. bih007@ucsd. edu
mingming. gonggunimgg@mm edu. au kunzl@cmu. edu
Abstract
Extracting stable and compact representation of the environment is crucial for efficient reinforcement learning in high-dimensional, noisy, and non-stationary environments. Different categories of information coexist in such environments how effectively extract and disentangle the information remains a challenging problem. In this paper, we propose IFactor, general framework to model four distinct categories of latent state variables that capture various aspects of informa- tion within the RL system, based on their interactions with actions and rewards. Our analysis establishes block-wise identifiability of these latent variables, which not only provides stable and compact representation but also discloses that all reward-relevant factors are significant for policy learning. We further present a practical approach to learning the world model with identifiable blocks, ensuring the removal of redundancies but retaining minimal and sufficient information for policy optimization. Experiments in synthetic worlds demonstrate that method accurately identifies the ground-truth ttuut latent variables, substantiating our thcoretical findings. Moreover, experiments in variants of the DeepMind Control Suite and RoboDesk showcase the superior performance of our approach over baselines.
Introduction
Humans excel extracting various categories of information from complex environments M1 22. By effectively distinguishing between resk-relevant information and noise, humans can learn efficiently and avoid distractions. Similarly, in the context of reinforcement learning, it's crucial for an agent to precisely extract information from high-dimensional, noisy, and non-stationary environments.
tackle this challenge by learning compact representations from images and modeling dynamics using oow-dimensional features. Recent research has demonstrated that learning policies through latent imagination in world models significantly enhances sample efficiency (5 6 7 g. However, these approaches often treat all information as an undifferentiated whole, leaving policies susceptible to irrelevant distractions [9] and lacking transparency in dicision-making [10].
World models
This paper investigates the extraction and disentanglement of diverse information types within an environment. To address this, we tackle two fundamental questions: 1) How can we establish a
Equal contribution. t: Corresponding authors.
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

Usconrofibble
Revard Relevant
Spel poodias aal daloncokas
Survendng vdacke
Masc coediinne
Remade
Inclevani
seciery
(a) Car-driving task. (b) Four types of states.
(c) Iaacoo (Ours).
(d) Denoised MDP.
Figure I: (a) An illustrative example of the car-driving task. (b) Four categories of latent state variables in the car driving task. (c) The structure of our world model. Grey nodes denote observed variables and other nodes are unobserved. We allow causally-rekie variables and prove that they are respectively identifiable up to block-wise invertible transformations. We show that both sy" and sf" are essential for policy optimization. The inclusion of the gray dashed line from (s.", """ to a signifies that the action could be determined by reward-relevant variables. (d) The structure of Denoised MDP 19]. It assumes the latent processes of X, and y, are independent and uses only x, for policy optimization. The gray dashed line from x,1 to a,- shows that the action could be determined only based the controllable and reward-relevant latent variables. Further, the existence of instantaneous causal effect from X; and y; to 2, renders the latent process unidentifiable without extra intervention on latent states
latent processes for four types of latent
comprehensive classification system for different information categories in diverse decision scenarios? 2) Can the latent state variables, classified according to this system, be accurately identified? Prior research has made progress in answering the first question. Task Informed Abstractions [I]] partition the state space into reward-relevant and
features, assuming independent latent processes for each category. Iso-Dream 12] learns controllable and noncontrollable sources of spatiotemporal changes on isolated state transition branches. Denoised MDP ( further decomposes reward-relevan states into controllable and uncontrollable components, also assuming independent latent processes. However, the assumption of independent latent processes is overly restrictive and may lead to decomposition degradation in many scenarios. Moreover, none of these approaches guarantee the identifiability of representations, potentially leading to inaccurate recovery of the underlying latent variables. While discussions on the identifiability of representations in reinforcement learning (RL) exist under linear assumptions [33, the question of identifiability remains unexplored for general nonlinear cases.
In this paper, we present IFactor, general framework to model four distinct categories of latent state variables within the RL system. These variables capture different aspects of information based on their interactions with actions and rewards, providing transparent representations of the following aspects: (i) reward-relevant and controllable parts, (ii) wwwrrdirrelevant but uncontrollable parts, (ii)) controllable but -ireeeiirreeevat parts, and (iv) unrelated noise (see Section 2). Diverging from prior methods, our approach employs general factorization for the latent state variables, allowing for causally-related latent processes. We theoretically establish the block-wise identifiability of these four variable categories in general nonlinear cases under weak and realistic assumptions. Our findings challenge the conclusion drawn in the Denoised MDP (9) that only controllable and -weard-reeevatt variables required for policy optimization. We emphasize the necessity of considering states that directly or indirectly influence the reward during the controllability. To learn these four categories of latent variables, we propose a principled approach that involves optimizing an evidence lower bound and integrating multiple novel mutual information constraints. Through simulations on synthetic data, we demonstrate the accurate identification of true latent variables, validating our theoretical findings. Furthermore, our method achieves state-of-the-art performance on variants of RoboDesk and DeepMind Control Suite.
process, irrespective of their
2 Four Categories of (Latent) State Variables in RL
For generality, consider tasks in the form of Partially Observed Markov Decision Process (POMDP) [TS), which is described as M + (S, A. Ω2 R, T, O, y). where S is the latent state space, A is the action space, Ω2 is the observation space, R 5 x A - R defines the reward function, T: XA - ASS is the transition dynamics, O: S A - A(S) is the observation function, y E [0. 1] is the discount factor. We use A(S) to denote the set of all distributions over S. The agent can interact with the environment to get sequences of observations ((,, a, r;>) The objective is to find a policy acting based on history observations, that maximizes the expected cumulative (discounted) reward.

---

In the context of POMDPs, extracting a low-dimensional state representation from high- digh--imensionall observations is crucial. Considering that action and reward information is fundamental to decision- making across various scenarios, we disentangle the latent variables in the environment into four distinct categories, represented S, [s sar, str s/"), J. based on their relationships with action and reward (see Figure [[c] as a graphical illustration):
Type I: sf" has an incident edge from 4,11, and there is a directed path from st" to ry- Type 2: sfi has no incident edge from 4,-1, and there is a directed path from of" to ri-
Type 3: sqf has an incident edge from 4:11, and there is no directed path from sf to r.-
Type 4: syl has no incident edge from 4-1, and there is no directed path from syi to ri-
sf" represents controllable and reward-relevant state variables that are essential in various scenarios. Taking the example of a car driving context, s²" encompasses driving-related states like the current speed, position, and direction of the car. These latent variables play a critical role in determining the driving policy since actions such steering and braking directly influence these states, which, in turn, have a direct impact on the reward received.
o refers to reward-relevant state variables that are beyond our control. Despite not being directly controllable, these variables still necessary for policy learning. In the context of car driving, st" includes factors like surrounding vehicles and weather conditions. Although we cannot directly control other cars, our car's status influences their behavior: if we attempt to cut in line unethically, the surrounding vehicles must decide whether to yield or block our behavior. As a result, we need to adjust our actions based on their reactions. Similarly, the driver must adapt their driving behavior to accommodate various weather conditions, even though the weather itself is uncontrollable.
The state variables denoted as xy" consist of controllable but reward-irrelevat factors. Examples of sf could be the choice of music being played or the positioning of ornaments within the car. On the other hand, s/d represents uncontrollable and reward- -ererdt latent variables such as the remote scenery. Both se" and se" do not have any impact on current or future rewards and are unrelated to policy optimization.
We next build a connection between the graph structure and statistical dependence between the variables in the RL system, so that the different types of state variables can be characterized from the data. To simplify symbol notation,
define sy := (s.", sf"), s. (s/² )²²), sf == (s57 sf") and sf := (s,ir s("). Specifically, the following proposition shows that sf that has directed paths to rrrr (for T> 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, = Er-r Tre. which has also been shown in [B]
Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data, si C is minimal subset of state dimensions that are sufficient for policy learning, and Sig E s, if and only if si, + R.aa-12- SH-1
Proposition 1.
Moreover, the proposition below shows that sf, that has a directed edge from a-1, can be directly controlled by actions and can be characterized by conditional dependence with the action variable. Proposition 2. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data. S, S St is minimal subset of state dimensions that are sufficient for direct control, and Si, si if and only if Sig 4 a4-1/51-1-
Furthermore, based on Proposition and Proposition 2 we can further differentiate s!", d0 ss" from sf and Si which is given in the following proposition.
Proposition 3. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithfull to the measured data, we can build a connection between the graph structure and statistical independence of causal variables in the RL system, with (1) su E st" if and only if St, M. Reaa-11, s/ا and S₁, 4L a1-1[Sr-,, (2) sig E StG if and only if sis th Reaa, & and Sis LL a--1lst-1, (3) Si s(" if and only if st, 1 Rla,, S/ and stg 4 a-11/8 and (4) sia E s.r if and only if St, IL R144-125 $/-1 and Sis IL a-111s--1
By identifying these four categories of latent state variables, we can achieve interpretable input for policies, including Soreward-relevan and directly controllable parts, (ii) relerntrelevant but not controllable parts, (iii) controllable but non- reward-releevan parts, and (iv) unrelated noise. Furthermore, policy training will become more sample efficient and robust to task-irrelevant changes by utilizing only 5/" and s/ for policy learning. We show the bise identifiability of the four categories of variables in the next section.
3

---

3 Identifiability Theory
Identifying causally-related latent variables from observations is particularly challenging, as latent variables are generally not uniquely recoverable [7]. Previous work in causal representation learning and nonlinear-ICA has established identifiability conditions for non- con--aamettri temporally latent causal processes [8 19]. However, these conditions are often too stringent in reality, and none of these methods have been applied to complex control tasks. In contrast to component-wiss identifiability. our focus lies on the block-wise identifiability of the four categories of latent state variables. From a policy optimization standpoint, this is sufficient because we do not need the latent variables be recovered up permutation and component-wis invertible nonlinearities. This relaxation makes the conditions more likely to hold true and applicable to a wide range of RL tasks. We provide proof of the block-wise identifiability of the four types of latent variables. To the best of our knowledge, we are the first to prove the identifiability of disentangled latent state variables in general nonlinear cases for RL tasks.
According the causal process in the RL system (as described in Eq.1 in I3), build the following mapping from latent state variables s, to observed variables on and future cumulative reward R₂
R, f(s,, s,.
(1)
where
O, = fis,, R = f1(s, n.).
(2)
Here, note that to recover sy, it is essential to take into account all future rewards frr, because any state dimension Si E S; that has a directed path to the future reward Fers for t > 0, is involved in s/. Hence, we consider the mapping from s to the future cumulative reward R, and 7/, represents residuals, except sy, that have an effect to R,
Below, we first provide the definition of blockwise identifiability and relevant notations, related to 20 21 22 23
Definition (Blockwise Identifiability). A latent variable St is blockwise identifiable if there exists one-to-one mapping h(-) between &t and the estimated S, i.e., S = h(s)).
Notations. We denote by := (sf, s/. m)) and by |s] the dimension of variable s. We further denote d, := 1s d, == 1s/1. d; = 1el. da == lorl. and de == [Ry]. We denote by F the support of Jacobian Js,,, by of the support of J (S,), and by T the support of T(s;) with J (S) = (s)s) We also denote T as a matrix with the same support as T.
In addition, given subset 2 E (1,.... dz), the subspace R is defined as R (2 E R. = 0). In other words, RS refers to the subspace of Rd indicated by index set D.
We next show that the different types of states sf", st" S. and S.T are blockwise identifiable from observed image variable o, reward variable r, and action variable a, under reasonable and weak assumptions, which is partly inspired by 20 21 22 23).
Theorem 1. Suppose that the causal process in the RL system and the four categories of latent state variables can be described as that in Section 2 and illustrated in Figure][(e [[c] Under the following assumptions
Al. The mapping f in Eq. 7 is smooth and invertible with smooth inverse.
(Ficl A2. For all i (l....,do + dRl and j E Fiz, there exist 15, 1E1 so that spantJ (S") Da Rdy and there exists matrix T with its support identical to that I. so that (,,6 RR
Then, reward-relevant and controllable states S.T, reward-relevant but not controllable states s!r, but controllable states si³, and noise sif are blockwise identifiable.
In the theorem presented above, Assumption A1 only assumes the invertibility of function f, while functions fi and f2 are considered general and not necessarily invertible, as that in (23) Since the function ff is the mapping from all (latent) variables, including noise factors, that influence the observed variables, the invertibility assumption holds reasonably. However, note that it is not

---

reasonable to assume the invertibility of the function f2 since usually, the reward function is not invertible. Intuitively, Assumption A2 requires that the Jacobian varies "enough" so that cannot be contained in proper subspace of R This requirement necessary to avoid undesirable situations where the problem becomes ill -oo-pd and is essential for identifiability. A special case when this property does not hold is when the function f is linear, as the Jacobian remains constant in such with the detailed proof given in
cases. Some proof techniques of Theorem follow from 20 22
Appendix A.4
World Model with Disentangled Latent Dynamics
Based on the four categories of latent state variables, formulate world model with disentangled latent dynamics. Each component of the world model is described as follows:
Observation Model:
pe (o; St)
Reward Model:
pe (r s,)
Transition Model:
Py (5t (1-1,0-1)
Representation Model:
gs (st 0,, a1-1,
Specifically, the world model includes three generative models: an observation model, reward function, and a transition model (prior), and representation model (posterior). The observation model and reward model are parameterized by The transition model is parameterized by y and the representation model is parameterized by Φ. We assume noises are i.i.d for all models. The action a,1 has a direct effect on the latent states St, but not on the perceived signals o,, The perceived signals, or, are generated from the underlying states St, while signals r, are generated only from reward-relevant latent variables st--
According to the graphical model depicted in Figure [(c] the transition model and the representation model can be further decomposed into four distinet sub-models, according to the four categories of latent state variables, shown in equation 4
Disentangled Transition Model:
Disentangled Representation Model:
P' (st
4 (st Ot, S/-1 a--1
(st
s/-1)
(s
O0, 5/-1)
(4)
81-1 Q-1)
qs, (s.t
O1, ((-1.-1))
Py, (5." 1t-1)
Jo. (5," S(-1)
Specifically, we have p. Pyi . Py: Pys Py. and qs = qs: qo> 9s; 9a. Note this factorization differs from previous works [1]] 9 that assume independent latent processes. Instead, we only assume conditional independence among the four categories of latent variables given $ which provides more general factorization. In particular, the dynamics of S and s are dependent on S ensuring that they are unaffected by any inrelvaatteevantt latent variables present in s On the other hand, of may be influenced by all the latent variables from the previous time step. Note that the connections
between S and sf can be adjusted based on the concrete problem. Since s!" and spi are controllable variables, their determination also relies on a,11
4.1 World Model Estimation
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 [3] 26). The bound encompasses reconstruction terms for both observations and rewards, along with KL regularizer:
T' (o, Th = In (r, s) T. -Klll((()))l
(5)
Further, the KL regularizer Ih 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:
OKL T. = - (L..lll)))) KL (g..117))) -Bs KL (qo.ll)))))) -B4 KL (ge.llp)))
(6)
Additionally, we introduce two supplementary objectives to explicitly capture the distinctive char- acteristics of the four distinct representation categories. Specifically, we characterize the reward-
5

---

X
Vi+2
AVT
A
Aym
A
(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 one-one mapping from h; to sy 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 sf (reward-relevant) are used for policy optimization.
relevant representations by measuring the dependence between s and R., given d,1 and 5y-1- that is I(s,, R₂ dr-12, s To ensure that of are minimally sufficient for policy training, we maximize I(s,, R₁ a-11a- s while minimizing I(s R, a,-11. 5/-1) to discourage the inclusion of redundant information in sy concerning the rewards:
I(s/; R2 la-1a, s 1(5/; R, ap-1a. $/-1).
The conditional mutual information can expressed as the disparity between two mutual information.
1s5/; R, 11 1 [RR; 41 $/-1) ((Ria,,115 S/-1)
(8)
1(5/; R, (a-1: ( I(R;; s/ Or-1a- s'y1)) I(R; 4--1s $(-1).
After removing the common term, we leverage mutual information neural estimation 27 to approxi- mate the value of mutual information. Thus, we reframe the objective in Eq 7 as follows:
Tes = di Mo (Rr; s/, 1m sg(s; log (Rri s/ a 1. sg(s)-)))).
We employ additional neural networks, parameterized by to estimate the mutual information. To incorporate the conditions from the original objective, we apply the stop stop_ graddiet operation sg to the variable S/11- Similarly, to ensure that the representations s are directly controllable by actions, while s, are not, we maximize the following objective:
JAs d2 H4y (a, st. sg(s--1)) (a,1; $7. sg(ss-))).
(10)
Intuitively, these two objective functions ensure that s/ is predictive the reward, while s/ not; similarly, sf can be predicted by the action, whereas cannot. The total objective for learning the world model summarized as:
TTOTAL Fig, E Tb Jks + Tis) ( const.
(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. thorough derivation and discussion of the objective function are given in Appendix
It is important to note that approach learning world models with identifiable factorization is independent of the specific policy optimization algorithm employed. In this work, we choose to build upon the tfte-offthee-rr method Dreamer 6 77. which iteratively performs exploration, model-fittin, and policy optimization. As illustrated in Figure 2 the learning process for the dynamics involves the joint training of the four categories of latent vanables. However, only the variables of s are utilized for policy learning. We provide the pseudocode for the IFactor algorithm in Appendix

---

5 Experiments
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 environ- ment 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 exper- iment details and hyperparameters can be found in Appendix E The source code is available at http://dddubbcoomA/lll/uuuuuuee
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.I. 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² from Sy to S,, as well as from S to sy, for E (ar, är. ař, ar). R²
can be viewed as the identifiability score, R² 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 ( We also include baselines along the line of nonlinear ICA:
Beta VAE [29] and FactorVAE (30 which do not consider temporal dependencies; SlowVAB [31] and PCL which leverage temporal constraints but assume independent sources; and TDRL [I9
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² values (> 0.9). In contrast, the baseline methods exhibit distortions in the identification results. Particularly, Denoised MDP assumes in- dependent latent process for xr, y,, and the existence of instantaneous causal effect from yr,;, z (see Figure I leading to unidentifiable latent variables without further intervention. Beta VAE, FactorVAE, VAE, 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. Aviragaa R & to S.
R²
0.9839
0,9850
Average R
0.98
0.01
0.03
3-
0:13
0.09
0.06
10
Extinated latents
0:10
1.03
0.08
0.03
True latents
0.55
1.03
0.07
0.05
0.05
0.03
9.02
0.08
0.99
0.02
3.02
0.32
0.99
+
Thue latents
Estimated latents
Trainig Steps
Trainng Steps
(a)
(b)
(c)
(d)
Figure 3: Simulation results: (a) The coefficient of determination (R²) obtained by using kernel ridge regression to regress estimated latents true latents.(b) The R² obtained by using kernel ridge regression (33] to regress true latents on estimated latents. (c) The average over four types of representations during training (True latents Estimated (dt The average R2 over four types of representations during training (Estimated latents - 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 uncontrollable Cartpole located in the upper portion of the image, which is irrelevant to the rewards. The second distractor is a controllable but eoward-irelevamm green light positioned below the reward-relevant Cartpole in the lower part of the image.
7

---

After estimating the representation model, we utilize latent traversal to visualize the four categories of representations the modified Caripole environment (see Figure 4 Specifically, the recovered si" corresponds to the position of the cart, while sf" corresponds to the angle of the pole. This demonstrates the ability of our method to automatically factorize eeward-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 the position of the cart remains fixed, even though the pole originally moves with the cart in the video. Additionally, st" corresponds to greenness of the light, and st" corresponds to the uncontrollable and reward-irreleval cartpole. These findings highlight the capability of our representation learning method to disentangle and accurately recover the ground latent variables from videos. The identifiability scores of four categories of representations provided in Appendix E.2
Obs
st
Cart position
gär
Pole angle
SaP
Greeness
Distractor
Figure 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 variant of Robodesk [34], which includes realistic noise elements such flickering lights, shaky cameras, and dynamic video background. In this task, the objective for the agent is to change the hue of TV screen to green using button press. We also consider variants of DeepMind where distractors such as dynamic video background, noisy sensor readings, and jittering camera are introduced to the original DMC environment. These additional elements aim to challenge the agent's ability to focus 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.
Control Suite (DMC) 35 g
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 modell-ased methods, we include Denoised MDP [9] which is the state-of-ther thef-th-rrtt method for variants of Robodesk and DMC. We also include DreamerPro (36), TIA [1I] and Dreamer (6) as additional model-bard -ased baselines. For the model-free ffree methods, we include DBC [37]. CURL(38), and PI-SAC [39]. To ensure fair comparison, we have aligned all common hyperparameters and neural network structures with those used in the Denoised MDP 99.
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), investigate and the replication of Denoised MDP results is given Appendix E.3 Furthermore, whether sf 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/ exhibits the best performance, while the performance of other policies degrades due to insufficient information (c.g., s and s/) or redundant information

---

Robodiskk with Various Distractors
Using Ropresentations Loarned IFaztor
700
Rollout
611
Eqaand Raturn
Robor Arm
Eskada
200
400
Screee
300
Flickeiin Liglt
Environment Stops
Environment Steps
Figure 5: Results on Robodesk. We have omitted the standard deviation of the performance of baselines for clarity. The left image shows that method IFacoor 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 si input achieve the best performance. The right image shows the latent traversal of the representations.
(e.g., s/ s," and sj). 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"", 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 se" This empirically confirms the minimality of x for policy optimization.
5.2.2 Deep.indd 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 presented in Appendix [B]] and Appendix [E.3] which shows that the inclusion of mutual information constraints is essential promote disentanglement and enhance policy performance. The policy performance that includes the standard error is provided Appendix E.4 Visualization of the learned representations is also provided in Appendix E.5
6
Related Work
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 S adopt 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
5
and -ireediiiaeeaan features. Denoised MDP [ takes a step further and decomposes the reward-relevan 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 identifiablity. Recent research also explores reconstruction-free representation learning
methods [6, 49 50 51.. DreamerPro [6 combines Dreamer with prototypes, which distills temporal structures from past observations and actions. Temporal Predictive Coding (50) encodes 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 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
elements in the environment that can be predicted across time. Contrastivelly-rrain
Identifiability in Causal Representation Learning. Temporal structure and nonstationarities were recently used to achieve identifiability in causal representation learning [52]. Methods such as TCL

---

Video Background Noisy Sensor
Vidco Background + Camera Jittering
Noiseless
Video Background
Cheetah Run
880
403
Epioods
083
203
183
100
503
Epieods Return
583
Walker Walk
880
403
100
403
203
201
203
08
62
1300
B
Reacher Easy
Epradd
533
309
Envionment Steps
Envronment Steps
Envronnem Steps
Envronment Stops
IFactor (Ours)
DreamerPro
DenoisedMDP
TIA
Dreamer
CURL
DBC
PI-SAC
Figure 6: Policy optimization variants of DMC with various distractors.
53 and PCL 332 leverage nonstationarity in temporal data through contrastive learning identify independent sources. CITRIS (54) and iCITRIS (55) takes advantage of observed intervention targets to identify causal factors. LEAP I8 introduces identifiability conditions for both parametric and nonparametric latent processes. TDRL 19 explores identifiability of latent processes under stationary the
environments and distribution shifts. ASRs [ establish the identifiability of the representations in
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.
7 Conclusion
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 block- wise 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-an 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.
8 Acknowledgements
This work is supported by National Key Research and Development Program of China (2020AAA0107200) the Major Key Project of PCL (PCL2021A12), China Scholarship Coun- cil, NSF Grant 2229881, the National Institutes of Health (NIH) under Contract R01HL159805, grant from Apple Inc., 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.
10

---

References
[1] Le Chang and Doris Tsao. The code for facial identity in the primate brain. Cell, 169 1028, 2017.
[2] Quian Quiroga, Leila Reddy, Gabriel Kreiman, Christof Koch, and Itzhak Fried. Invariant visual representation by single neurons in the human brain. Nature, 435(70455:1102--1107 2005.
[3] Jürgen Schmidhuber. On learning to think: Algorithmic information theory for novel com- binations of reinforeement learning controllers and recurrent neural world models. CoRR, abs/1511. 1,.9249,, 2015.
[4] David Ha and Jürgen Schmidhuber. World models. CoRR, abs/ 0bss/1803.10122.. 2018.
[5] Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning, pages 2555 2565, Long Beach, California, 2019.
[6] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. In Proceedings of the 8th International Conference on Learning Representations, Addis Ababa, Ethiopia, 2020.
[7] Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In Proceedings of the 9th International Conference on Learning Representations, Virtual Event, Austria, 2021.
[8] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse domains through world models. CoRR, abs/2301.04100 2023.
[9] Tongzhou Wang. Simon Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. In Proceedings of the 39th International Conference Machine Learning, pages 22591-2261 Baltimore, Maryland, 2022.
[10] Claire Glanois, Paul Weng, Matthieu Zimmer, Dong Li, Tianpei Yang, Jianye Hao, and Wulong Liu. A survey on interpretable reinforcement learning. CoRR, abs/21 12.1112,. 2021.
[11] Xiang Fu, Ge Yang. Pulkit Agrawal, and Tommi S. Jaakkola. Learning task informed abstrac- tions. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 3480-3491, Virtual Event, 2021.
[12] Minting Pan, Xiangming Zhu, Yunbo Wang, and Xiaokang Yang. Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models. Advances neural information processing systems, volume 35, pages
New Orleans, Louisiana, 2022.
[13] Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bern- hard Schölkopf, and Kun Zhang. suutiiinuffi state representation learning for control with structural constraints. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pages
Baltimore, Maryland, 2022.
[14] Aapo Hyvärinen, Shohei Shimizu, and Patrik O. Hoyer. Causal modelling combining instanta- neous and lagged effects: an identifiable model based on non- -an-gasssiannt. In Proceedings of the 25th International Conference on Machine Learning, volume 307, pages 424-431, Helsinki, Finland, 2008
[15] Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134, 1998.
[16] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pages 4114-4124, Long Beach, California, 2019.

---

[17] Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural networks, (2(3):4294333 1999.
[18] Weiran Yao, Yuewen Sun, Alex Ho, Changyin Sun, and Kun Zhang. Learning temporally causal latent processes from general temporal data. In Proceedings of the 10th International Conference on Learning Representations, Virtual Event, 2022.
[19]
Weiran Yao, Guangyi Chen, and Kun Zhang. Temporally disentangled representation learning. In Advances in Neural Information Processing Systems 35, New Orleans, Louisiana, 2022.
[20] Sébastien Lachapelle, Pau Rodríguez, Yash Sharma, Katie Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ICA. In Proceedings of the 1st Conference on Causal Learning and Reasoning, volume 177, pages 428 484, Eureka, CA, 2022.
[21] Lingjing Kong, Shaoan Xie, Weiran Yao, Yujia Zheng, Guangyi Chen, Petar Stojanov, Victor Akinwande, and Kun Zhang. Partial disentanglement for domain adaptation. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pages 11455-11472, Baltimore, Maryland, 2022.
[22] Yujia Zheng, Ignavier Ng, and Kun Zhang. On the identifiability of nonlinear ica: Sparsity and beyond. In Advances in Neural Information Processing Systems 35, New Orleans, Louisiana, 2022.
[23] Lingjing Kong, Biwei Huang, Feng Xie, Eric Xing, Yuejie Chi, and Kun Zhang. Identification of nonlinear latent hierarchical models. CoRR, abs/2306.07916, 2023.
[24] Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduc- tion variational methods for graphical models. Machine learning, 37:183-233, 1999.
[25] Naftali Tishby, Fernando C. N. Pereira, and William Bialek. The information bottleneck method. CoRR, physics/0004077, 2000.
[26] Alexander A. Alemi, lan Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. Proceedings of the 5th International Conference on Learning Representations, Toulon, France, 2017.
[27] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, R. Devon Hjelm, and Aaron C. Courville. Mutual information neural estimation. In Proceed- ings of the 35th International Conference Machine Learning, volume 80, pages 530-539, Stockholm, Sweden, 2018
[28] Julius von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. In Advances in Neural Information Processing Systems 34, pages 16451-16467, Virtual Event, 2021.
[29] Irina Higgins, Loic Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In Proceedings of the 5th International Conference on Learning Representations, Toulon, France, 2017.
[30] Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Proceedings of the 35th International Conference Machine Learning, volume 80, pages 2654-2663, Stockholm, Sweden, 2018
[31] David A. Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge, and Dylan M. Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding. In Proceedings of the 9th International Conference on Learning Representations, Virtual Event, Austria, 2021.
[32] Aapo Hyvärinen and Hiroshi Morioka. Nonlinear ICA of temporally dependent stationary sources. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54, pages 460-469, Fort Lauderdale, FL, 2017.
12

---

[33] Vladimir Vovk. Kernel ridge regression. In Empirical Inference Festschrift in Honor of Vladimir N. Vapnik. pages 105-116, 2013.
[34] Tongzhou Wang. Robodesk with diverse set of distractors. robodesk 2022.
[35] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, abs/1801.0066 2018.
[36] Fei Deng, Ingook Jang, and Sungjin Ahn. Dreamerpro: Reconstruction craettrrccioo-free model-based rein- forcement learning with prototypical representations. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pages 4956-4975, Baltimore, Maryland, 2022.
[37] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In Pro- ceedings of the 9th International Conference on Learning Representations, Virtual Event, Austria, 2021.
[38] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised repre- sentations for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, volume 1 119, pages 5639-5650, Virtual Event, 2020.
[39] Kuang-Hung Hueni Lee, lan Fischer, Anthony Z. Liu, Yijie Guo, Honglak Lee, John F. Canny, and Sergio Guadarrama. Predictive information accelerates learning in RL. In Advances Neural Information Processing Systems 33, Virtual Event, 2020.
[40] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International conference on machine learning, pages 1861-1870, Stockholm, Sweden, 2018.
[41] Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin A. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems 28, pages 2746-2754, Quebec, Canada, 2015.
[42] Niklas Wahlström, Thomas B. Schön, and Mare Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. CoRR, ab22150..022111 2015.
[43] Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, and Sergey Levine. Time-contrastive networks: Self-supervised learning from video. In Proceedings of the 35th International Conference Robotics and Automation, pages 114-1144, Brisbane, Australia, 2018.
[44] Aäioo van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CORR, abs/1807 507.03748, 2018.
[45] Wilson Yan, Ashwin Vangipuram, Pieter Abbeel, and Lerrel Pinto. Learning predictive rep- resentations for deformable objects using contrastive estimation. In Proceedings of the 5th Conference on Robot Learning, pages 564-574, 2021.
[46] Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforcement and infomax learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Floring Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33, pages 3686-3698, Virtual Event, 2020.
[47] Zheng-Mar Zhu, Shengyi Jiang, Yu-Ren Liu, Yang Yu, and Kun Zhang. Invariant action effect model for reinforeement learning. Proceedings of the 36th AAAI Conference on Artificial Intelligence, pages 9260-9268, Virtual Event, 2022.
[48] Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Infor- mation prioritization through empowerment visual model-based RL. In Proceedings of the Tenth International Conference on Learning Representations, Virtual Event, 2022.
13

---

[49] Arna 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. Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 8130-81 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-66 6344, 2021.
[53] Aapo Hyvärinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ICA. In Advances Neural Information Processing Systems 29, pages 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 instantancous and temporal effects interactive systems. In Proceedings of the I 1tth International Conference Learning Representations, 2023.
[56] P. Spirtes, C. Glymour, and R. Scheines. Causation, Prediction, and Search. Spring-Verlays Lectures Statistics, 1993.
[57] Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, Cambridge, 2000.
[58] Julius von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, 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. babll-- Reliable reinforcement learning implementations. Journal of Machine Learning Research, 12-26811-268:. 2021.
14

---

A Theoretical Proofs
A.1 Proof of Proposition
Proposition I shows that sy, which has directed paths to Fitr (for r 2 0), is minimally sufficient for policy learning that aims maximize the future reward and be characterized by conditional dependence with the cumulative reward variable R.
Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data, si is minimal subset of state dimensions that are sufficient for policy learning, and Si s and only if si Rila-ia- $/
Proposition 1.
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 over of variables V satisfies the global Markov property on graph G if for any partition (A, B, C) such that if d-separates from C, then p(A, C)B) = p(A)B))(CCIB).
Definition (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 8;, is in s,, that is, has a directed path to ret. and only if St # R1 8:-1-
In step 2, show that for St with st HL R,la,11 S-11 if and only if su 4h R,la,-1u $/-11
In step 3, we show that sf are minimally sufficient for policy learning.
Step 1: We first show that a state dimension Si, is in sf, then Sn 4L R110;11::S-111
We prove it by contradiction. Suppose that St independent of R, given 4,tr and 5,-1- Then according to the faithfulness assumption, path to Yir,s which contradicts the assumption, because, otherwise, a,11 and s,-1 cannot break the paths between S and R, which leads to the dependence.
from the graph that st does not have directed
We next show that if st, 4L R,la,-1. S,11 then St S/-
Similarly, by contradiction suppose that St, does not have directed path Ti+.- From the graph, is easy to see that a,11 and s,- must d-separate the path between Sú and R, According to the Markov assumption, 8i, is independent of R, given dg-t and S,-1, which contradicts we have a contradiction, must be that Sú, has a directed path to YHt,, i.e. Si, E s,
the assumption. Since
Step In step 1, we have shown that Sü + R, 8,-1, if and only if it has directed path rer- From the graph, it is easy to see that for those state dimensions which have directed path to Ti-r- a-1x and 5,-1 cannot break the path between Si, and R, Moreover, for those state dimensions which do not have a directed path to ri+r, 4-11æ and s are enough to break the path between Sú, and R,
Therefore, for St,r ML R,a,,11,, S,-1, if and only if sn, ML R 5/-1'
Step 3: In the previous steps, has been shown that if a state dimension Siy is in s, then St + Rela,11, Sh_ and if state dimension Sú, is not in sf, then 8ú Jl Rla,-1,, S/-1" This implies that sf are minimally sufficient for policy learning to maximize the future reward.
A.2 Proof of Proposition
Moreover, the proposition below shows that si. which receives edge from a,11, can be directly controlled by actions and can characterized by conditional dependence with the action variable.
Proposition 2. Under the assumption that the graphical representation, corresponding the environment model, Markov and faithful to the measured data, Si S S₁ is a minimal subset of state dimensions that sufficient for direct control, and St S if and only if si, HL a--1/1-1--
15

---

Below, we give the proof of Proposition 2.
Proof. The proof contains the following two steps.
In step 1, we show that state dimension Sy is in st, that is, receives an edge from a,-1 and only if # 4,-18%-11-
In step 2, we show that sf contains a minimally sufficient subset of state dimensions that can be directly controlled by actions.
Step 1: We first show that a state dimension Si, is sy, then St, 46 d1-11st-1-
We prove by contradiction. Suppose that Si is independent of a,-1 given $,-1- Then according to the faithfulness assumption, which contradicts the assumption, because, otherwise, S,-1 cannot break the paths between Si, and a,- which leads to the dependence.
from the graph that St does not receive edge from d,-1,
We next show that if Str #h a,-1l:-1, then St, E s.y
Similarly, by contradiction suppose that St does not receive an edge from d,1 From the graph. easy to see that 5,-1 must break the path between Str and a-11 According to the Markov assumption, Sty is independent of a-1 given S,-1, which contradicts to the assumption. Since we have a contradiction, it must that Si has an edge from a;-1.
Step 2: In the previous steps, it has been shown that if a state dimension Str is in sf, then St, 4 d,-18s-1,, and if state dimension St is in sy, then St. a,-188-11- This implies that sf minimally sufficient for one-step direct control.
A.3 Proof of Proposition
Furthermore, based on Proposition I and Proposition 2 we can further differentiate s." ar st" from s and s%. which is given in the following proposition.
Proposition 3. Under the assumption that the graphical representation, corresponding to the environment model, Markov and faithful to the measured data, we can build connection between the graph structure and statistical independence of causal variables in the RL system, with (1) E st" and only Sig H. 1,10,11,/,, and St, HL. 0,-1(5t-1, (2) St, ser if and only if Sta + R,aa,1, S and Sis 0;-1(5t-1. (3) Sis se if and only if Si,, R,la,1, and Sig a,1188i-,, and (4) St E s." if and only if St, R, S/11 and Sig a--1/8t-1-
Proof. This proposition can be easily proved by levering the results from Propositions I and 2.
A.4 Proof of Theorem
According the causal process in the RL system (as described in Eq. 1 in 1 following mapping from latent state variables s, to observed variables and future cumulative reward R;
build the
[o, R,1 f(s,, sT,)))
(12)
where
O, = fil,, s
(13)
R = f2(s, nT).
Here, note that to recover sf. it is essential to take into account all future rewards r.r, because any state dimension Sy E s, that has a directed path to the future reward Ker. for > 0, is involved in sf. Hence, we consider the mapping from si to the future cumulative reward R, and 71, represents residuals, except s/. that have an effect to R.
The following theorem shows that the different types of states sf", StT Sp3 and S are blockwise identifiable from observed image variable o,, reward variable r,, and action variable a, under reasonable and weak assumptions.
Theorem 1. Suppose that the causal process in the RL. system and the four categories of latent state variables can be described as that in Section and illustrated Figure 1(c). Under the following assumptions
16

---

Al. The mapping f in Eq. 12 smooth and invertible with smooth inverse.
A2. For all i (l....,do + dæ) and j E Ti. there exist (s, tFul By and there exists matrix T with its support identical to that of J so that (J T ER
)allel ( so that spantJ G G Fid
Then, reward-relevant and controllable states s reward-relevant but not controllable states s reward-irrelevanl but controllable states sf and noise sQ", are blockwise identifiable.
In the theorem presented above, Assumption Al only assumes the invertibility of function f. while functions fi and f2 are considered general and not necessarily invertible, as that in [23]. Since the function is the mapping from all (latent) variables, including noise factors, that influence the observed variables, the invertibility assumption holds reasonably. However, note that it is not reasonable to assume the invertibility of the function f2 since usually, the reward function is not invertible. Assumption A2, which is also given 220 22 23), aims to establish a more generic condition that rules out certain sets of parameters to prevent ill- ill-posed conditions. Specifically, it ensures that the Jacobian is not partially constant. This condition typically satisfied asymptotically. and it is necessary to avoid undesirable situations where the problem becomes ill ill-posed.
Proof. The proof consists of four steps.
1. In step we show that sf = 5." U .tt is blockwise identifiable, by using the characterization that the action variable only directly influences s," and s!"
2. In step 2, we show that sy s." se" is blockwise identifiable, by using the characterization that the future cumulative reward R, is only influenced by s/" and s."
3. In step 3, we show that s/" is blockwise identifiable, by using the identifiability of s." U sq and s." S."
4. In step further show the blockwise identifiability of sf' s.f, and s
Step 1: prove the block identifiability of st.
For simplicity of notation, below, we omit the subscript
Let h f-1 We have
I(().
(14)
where h f-1 o f is the transformation between the true latent variable and the estimated one, and S:S X denotes the estimated invertible generating function. Note that as both f and are smooth and invertible, h and h is smooth and invertible.
Since h(-) is smooth over S, its Jocobian can be written as follows:
A
Jn-
(15)
C
da
D
The invertibility of h- implies that Jn is full rank. Since 5 has changing distributions over the action variable while d has invariant distributions over different values of a, we can derive that C= Furthermore, because J4-1 is full rank and C is zero matrix, D must be of full rank, which implies 11-1 is invertible, where h denotes the first derivative hー! Therefore, s" is blockwise identifiable up to invertible transformations.
Step prove the blockwise identifiability s/c
Recall that we have the following mapping:
[o, R,1 = f(S), $..1).
where
0, f(S,,
R = f2(s,,
17

---

Note that here to recover s,, we need to take into account all future rewards, because 5 contains all those state dimensions that have directed path future rewards 1+:17- "i represents all other factors, except 8,, that influence R, at time instance 1. Further note here we assume the invertibility of f, while fi and f2 are general functions not necessarily invertible.
We denote by (s', s .n)) We further denote by the dimension of by dy, the dimension of s' by dy, the dimension of š by d3. the dimension of o by do. and the dimension of R, by dR-
We denote by F the support (s),, by f the support of J j(),, and by T the support T(s). We also denote T as a matrix with the same support as T. The proof technique mainly follows 23 as well as the required assumptions, and is also related to [20 (2).
Since h f o f, we have f = f h (5). By applying the chain rule repeatedly, we have
( = ()(6) (h()).
(16)
Fi.l With Assumption A2, for any i [1, d dxl, there exists (5 s.t. )s
(F
R
Since Di forms a basis of R for any jo E Flss we can write canonical basis vector
ejo F. as:
Σ aj J,(5 k. FeF
(17)
where ay R is coefficient. Then, following Assumption A2, there exists a deterministic matrix T such that
T T
Σ
).
Rd fi'
(18)
JEF
where is due to that each element in the summation belongs to R Fi. Therefore,
Yj Ei. Fi.. E R
Equivalently, we have:
V(, x
(19)
We would like to show that 3" does not depend on s' and 1J, that is, Tú, = for i (l. and j (d,- l.....dy).
We prove it by contradiction. Suppose that 8" had dependence on s, that is, 3(rr,j) T with ju dyy) and jy E (dy- + I...., dy + dy).
Hence, there must exist i, E (d, 1......de do + dgl. such that, (i,, ET.
It follows from Equation 19 that:
(ir) xT jos Ef (...))
(20)
However, due to the structure of L2. [Jglis O, which results in a contradiction. Therefore, such (ir, jz) does not exist and 3 does not depend on s². The same reasoning implies that 3' does not dependent on 17. Thus, 5, does not depend on (s², 17). In conclusion, 3 does not contain extra information beyond y".
Similarly, we can show that (3", (f) does not contain information of s,-
Therefore, there is one-to-one mapping between s" and 8r.
Step 3: prove the blockwise identifiability of s.".
18

---

In Step and Step 2, we have shown that both 5 and s" are blockwise identifiable. That is,
S h,(s"),
(21)
So ha(s));
where ha and h, are invertible functions. According to the invariance relation of 5th We have the following relations:
8 1 ha(5 (idger
(22)
remains show that both h, == 1, and ha ha(-) do not depend on -"r and s in their arguments.
We will prove this by contradiction, following the proof technique in [88 Without loss of generality, we suppose 3/ E (1.... dgrr, 5" S s.t., derivatives. Thus, dh dop
dh (5**) # O. As h smooth, it has continuous partial Txp
# O holds true in neighbourhood of s" i.e.,
En s. Sf (s (5"1" s(")) is strictly monotonic on (str -7, St" 17).
(23)
where " denotes variable -" excluding the dimension I. We further define an auxiliary function v Sur x Str x S R20 follows:
(5 8 = (5) - h(()))
(24)
To obtain the contradiction to the invariance, it remains to show that & with a probability greater than zero w.r.t. the true generating process.
There are two situations at (sar= grr saF*)
where gar* is an arbitrary point in S
situation 1: arrr säre sa7 0:
situation 2: (are gāre gar )ar)==
In situation 1, we have identified a specific point yestrr "rr" that makes >
In situation 2, Eq. 23 implies that Vs (sfrh sq TD, Msaar, (.ST s(r) sap O.
Thus, in both situations, we can locate a point (5 are' 5" such that "h" pare .5ary) ) > 0, where gar" = =a"" in situation and s²r Essf",, sxp'* 17), sare s in situation 2.
Since composition continuous functions, always open for continuous functions, the open set Rso has an open set U E S x Sür x Sar as its preimage. Due to (sar* _ar s U, U/ is nonempty. Further, U is nonempty and open, U has a Lebesgue measure of greater than zero.
continuous. As pre-image of open sets are
As we assume that p s that P PA[П]] O. That is, & > 0 with probability greater than zero, which contradicts the invariance condition, Therefore,
fully supported over the entire domain S S x S we can deduce
can show that (s") does depend on _är
Similarly, we can show that h,(s") does not depend on S"
Finally, the smoothness and invertibility h, and ha follow from the smoothness and invertibility of h, and ha over the entire domain.
Therefore, (ha) is smooth invertible mapping between 5 and gnr That is, S is blockwise invertible.
Step 4: prove the blockwise identifiability of s/", sf" and shr
We can use the same technique in Step 3 to show the identifiability of Jär and Sa Specifically, since s, and or are identifiable, identifiable, we can show that "F is identifiable. Furthermore, since " shr and " are identifiable, can show that d² is identifiable
show that aar is identifiable. Similarly, since s and ""
19

---

B Derivation of the Objective Function
We start by defining the components the world mode as follows:
Observation Model:
polor st)
Reward Model:
po(r, 5,)
(25)
Transition Model:
(5 4-1-00-1
Representation Model:
(8 O₁, St-1, a-11)
The latent dynamies be disentangled into four catogories:
Disentongled Transition Model:
Disentangled Representation Model:
Pn (st 5 a2-1)
qs, (s."
O 5 5111
Pez (sf 1-1
(5,(5
Op S 4-1
(26)
Pys 4-1.0--1)
qss 9ss (1
Op. S1-1
(s) St-1)
qd. 90, (s,
((-1)
We follow the derivation framework in Dreamer 6 and define the information bottleneck objective for latent dynamics models [25]
max (S1T: (ou. M1.T) (1:)) B (SET, he aa:r)
(27)
where B is scalar and i, are dataset indices that determine the observations p(o, | i;) = 8(0, o,) as in (66.. Maximizing the objective leads model states that can predict the sequence of observations and rewards while limiting the amount of information extracted at each time step. We derive the lower bound of the first term in Equation 27
I (SETi (01T, (ITT a1:))
Σ In (OLr, Spr.arr)) In p( fer dxr))
=E Be
iB Σ In p (0 Fur (prrarr)
(28)
2E Σ In p (OEr Fur SLT.arr)) KL. p p (STr.41))|| [ (0, 1s) po(ro st)
= Σ In po (o, s)) In po (n
Thus, we obtain the objective function:
To = In (o, s))
Th po (r, s))
(29)
For the second term in Equation 27 we use the non-negativity of the KI. divergence to obtain an upper bound,
I (S air)
Σ In In a... p (st st-1, a-1)
For Avr)
Σ In (St s. In (5-1,
-E
(30)
Σ Io (Sr S. py (St St.1. a-1)
<B
Σ KL (go (Se (.11 1l py (St S- (1-1.
=B
20

---

According to equation 26 we have Py = Pyi Py: Pys Py, and qs Qo:
9s. 9d, 96.
KL(qs pạ) = Kliqe,
90: 90: 9a,
Pys Pyz Py: Pyaa)
= Egs [In qs + In
9
qs: In
9da
+ In
(31)
Py
Pn
Py:
Pya
= KL ll Py,) + KLiqs, ll Py.)) + KL ll Py,) + KL(qg, Py:)
We introduce additional hyperparameters to regulate the amount of information contained within each category of variables:
Tb -Bi KL (gs) ll pr.) B2 KL 90: pr.) -Bs KL (go. |1 Pr.) - -B+ KL KL (go. Pr.). (32)
KL(90.
Additionally, we introduce two supplementary objectives to explicitly capture the distinctive char- acteristics of the four distinct representation categories. Specifically, we characterize the reward- relevant representations by measuring the dependence between s; and R, given 4,-1> and 5 that is Iss,, a,-1u 5²-1) (see Figure (5(a) Note that if the action not dependent on s,, such when is randomly chosed. a, does not need to be conditioned, In this case, the mutual information turns into I(s,,, R2 (4-1, s/-)). ). To ensure that of are minimally sufficient for policy training, we maximize to discourage the inclusion of redundant Iss R₁ a--1a- s while minimizing I(s R, d,-11, 5/-1) ) information in sy concerning the rewards:
((s;; R la-ta, s - 1(5/; R, a--1a $1-1).
(33)
The conditional mutual information can expressed as the disparity between two mutual information values.
11-12 s/ I(5;; |
= [(R; s/ a-11 s/1,., (RRia,111 s( I(5;; R; (a-1aa..// ) I(R;; s/. a-11 s/21) I(Rçi a-11: s(
(34)
Combining the above two equations, eliminated the identical terms, ultimately yielding the following formula
I(R; am.1, I[R; 4-11 $
(35)
We use the Donsker- Varadhan representation to express mutual information as supremum over functions,
I(x; Y) = Dra(p(x,))) p(x)p())
sup Bpayy [T (), log Eu
(36)
Ter
We employ mutual information neural estimation [27 to approximate the mutual information value. We represent the function 7' using neural network that accepts variables (x. y) as inputs and is parameterized by a. The neural network is optimized through stochastic gradient ascent to find the supremum. Substituting and with variables defined in Equation 35 our objective is reformulated as follows:
Ths = 41 - Ma (Ry;
s/, a- sg(s,_)) las (R,; Sy. 4- sg (g)()))))
(37)
To incorporate the conditions from the original objective, we apply the stop gradienn operation to the variable S- Similarly, to ensure that the representations s are directly controllable by actions, while se are not, we maximize the following objective:
T(, a,- St-1) [6 4:11 St-1),
(38)
By splitting the conditional mutual information and eliminating identical terms, we obtain the following objective function:
I's d2 (gg(s- Iag (ar11 S (g(r.)))))))...
(39)
where 01, 022 33 Q4 can be obtained by maximizing Equation 36 Intuitively, these two objective functions ensure that sf is predictive the reward, while sf is not; similarly, sf can be predicted by the action, whereas sf cannot.
21

---

Combine the equation 29 equation
32
equation 37 and equation 39 the total objective function is:
+ Tks + +Jas) const
TTOTAL=
max min Ege
an agus
max min Bog log po(o;) St) log ps(ri st)
(40)
Σ B: KL(q0, Px.) d1 (la - lay) ) + d2 (lay la,) )} const.
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, 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::
To in pa (0, 8t), T2 In po 1s), Tn =-=KLL pt)).
JAS = di ,I,, (Rri Sy, a, 8,, JAs = d2 (1, (a-1; s sg((.-1) ))- - lo. (a,11 s, sg(S 11))
sg(( - Ing (Rri S/. a. sg( (/(55/1)))
(41)
and the KL divergence term can be further decomposed into components:
Jb = -Bi KL(90 (S," (,, 1,11 Pr (5."
Tb. = - KL(qs, (3 10,0 Pyz (sp" (())))
(42)
Th = -B:- KLLg0s (5." 1,, S1,10-1) |1 Pys (s (t-1, a-1)
Jb. = -B4 K.-KLIqu (s 10, P:4 (5 st-1)).
Specifically, maximizing lay in Tks enhances the predictability of R, based the current state x'_ conditioning on (5"_1+4-1:x). However, notice that this objective can be partially accomplished by optimizing Tg. When learning the world model, both the transition function and the reward function are trained: the reward function predicts the current reward r, using sy, while the transition model predicts the next state. These combined predictions contribute the overall prediction of R,
Minimizing log in Tks eliminates extraneous reward-related information present in sy According to our formulation, s/ can still be predictive of R, as long as it does not introduce additional predictability beyond what is already captured by (s" 4111). This is because we only assume that s is condition- ally independent from R, when conditioning
(s"_)= (4-1z). If we don't condition on (s/_)- (5/-j= 4z-1z), introduces 5"=_ the confounding factor between sy and R, establishing association between si and R, (refer to Figure Figure[[5(a) [S(a)). Note that the KL divergence constraints govern the information amount within each state category. By amplifying the weight of the KL constraints s/ the value of loz can indirectly be diminished.
By maximizing la, and minimizing la, in T AS we ensure that s can be predicted based on (a,11, St-1) while s² cannot. The KL constraints on s," and s/d incorporate the action a,- into the prior and posterior, implicitly requiring that x should be predictable given 4-1- Conversely, the constraints on xhr and s/² do not include the action 4-1 in the prior and posterior, implicitly requiring that of should not be predictable based on a,11 However, relying solely on indirect constraints can sometimes be ineffective, it may lead entangled representations that negatively impact policy performance (see Figure [3]
Ablation the mutual information constraints. The inclusion of both Jkks and T is essential in certain environments to promote disentanglement and enhance policy performance, despite sharing
22

---

Reacher-Video-Caaeee ReccherrVider--Cnnerraaiittnnn
800
Factor wio
DenelsedMDP
Episode Return
800
400
200
0.0
0.2
04
0/
0.5
1.0
Environment Steps
Figure 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 million environment steps.
some common objectives. We have observed improved training stability the variant of Robodesk enviromment (see Figure and significant performance gains in the Reacher environment with video background and camera jittering (see Figure 7. When two mutual information objectives we notice that entangled representations emerged in these environments, as depicted in Figure [3 We assign values of 1 to d and d in the environment of modified Cartpole, variant of Robodesk and Reacher with video background and jittering camera. Empirically, value of 0.1 has been found to be preferable for both d and dz. 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 (21,B2,8s.Ba) depends on the specific goals of representation learning and extent of noise interference 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, is recommended to set the values of ß1 and f2 higher than B3 and B4 (for experiments 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.
removed,
C Algorithm
Algorithm 1: IFactor
Input: Representation model: (s(s) 8.1.4.1.00 9: 4oz 94s qui
Transition model: (y(s; 5.11 a.1) Pri Pn Pr. Pyc
Observation Model: Ps (a, sy): Reward model: p.(r, s
Policy Model: (a, s/); Value model: vg(s;):
Mutual information estimator: le, lay. Auii
Policy optimization algorithm Pt-OPr, which is in default the same as that used in Dreamer.
Output:
Transition Model Py. Representation Model 4o. Policy Model ny.
I: while training do
2: EExploration
a, r;), with æ acting on qu encoded outputs 4: Add experience to the replay buffer D DUDD
Collect new trajectories
((,,,,,
5:
// Model learning
6: Sample batch of data sequences ((o,, r;),, from the reply buffer D
Obtain 5, by the representation model and estimate mutual information terms lo Aun lang n do
9:
Train mutual information neural estimators by maximizing lay loys lozs los end for
10:
11:
Freeze the parameters in la. loy lay lo, and Train qo Py and pe with Equation
12:
Policy learning by dreaming
13:
Imagine trajectories of s; using the learned world model.
14:
Train n. and by running Pi-OrT
15:
end while
23

---

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 r 2 in the synthetic environment. When r = 1, it reduces common MDP. The ground-truth generative model of the environment follows::
Pyi Si I 5-ra-1:------
Observation Model:
(0; st)
Pra stir 1 (
Reward Model:
Po (ri sy)
Transition
Transition Model:
(St St-rt-1, a-r1-1)
Pys (s St-rtt- 10-rtt-1-4--1--1 a1-r1-1)
Py: (5) St-rt11)
(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 [ The latent variables St have 8 dimensions, where st" . s S 2. At each timestep, a one-hot action of dimension 5, denoted as a, is taken. The lag number of the process is set to The observation model Pa (o,, st) is implemented using a random three-layer MLPs with LeakyRel.U units. The reward model po (r:-l s,) is represented by random one-layer MIP. It's worth noting that the reward model is not invertible due to the scalar nature of r, Four distinct transition functions, namely Py. Pys, Pyy. and Px.. are employed and modeled using random one-layer MLP with LeakyRel.U units. The process noise is sampled from an i.i.d. Gaussian distribution with standard deviation of o 0.1. simulate nonstationary noise for various latent variables in RI,, 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 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 inreeinteeevantt cartpole undergoes random actions 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 large TV that displays natural RGB videos. On the desk, there is green button that controls both the hue of the TV and a light. The agent's 's objective is to manipulate this button in order to change the TV's hue green. The agent's reward determined based on the greenness of the TV image. In this environment, all four types of information are present (see Table I)).
Modified Cartpole
Robodesk
Reacher Easy
Cheetah Run
Walker Walk
Figure Visualization of the environments used in our experiments.
24

---

Ctrl + Rew
Ctrl Rew
Ctrl + Rew
Ctrl + Rew
Modified Cartpole
Agent
Agent
Green Light
Distractor cartpole
Jittering and flickering environment lighting.
Blocks desk,
TV content,
Agent, Button. Light on desk
Handle on desk,
Robodesk
Button sensor noise
Other movable objects
(Green hue of TV)
Jittering camera
Noiseless
Agent
(Agent)
Video Background
Agent
(Agent)
Background
Video Background
(Agent)
DMC
+ Noisy Sensor
Agent
Background
Video Background
Background Jittering camera
Agent
(Agent)
Camera Jittering
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 De- noised MDP that assumes independent latent processes, IFactor allows for causally-related processes. Therefore, in this paper, the term "controllable" refers specifically one-step controllability, while "reward-relevant" is characterized by the conditional dependence between s and the cumulative reward variable R, when conditioning on (s,-1,Q:11s). 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 ceward-irreleve factors, as they are independent of the reward given the state of the robot arm and green button, aligning with the definition of spa
D.4
Variants of DeepMind Control Suite
Four variants p introduced for each DMC task:
Noiseless: Original environment without distractors.
Video Background: Replacing noiseless background with natural videos (37 Ctrll Rew).
Video Background Sensor Noise: Imperfect sensors sensitive to intensity of a background patch (Ctrl Rew).
Video Background + Camera Jittering: Shifting the observation by a smooth random
walk (Ctrl Rew).
The video background in the environment incorporates grayscale videos from Kinetics-400, where pixels with high blue channel values replaced. Camera jittering is introduced through smooth random walk shift using Gaussian perturrbngg acceleration, velocity decay, and pulling force. Sensor noise is added by perturbing specific sensor based on the intensity of a patch in the background video. The perturbation involves adding the average patch value minus 0.5. Different sensors are perturbed for different environments. These sensor values undergo non-linear transformations. primarily piece-wiss wise linear, to compute rewards. While the additive reward noise model may not capture sensor behavior perfectly, it is generally sufficient as long as the values remain within moderate ranges and stay within one linear region. (Note: the variants of Robodesk and DMC not the contributions of this paper. We kindly refer readers to the paper of Denoised MDP 9 for more detailed introduction.)
25

---

E Experimental Details
Computing Hardware We used a machine with the following CPU specifications: Intel(R) (e) Silver 4110 CPU @ 10GGHZ 32 CPUs, eight physical cores per CPU, a total of 256 logical CPU units. The machine has two GeForce RTX 2080 Ti GPUs with IIGB GPU memory.
Reproducibility We've included the code for the framework and all experiments in the supplement. We plan to release our code under the MIT License after the paper review period.
E.1 Synthetic Dataset
Hyperparameter Selection and Network Structure We adopt similar experimental setup to TDRL. [19), while extending it by decomposing the dynamics into four causally related latent processes proposed in this paper (refer Equation 43). For all experiments,
assign β1 BB 33 = B4 = 0.003 as the weights for the KL divergence terms. In this particular experiment, we set d1 and 22 to 0 because the utilization of the ELBO loss alone has effectively maximized J and J as illustrated in Figure 10 Here, Js represents la L.., and J' represents la, - la. The network structure employed in this experiment is presented in Table 2
Training Details The models are implemented in PyTorch 1.13.1. The VAE network is trained using A optimizer for 100 epochs. learning rate of 0.001 and a mini-batch size of 64 are used. We have used three random seeds in each experiment and reported the mean performance with standard deviation averaged across random seeds.
E.1.
The identifiability socres for baselines
0.3430
R² 0.3700
R
0,2274
0.2643
--
034
0.12
0.43
0.12
1-
0.32
9.22
0.38
0.32
1--
0.04
9.16
0.36
0.34
5-
0.05
0.34
0.14
0.25
Estimatted latents
Estimatted Estimattedlattn latents
0.2
0.78
0.19
0.12
0.46
0.32
0.8
0.28
0.07
True lalent
0.14
025
0.30
0.26
0.03
True
0.5
0.27
5.1
0.35
0:23
0.38
0.45
*
0.17
0.18
0.38
0:45
0.17
9.21
0.13
0.47
0.27
0:38
0.35
0.24
0.22
0:12
0.28
0.31
latents
Estimated latants
latents
Estimated latents
(a) DenoisedMDP
DenoisedMDP
(c) Beta
(d) VAE
0.1599
0.1947
R 0.1930
R 0.2122
a
00
0.27
0.36
0.15
8-
0.22
9.22
0.17
0.31
0.08
9.16
0.50
0.08
3-
0.14
0.17
0.15
0.08
Estinatted latents
Estimatted latents
028
0.11
0.33
0.17
0.15
0.34
0.12
0.10
True stentt
0.34
0.4
0.6
0.28
0.14
True
033
0.27
0.33
0.15
0.14
0.43
0.35
*
025
0.32
0.33
011
025
2.26
0.35
0.18
0.07
0.33
0.28
0.15
0.2
0.19
0.15
0.32
True Iatents
Estimated latents
True latents
Estimated latents
(e) FactorVAE
(f) FactorVAF
(g)
PCL.
(h) PCT.
R² = 0.2136
R² 0.2500
R² = 0.2459
R 0.2508
*
D41
0.34
0.37
0.15
0.36
3:88
0.13
0.32
>:
0.36
5.67
0.21
0.05
0.39
0:17
0.3
0.23
Estimated latents
Estimeted latents
3.06
0.13
0.51
0.28
9.59
0.32
0.21
0.11
0.37
0.52
0.21
0.28
Thg 1onii
0.13
0:33
True
0.6
0.23
0.07
0.09
0.40
0.22
0.2
0.13
0.18
0.14
0.02
0.33
0.28
9:33
0.43
0.32
*
0.24
0.2
0.15
0.22
0.05
0.5
0.39
0.22
latents
Estimated latants
latents
Estimated latents
11) VIOWVAE
SlowVAE
(k) TDRL
(1) TDRL
Figure 9: The identifiability socres for baselines in the experiments on synthetic data set. It can be observed that all baselines do not enjoy the property of block-wise identifiability.
26

---

Table 2: Architecture details. BS: batch size, T: length of time series, o_dim: observation dimension, s_ dim: latent dimension, sy dim: latent dimension for sy ser dim: latent dimension for siar St dim: latent dimension for s4 sy" dim: latent dimension for s," (s dim sf" dim + s." _dm + sp _imm + sf" __dmm LeakyReLU: Leaky Rectified Linear Unit.
Configuration
Description
Output
MLP-Obs-Encooee
Obscrvation Encoder for Synthetic Data
Input: Our
Observed time series
BS x T x o o_dim
Dense
128 neurons, LeakyReLU 128 neurons, LeakyReLLU 128 neurons, LeakyRel.U
BS x T x 128
Dense
x x 128
Dense
BS x T x 128
Dense
Temporal embeddings
BS X T x s_dim
2. ML.P-Obs-Decoder
Observation Decoder for Synthetic Data
Input: $1rr
Sampled latent variables
BS x T x dim
Dense
128 neurons, LeakyReLU 128 neurons, LeakyReLU 0. dim neurons, reconstructed 6
BS x T x 128
Dense
x 128
Dense
BS x T x O dim
3. MLP--ewarddDecodee
Reward Decoder for Synthetic Data
Input: Sur
Sampled latent variables neurons, LeakyReLU
BS x T x s_dim
Dense
BS x Tx
4. Disentangled Prior for sf"
Nonlincar Transition Prior Network
Input Dense
Sampled latents and "rtionns( sp dim neurons, prior output
x X( <__im + a_dim)
BS x T x x "_dim
5. Disentangled Prior for sf"
Nonlinear Transition Prior Network
Input
Sampled latent variable sequence si
BS x T sf dim
Dense
o dim neurons, prior output
BS x T
_dim
6. Disentangled Prior for s/
Nonlinear Transition Prior Network
Input Dense
Sampled latents and actions S1T 41:
BS x T x (s dim + a dim)
of "ddimm neurons, prior output
BS x T sf __dim
7. Disentangled Prior for x2P
Nonlinear Transition Prior Network
Input Dense
Sampled latent variable sequence S1T
BS x T x dim
z __diim neurons, prior output
BS Tx 3 _dim
E.1.2
Extra Results.
During the training process, record the estimation value of four mutual information (MI) terms. The corresponding results are presented in Figure I0) Despite not being explicitly incorporated into the objective function, the terms l, - laz and /a To, exhibit significant maximization. Furthermore, the estimation values of lu and la found to be close 0. These findings indicate that the state variable sf contains little information about the reward, and the predictability of se by the action is also low.
E.2 Modified Cartpole
In the modified Cartpole environment, we configure the values as follows: B1 B2 B3 B4 = 0.1 and A; = 22 = 0.1. Recurrent State Space Model (RSSM) uses deterministic part and a stochastic part represent latent variables. The deterministic state size for four dynamies are set to be (15, 15. 15, 15), and the stochastic state size are set to be (2, 2, 1, 4). The architecture of the encoder and decoder for observation is shown in Table and and Table (64 64 resolution). Reward model uses 3-layer MLPs with hidden size to be 100 and four mutual information neural estimators are 4-layer MLPs with hidden size to be 128. Unlike the synthetic dataset, where there are clear categories
27

---

Essimated MIl Value for
Estimated MI Value for
Essimated MI Value for
Ectimated MI Value for
0.00
0.00
22
0044
Estimation
0.00
21
-000
Training Stepe
Training Steps
Training Steps
Traning Steps
Estmation Valua
Training
Trainin Stops
Figure 10: Estimation of the value of four mutual information terms and their differences in experi- ments synthetic data.
for latent state variables, real-world situations pose a challenge due to the potential ambiguity in categorizing true variables. In our modified Cartpole environment, defined the ground-truth s!" as the cart's s position, s as the pole's angle, s." as the light's ' greenness and s as the distractor Cartpole's s state (including cart position and pole angle). The disentanglement scores for individual sf" and sy", as well as the combined sf, are shown in Figure latent variables can be clearly identified.
We can obviously see that the true
Input Shape
Kernel Size
Operator
Stride
Padding
Input
96, 96]
Conv. + ReLU
[32, 47]
4
2
0
Conv. + ReLU
[64, 22, 22]
4
2
0
Conv. + ReLU
[128, 10, 10]
4
2
0
Conv. + ReLU
[256, 4, 4]
4
2
0
Conv. + ReLU
[256, 2,
3
1
0
Reshape FC
[1024]
Table 3: The encoder architecture designed for observation resolution of (96 x 96). Its output is then fed into other networks for posterior inference. The default activation function used in the network is RELU. For observations with a resolution (64 x 64), the last convolutional layer(*) removed.
R
0.8544
0.8960
*
0.74
9:53
-8.88
0.05
0.91
9.36
0.
0.29
Estimated latents
0.93
-0.01
0.04
0.97
-0.03
0.38
Estimated latents
True latents
0:00
0.50
0.05
0.8
0.00
0.47
Av-
-0.01
0.79
.0.00
0.18
0.79
.0.00
0.01
0.00
0.13
0.00
0.85
.0.02
-0.00
0.93
D144
0.12
0.00
0:5
0.03
08
40.00
20
2.41
True latents
Estimated latents
True latents
Estimated latents (d) Cartpolc
(a) Cartpole
(b) Cartpole
(c) Cartpole
Figure 11: The identifiability socres for [Factor in the experiments cartpole.
28

---

Input Shape
Kernel Size
Operator
Stride
Padding
Input
[input_size]
FC ReLU + Reshape
[1024, 1,
Conv. Transpose ReLU
[138,3.3]
Conv. Transpose ReI.U
[128, 9,9]
5
Conv. Transpose ReLU
[64, 21,211
5
2
Conv. Transpose ReLU
[32, 46, 46]
Conv. Transpose ReLU
[3, 96, 96]
Table 4: The decoder architecture designed for (96 x 96)rrsolltion observation. For (64 x 64)-
resolution observation, the first transpose convolutional layer(**) removed.
E.3 Variant of Robodesk
In the variant of Robodesk. we conduct experiments with the following hyperparameter settings:
Bi = B 2, B: = B4 = 0.25, and d = d = 0.1. For the four dynamics, we set the deterministic state sizes to (120, 40, 40, 40), and the stochastic state sizes (30, 10, 10, 10). Denoised MDP utilizes two latent processes with deterministic state sizes [120, 120] and stochastic state sizes [20, 10). For the mutual information neural estimators, we employ 4-layer MLPs with hidden size of 128. To ensure fair comparison, we align the remaining hyperparameters and network structure with those in the Denoised MDP, We reproduce the results of the Denoised MDP using their released code, maintaining consistency with their paper employing the default hyperparameters. In order to evaluate the impact of the Mutual Information (MI) constraints, we conduct an ablation study. The results are shown is Figure I2 The constraints Tis and Tss are observed stabilize the training process of IFactor. The results of IFacoor are areaveraged over 5 runs, while the results of Denoised MDP and IFccor without MI are averaged over three runs.
Robodesk with Various Distractors
800
Fontor
Facorabb 51
Durcoud
700
Episode Retur
800
400
300
0.0
0.2
0.4
0.6
0.8
1.0
Environment Steps
Figure 12: Comparison between IFactor and Denoised MDP in the variant of Robodesk environment.
Policy learning based on the learned representations by IFactor We retrain policies using the Soft Actor-Criti Critic algorithm [40] with various combinations of the four learned latent categories as input. We wrap the original environment with visual output using our representation model to obtain compact features. In this process, both deterministic states and stochastic states are utilized form the feature. For instance, when referring S, we use both the deterministic states and stochastic states of s/. The implementation of SAC algorithm is based on with learning rate of 0.0002. Both the policy network and network consist of 4-laer ML.Ps with hidden size of 256. We use the default hyperparameter settings Stable-Basclines" Basclines3 for other parameters.
29

---

E.4 Variants of Deep Mind Control Suite
In the noiseless DMC environments, we set BI = B: =
=. For the DMC environments
with video background, we set Bi = fB = and B3 0 0.25. In the DMC environments with
video background and noisy sensor, we set Bi = = 5==2 and B3 = =B: = 0.25. Lastly, for the DMC
environments with video background and jittering camera, we Se = = = set = A2 0.1 for our experiments. For the other environments, we set A₁ = 12 = O. The deterministic state sizes for set to (120, 120, 60, 60), while the stochastic state sizes set to (20, 20, 10, 10). The four mutual information neural estimators utilize a 4-layer MI.Ps with a hidden size of 128. We align the other hyperparameters and network structure with those used in the Denoised MDP for fair comparison.
Regarding the Reacher environment with video background and jittering camera, and ß3 = =B. = 0.25.
the four dynamics
Environment Steps
Action Repeat Every
Train
Collection Intervals
Batch Sequence Size
Horizon
Length
Modified Cartpole
200,000
1
5
5
20
30
8
Robodesk 1,000,000
2
1000
100
50
50
15
DMC 1,000,000
2
1000
100
25
50
12
Table 5: Some hyperparameters of our method in the environment of Modified Cartpole, Robodesk and DMC. Environment Steps represents the number of interactions between the agent and the environment. Action Repeal determines how many times an agent repeats action in a step. Train Every specifies the environment step between adjacent training iterations. Collection Intervals defines the number of times the model is trained in each training iteration (including world models, policy networks and value networks). Batch Size refers to the number of trajectories in each mini-batch. Sequence Length denotes the length of the chuck used in training the world models. Horizon determines the length of dreaming when training the policy using the world model. Hyperparameters aligned with those used in the Denoised MDP for fair comparison.
E.4.1 Policy optimization results on variants of DMC
We present detailed evaluation results in Table showcasing both the mean values and the corre- sponding standard deviations for the final policy performance across each task. Results are averaged across three seeds. Denoised MDP performs well across all four variants with distinct noise types.
E.4.2 Mutual information
Figure 7 demonstrates the notable improvement in policy performance in the Reacher environment with video background and jittering camera due to the inclusion of the constraints Jhs and J'As To further investigate how they affects the model learning, we record the estimation values of four Mutual Information terms throughout the training process, as depicted in Figure 13 The results indicate that both L,, laz and las la, are maximized for both IFactor and IFactor without MI. However, IFactor exhibits significantly higher rate of maximizing la, lus compared to IFactor without MI. This increased maximization leads greater predictability of sf by the action, ultimately contributing to the observed performance gain.
E.5 Visualization for DMC
In this experiment, we investigate five types of representations, which can be derived from the combination of four original disentangled representation categories. Specifically, sf is the controllable and reward relevant representation. sf (s." s(x)) is the reward-relevant representation. s." is the controllable but reward-irrelene representation (noise). x/ (sahh , is the reward-irrelevant representation. Only representations of s, are used for policy optimization. We retrain 5 extra observation decoders to reconstruct the original image. which can precisely characterize what kind of information each type representation contains, surpassing the limitations of the original decoder that is used in latent traversal. The visualization results are shown in Figure 14] can be observed that si" captures the movement of the agent partially but not well enough; s captures the movement of the agent precisely but sf fails (Reacher and Cheetah) captures extra information of the background (Walker). This finding
representation. sh is the uncontrollable and reward-irreleant
30

---

aactor
Dreamer
Denoised
Cheetah
Dreamer
TIA
DBC
CURL
PI-SAC
(Ours)
Pro
MDP
Noiseless
874+339
803+75
771=60
714+348
766+29
182+50
171+34
369428
Video Background
573:1993
366+96
431+111
171#45
388+297
248+71
70+48
112+31
Video Background Noisy Sensor
456+88
134+41
400+190
166+42
227+13
141#13
199=7
151+14
Video Background + Camera Jittering
418+22
150+ 104 104+100
160+32
202+93
141+35
169+22 156+20
Faacor (Ours)
Dreamer
Denoised MDP
Walker
Dreamer
TIA
DBC
CURL
PI-SAC
Pro
Noiseless
966+5
941#16
947=13
955#6
955+5
614+111
417+296
203+92
Video Background
917+52
909+48
790+113 13 247+135
685+337
199+67
608+100
200+18
Video Background Noisy Sensor
701:174
242+65
661+120
279+145
425+281
95+54
338+92
222+21
Video Background + Camera Jittering
524+194 368+301 101 106+22 230+332
62+17
448+70
116#6
Factor (Ours)
Dreamer
Denoised MDP
Reacher
Dreamer
TIA
DBC
CURL
PI-SAC
Pro
Noiseless
924+337
924461
686+216
876+57
587+256
95458
663+221
166+235
Video Background
963+10
555+92
544=121
253=127
123+21
102+58
151
76+35
Video Background Noisy Sensor
839+51 675=137
561=182
202+82
264+280
97+39
607+260
85:5
Video Background + Camera Jittering
736+53 675+82 213=106 109=19 89+26
87+51
632+96
84+13
Table 6: Policy optimization evaluation on variants of DMC with various distractors. IFactor consistently performs well across all four variants with distinct noise types. Bold numbers show the best model-learning result for specific policy learning choices. Results are averaged over 3 runs.
suggests that sf contains sufficient information within the original noisy observation for effective control, while effectively excluding other sources of noise.
31

---

Estiatedd MI Value for lay
Estimated Ml Value for lr
Estimated MI Value for la
Estimated MI Value for
1.50
0.25
Estraation Vakie
1:23
0.20
1:92
0.75
0.15
0.00
0.00
0.00
Training Steps
Training Slepe
Training Steps
Trahiing Sleps
Facter
1:5
0.5
Estmation
0.
021
2.00
-13
02
9.4
3.
0.4
0.5
Traiinn Stops
Training Stops
Figure 13: Estimation of the value of four mutual information terms and their differences in the Reacher Easy environment with video background and jittering camera.
Reacher Dasy Video Background
Cheetah Run
Video Background + Jittering Camera
Env. Rollout S
Eny. Rollout st
sf
s
@F
s
Walker Walk Video Background Noisy Sensor
Env. Rollout
saa
sf
sp
aa
sf
Figure 14: Visualization of the DMC variants and the factorization learned by IFactor.
32

---

Comparison between Iaccor and Related Work
F.
Comparison between IFactor and Denoised MDP
While both IFactor and Denoised MDP share the common aspect factorizing latent variables based on controllability and reward relevance, it is crucial to recognize the numerous fundamental distinctions between them.
(a) IFactor (ours).
(b) Denoised MDP.
Figure 15: Graphical illustration of our world model and Denoised MDP. First and foremost, Denoised MDP adopts rather stringent assumption by solely considering three types of latent variables and assuming independent latent processes for x, (controllable and reward-relevant) and (uncontrollable but reward-relevane However, this strict assumption may hold in many scenarios where uncontrollable yet reward-relevant factors exhibit dependencies on controllable and reward-relevant factors. Take, for instance, the case of car driving: the agent lacks control over surrounding vehicles, yet their behaviors are indirectly influenced by the agent's s actions. In contrast, our approach encompasses four types of causally related latent variables while only assuming conditional independence when conditioning on the state in the previous time step. This assumption holds true naturally within the (PO)MDP framework.
Secondly, Denoised MDP is limited to only factoring out additive rewards ry, from Tx,, disregarding the possibility of non-additive effects in many uncontrollable yet reward-relevant factors. In contrast, our method embraces the inclusion mon-additive effects of x,"r on the reward, which is more general.
Thirdly, Denoised MDP uses only controllable and reward-relevar latent variables x, for policy optimization, which we show in the theoretical analysis that it is generally insufficient. In contrast, our method utilize both controllable and uncontrollable reward-relevant factors for policy training.
Finally, Denoised MDP makes the assumption of an instantaneous causal effect from X; to 2 and from 2, which is inherently unidentifiable without further intervention. It is worth noting that imposing interventions on the latent states is unrealistic most control tasks, as agents can only choose actions at specific states and cannot directly intervene on the state itself. In contrast, our method assumes that there exists no instantancous causal effect for latent variables. In conjunction with several weak assumptions, we provide a proof of block-wise identifiability for our four categories of latent variables. This property serves two important purposes: (1) it ensures the removal of reward- -irrelevant factors and the utilization of minimal and sufficient reward-relevam variables for policy optimization, and (2) it provides potential means for humans to comprehend the learned representations within the reinforcement learning (RL) framework. Through latent traversal of the four types of latent variables, humans can gain insights into the specific kind of information that each category of representation contains within the image.
From the perspective of model structure, it is worth highlighting that the architecture of both the transition model (prior) and the representation model (posterior) in IFactor differs from that of Denoised MDP. The structure of prior and posterior of IFactr is shown follows:
Prior:
Posterior:
Pr sar Sf s/-1.
g
sar Or, ((-1
Pyz ar s/-1)
S 0,,5 s/1)
(44)
Pr (saF (t-1,
(s o,, Sp-1. Sp-1.1111)
Py, (str St-1)
4s. (s Sp11)
While Denoised MDP has the following prior and posterior:
33

---

Prior:
Posterior:
Pyi (Xy 3--1,0--1)
Poi (x, Nr
(45)
Pr: (y Y--1)
Pe: Oyl 2-11)y-1-2-11100; a-11)
Py (2/l 2-11)
Ps. (Z 0,,
notable distinction can be observed between Denoised MDP and IFactor in terms of the assumptions made for the prior and posterior structures. Denoised MDP assumes independent priors for x, and yr, whereas IFactor only incorporates conditional independence, utilizing 5/1 as input for the transition of both s/" and s.", Moreover, the posterior of y, receives a,11 as input, potentially implying controllability. Similarly, the posterior of incorporates 2r-1 input, which may introduce noise from 2 into x,- These implementation details can deviate from the original concept. In contrast, our implementation ensures consistency between the prior and posterior, facilitating a clean disentanglement in our factored model.
From the perspective of the objective function, IFactor incorporates supplementary mutual information constraints, namely Jks and J's to promote disentanglement and improve policy performance.
F.2 Comparison between IFactor, InfoPower and IsoDream
Iaccor learns different categories of state representations according their relation with action and reward, which is different from InfoPower [48] and IsoDream [[2], and moreover, IFactor emphasizes block-wise identifiability for the four categories of representations while InfoPower and do not. Specifically, InfoPower learns types of latent variables, including s Sf and s IsoDream uses three branches for latent dynamics, distinguishing controllable, noncontrollable, and static parts.
Other differences with InfoPower:
Reconstruction Basis: InfoPower is reconstruction-ffre, while IFactor is reconstruction- based.
Objective Functions: InfoPower prioritizes mutual information and empowerment, while IFactor utilizes reconstruction, KL constraints, and mutual information constraints for disentanglement. InfoPower formulates policy using task reward value estimates and the empowerment objective, while IFactor learns policy by maximizing the estimated value and dynamics backpropagating.
Other differences with IsoDream:
Objective Functions and Dynamics Modeling: IsoDream models controllable transitions using inverse dynamics, while IFactor disentangles with multiple mutual information and KL divergence constraints. IsoDream learns policy using future-state attention mechanism rooted in present controllable and future uncontrollable states. contrast, IFactor focuses reward-relevant states, ensuring sf variables are optimal for policy optimization.
34


Reviews:
[{"review_id": "ohgOYvw4bN", "reviewer": "Reviewer_gMe8", "summary": "The authors propose an alternative way to create world models in an RL system by separating them into blocks dependant on their casual effects on future observations, states and rewards. They theoretically show that under some assumptions it is possible to identify such classes of variables, even if specific instances cannot be identified. They go on to show how their approach can increase robustness and disentanglement in several control tasks.", "questions": "In general my questions revolve around the quality of the plots and some results that I think are missing.\n\n1. I would like to sees disentanglement scores for the comparison models in Figure 3. While it is true that these do not posses explicit block level disentanglement, this could potentially emerge from training. I personally don't believe this is likely, but it is still important that the authors show that it is not the case.\n2. Why are there no disentanglement scores for the cart pole dataset? The plot is very small too. This needs must be larger otherwise it is not easy to read. Figure 5 has the same issue.", "limitations": "1. Section 2 could use a slight rewrite. Each of the propositions in lines 98 to 113 start the same way. Seems like it should be possible to just use one heading and enumerate the propositions afterward.\n2. I would prefer if the authors gave more intuitions in section 3 and left the theorem for the appendix. I find that when reading articles it's much easier to reason about the approaches if a good intuition is given than if I am presented with a technical derivation of a theorem. Speaking of which it is not clear why property A2 in Theorem 1 is needed. Also, I cannot find what the definition of $\\hat{s}$ in the text until it is used in Definition 1.\n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is well motivated. I especially liked that instead of focusing on variable level disentanglement the authors have opted for block level disentanglement. This is a good way to relax the assumptions of the model without being completely unconstrained.\n2. The theoretical motivation seems good (but see below). It is always a bonus when researchers can provide some theoretical justification for their work.\n3. Experimental evidence is convincing. I especially liked the analysis of the disentanglement of structured datasets and the results look extremely promising.", "weaknesses": "No serious weakness. I have some comments regarding how section 2 and 3 are written but otherwise it is a solid paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Dear Reviewer,**\n\nThank you for your detailed feedback on our paper. Your helpful comments have helped us further improve the paper. We've worked hard to deal with all the things you pointed out, and here, we explain the changes we made based on your advice.\n\n**Q1: Disentanglement scores for comparison models in Figure 3**\n\n*Response:* Thanks for your comments. To clarify, we provided averaged disentanglement (or, as we prefer, 'identifiability') scores ($R^2$) in the right two panels of Figure 3, which are derived as the mean of the four diagonal values. This was intended to show how the identifiability scores change during training. Following your feedback, we have now computed and appended the identifiability scores for the optimal model iteration for all comparison models. You can find the results in the newly uploaded pdf file, and they will also be included in the appendix of the final version.\n\n**Q2: Disentanglement scores for the cart pole dataset and legibility issues in Figures**\n\n*Response:* Thanks for your careful reading. We have included the identifiability scores for the cart pole dataset in Figure 2 of the newly uploaded pdf file. Unlike the synthetic dataset, where there are clear categories for latent state variables, real-world situations pose a challenge due to the potential ambiguity in categorizing true variables. In our modified Cartpole environment, we defined the ground-truth $s^{ar}_t$ as the cart's position, $s^{\\bar{a}r}_t$ as the pole's angle, $s^{a\\bar{r}}_t$ as the light's greenness, and $s^{\\bar{a}\\bar{r}}_t$ as the distractor Cartpole's state (including cart position and pole angle). The disentanglement scores for individual $s^{ar}_t$ and $s^{\\bar{a}r}_t$, as well as the combined $s^r_t$, are shown in Figure 2 of the supplementary file. We can obviously see that the true latent variables can be clearly identified. Additionally, we have enlarged Figures 4 and 5 in the revised version to improve readability.\n\n**Limitations:**\n\n1. **Structure of Propositions**\n\n*Response:* We really appreciate your careful observation regarding the repetitive structure between lines 98 and 113. In line with your suggestion, we will simplify this section by consolidating it under a single heading and providing numerical listings for the propositions. This modification will be present in the paper's final version.\n\n1. **Intuition in Section 3 and Clarifications on Theorem 1**\n\nThanks for your great suggestion. Intuitively, property A2 in Theorem 1 requires that the Jacobian varies “enough” so that it cannot be contained in a proper subspace of R. This requirement is necessary to avoid undesirable situations where the problem becomes ill-posed and is essential for identifiability. A special case when this property does not hold is when the function f is linear, as the Jacobian remains constant in such cases. We will give a more detailed explanation in the main text, accompanied by a more intuitive outline of the proof.\n\nWe sincerely hope our revisions and clarifications align with your expectations. Your invaluable insights have definitely improved the quality of our research."}, {"review_id": "6kcbiTSd6p", "reviewer": "Reviewer_qxcU", "summary": "- The paper introduces a framework called IFactor for modeling latent state variables in reinforcement learning (RL) systems. The framework categorizes these variables into four distinct types based on their interactions with actions and rewards. The paper further establishes block-wise identifiability of these latent variables, which provides a stable and compact representation and discloses that all reward-relevant factors are significant for policy learning. Overall, the paper contributes a comprehensive framework, IFactor, that provides representations of different aspects of information in RL systems and highlights the importance of considering states that influence the reward during decision-making.", "questions": "- I understand that blockwise identification refers to the existence of a one to one mapping between a late variable and its estimated value. However, I lacked a thorough understanding of how decoupling the four latent components is related to the concept of block-wise identifiability. Could you please explain further? Or what are the benefits of defining these components blockwise identifiable? \n- Additionally, if this mapping is not blockwise identifiable, what impact will it have on the decoupling performance and how will it influence the model performance?", "limitations": "Negative societal impact is not discussed", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- IFactor models different latent state variables in based on their interactions with actions and rewards in RL systems. These categories include reward-relevant and controllable parts, reward-irrelevant but uncontrollable parts, controllable but reward-irrelevant parts, and unrelated noise.\n- The paper defines blockwise identifiability as the existence of a one-to-one mapping between a latent variable and its estimated value, such that the estimated value can be recovered from the latent variable and vice versa. This means that each latent variable can be uniquely identified and separated from other variables in the model, allowing for a more stable and compact representation of the environment.\n- The paper includes extensive experiments and ablation studies to validate the effectiveness of the proposed framework and to analyze the impact of different components. Meanwhile, the paper includes a thorough evaluation of the proposed framework on several benchmark tasks.", "weaknesses": "- Missing related work. The separate modeling of controllable and uncontrollable components has been explored in various previous works, such as InfoPower[1] and Iso-Dream[2], which are missing from related work. It would be better if the authors can discuss the differences between the proposed method and this type of work.\n  - [1] INFOrmation Prioritization through EmPOWERment in Visual Model-Based RL. ICLR 2022.\n  - [2] Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models. NeurIPS 2022.\n- The error bars of the baselines are not included.\n- The performance of TIA on variants of DMC sometimes encounters a sudden drop during training, which is worth a further check for the reason.\n- Except for DenoisedMDP, the compared approaches are somewhat \"out-of-date\". It would be nice if the authors could include stronger baseline methods for model comparison, especially the approaches that show impressive performance on generalization in noisy visual observations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your valuable feedback. Please see our responses to your questions point-by-point below.\n\n1. M**issing related work on separate modeling**:\n    \n    Thank you for bringing our attention to the related works, including InfoPower [1] and Iso-Dream [2]. We will make sure to give proper citations and provide comparisons in our appendix for the revised version. A brief discussion about differences is provided below.\n    \n    Main differences: IFactor learns different categories of state representations according to their relation with action and reward, which is different from InfoPower and IsoDream, and moreover, IFactor emphasizes block-wise identifiability for the four categories of representations while InfoPower and Iso-Dream do not. Specifically, InfoPower learns 3 types of latent variables, including  $s^{ar}_t, s^{\\bar{a}r}_t$ and  $s^{\\bar{a}\\bar{r}}_t$. IsoDream uses three branches for latent dynamics, distinguishing controllable, noncontrollable, and static parts. \n    \n    - Other differences with InfoPower:\n        - **Reconstruction Basis**: InfoPower is reconstruction-free, while IFactor is reconstruction-based.\n        - **Objective Functions**:  InfoPower prioritizes mutual information and empowerment, while IFactor utilizes reconstruction, KL constraints, and mutual information constraints for disentanglement. InfoPower formulates policy using task reward value estimates and the empowerment objective, while IFactor learns policy by maximizing the estimated Q value and dynamics backpropagating.\n    - Other differences with IsoDream:\n        - **Objective Functions and Dynamics Modeling**: IsoDream models controllable transitions using inverse dynamics, while IFactor disentangles with multiple mutual information and KL divergence constraints. IsoDream learns policy using a future-state attention mechanism rooted in present controllable and future uncontrollable states. In contrast, IFactor focuses on reward-relevant states, ensuring $s^r_t$ variables are optimal for policy optimization.\n2. **Absence of error bars for baselines**:\n\nThank you for your careful checking. We didn't include error bars for baselines due to the following reasons. Initially, we tried to replicate the results of Denoised MDP [3] using their available code. However, the outcomes we obtained were considerably poorer than those reported in their paper (refer to Figure 4 in the appendix). In order to maintain a fair comparison, we directly incorporated the reported results from Denoised MDP's image plots. The absence of error bars is because we were unable to extract them from these image plots. In light of your feedback, we have made the subsequent changes.\n\n1. We have included error bars for Figure 3 in the newly uploaded pdf file. They will be included in the final version.\n2. We have also listed the mean and standard deviation of the return value for all baselines, when policy optimization gets converged, at the end of the general response (these values are directly copied from the reported results). These tables will also be added to the appendix of the revised manuscript.\n\n1. **Performance inconsistencies of TIA**:\n    \n    To maintain fair comparisons, the results of TIA were directly copied from Denoised MDP. We will replicate and provide detailed explanations of TIA's results in our revised manuscript.\n    \n2. **Comparison with contemporary baselines**:\n    \n    Thank you for the suggestion. Considering the continuous progress in RL, we've added results from DreamerPro, a top model-based RL method that doesn't require reconstruction and can handle noisy visuals effectively, as a baseline. You can find the updated results in Figure 4 of the newly uploaded pdf file and at the end of the general reponse, where IFactor achieves superior performance.\n    \n\n## **Questions:**\n\nQ1: Block-wise Identifiability & Decoupling:\n\nA1: Basically, block-wise identifiability ensures the theoretical guarantee of the estimations of the four categories in relation to their ground truth values. With identifiability, we can ensure asymptotic correctness. Moreover, it is important to note that decoupling alone does not guarantee identifiability. Specifically, solely relying on decoupling may lead to missing reward-relevant information in the estimation $\\hat{s}^r_t$, as well as the presence of redundant information.\n\nQ2: Benefits of block-wise identifiability and consequences of non-blockwise identifiability\n\nA2: Since block-identifiability guarantees the recovery of ground truth variables, it additionally offers the advantages of improved interpretability and facilitating more efficient and effective policy optimization, which are also the downsides of missing identifiability.\n\n- **Interpretability**: It guarantees the recovery of underlying causal latent variables under certain conditions, thereby enhancing the clarity of decision-making processes by distinguishing latent state variables based on controllability and reward relevance.\n- **Policy Optimization**: By focusing on $s^r_t$, IFactor streamlines policy learning, keeping only essential information and discarding redundancy. If the latent variables lack block-wise identifiability, $s^r_t$ may have redundant or insufficient information, hindering efficient and effective policy learning.\n\nWe deeply appreciate the depth and breadth of your review. Your insights are crucial in improving our manuscript. We will integrate your feedback into our revised version.\n\nThe reference index can be found at the end of the general response."}, {"review_id": "h93bcm6Srb", "reviewer": "Reviewer_oFQM", "summary": "Learning efficient world models requires architectural priors to learn better representations that can capture different aspects of the environment. However, existing methods like Dreamer do not focus on learning disentangled representations and lack the ability to separate noise from the reward-relevant information in the observations. In this work, a new method is proposed where the latent state is divided into 4 blocks based on the dependence of action and reward function. The paper shows that using the blocks of states that affect the reward function can be used to learn policies. Experiments show the efficacy of the proposed algorithm on Robodesk and DMControl tasks with distractions and noise in the form of background or based on a sensor/camera. Moreover, ablation studies show that the model learns to disentangle the factors in the environment. Lastly, the paper also provides a theoretical justification for the proposed architecture.", "questions": "1. For learning efficient latent representations, recent works like DreamerPro have used auxiliary losses to extract relevant concepts from the observation [1]. How does the proposed method compare with reconstruction-free model-based methods?\n2. For experiments on control tasks, the noise used is the background noise. There can be other types of noises as shown in [2] where distractor objects are moving in the environment. Will the proposed method be able to disentangle such noises too?  \n3. Were there any experiments conducted on a few games of Atari (like MsPacman, PrivateEye, and Montezuma’s Revenge) and observe what the latent embeddings are learning? On the current task, the whole agent is visible in the observation, what happens when the agent gets a partial view of the environment?\n4. Related work should include [4] as they also plan to learn latent representations that disentangle different objects. Furthermore, how will the proposed method work in such scenarios where different rewards are associated with different objects? \n5. Will the proposed method work in multi-task learning or meta-learning scenarios where the rewards are changing because there is a part of latent space that depends on the reward function which may change with different reward functions?\n6. The current method assumes the presence of a reward function. Will this assumption make it difficult to learn unsupervised world models? Here the latent space cannot be conditioned on the intrinsic reward as it changes with training.\n\n### References\n[1] Deng, Fei, Ingook Jang, and Sungjin Ahn. \"Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations.\" International Conference on Machine Learning. PMLR, 2022.\\\n[2] Jain, Arnav Kumar, et al. \"Learning robust dynamics through variational sparse gating.\" Advances in Neural Information Processing Systems 35 (2022): 1612-1626.\\\n[3] Nguyen, Tung D., et al. \"Temporal predictive coding for model-based planning in latent space.\" International Conference on Machine Learning. PMLR, 2021.\\\n[4] Kipf, Thomas, Elise Van der Pol, and Max Welling. \"Contrastive learning of structured world models.\" arXiv preprint arXiv:1911.12247 (2019).\n", "limitations": "The paper briefly talks about a few limitations between lines 346-350 and needs to elaborate more on the discussion and why they are challenging. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The idea of disentangling the learned representations to different factors based on actions and rewards is interesting and nicely explained in the paper. \n2. The experiments are thorough and support the claims made in the paper that the agent is able to differentiate between parts of the state that determines the reward and parts of the state that change based on the action.\n3. The paper presents a theoretical derivation of the identifiability of the latent state representations.\n", "weaknesses": "1. World Models have been extended to scenarios with noisy distractions in [1][2][3]. They have shown improvements over the Dreamer method for such tasks. However, these methods are not used as baselines in the experiments. \n2. Some ablations and experiments are required to further understand the scenarios where the proposed method works (More on this in Questions below)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Dear Reviewer,**\n\nThank you for your valuable feedback. Below, we address your concerns in a point-by-point manner and have added various experiments, following your suggestions. \n\n**Weaknesses:**\n\n1. Lack recent works in world models such as [4,5,6] as baselines\n    \n    *Response:*  Thank you for bringing our attention to the related works. In our updated experiments, we have included DreamerPro [4] as a baseline for DMC variants. The results, available in the general response and newly uploaded pdf file, consistently show IFactor outperforming DreamerPro. Variational Sparse Gating (VSG, [5]) introduces a distinct latent dynamics model that deviates from the typical RSSM structure. We have not included [5] in our baseline comparisons because IFactor and VSG actually address different aspects, and we are currently working on combining them for potential performance improvement. Moreover, [6] proposes an alternative approach using temporal predictive coding, where latent states may not be identifiable due to the non-invertibility of the mixing function f. We have not included [6] in our baselines because its implementation isn't yet publicly available.\n    \n2. Requirements for experiments to further understand the scenarios where the proposed method works.\n    \n    *Response:*  Thank you for your suggestions. We offer detailed responses to each question below. \n    \n\n**Questions:**\n\n**Q1:**  Comparative Analysis with reconstruction-free model-based methods, such as DreamerPro:\n\n**A1:** We've expanded our experiments to incorporate DreamerPro as a baseline for DMC variants. Results are shown at the end of the general response and in Figure 4 of the newly uploaded pdf file. Notably, IFactor consistently surpasses DreamerPro's performance. \n\n**Q2:** Disentanglement of Noises such as  distractor objects in the environment:\n\n**A2:** Thank you for your insightful thoughts. We have indeed addressed the concept of distractor objects moving within the environment in Section 5.1.2, where we considered a modified Cartpole environment with a distracting cart that is neither reward-relevant nor action controllable. The identifiability score ($R^2$) is given in Figure 2 of the newly uploaded pdf file. Following your suggestions, we are now working on the environment presented in [5].\n\n**Q3:** Experiments on Atari games and Partial-view Environments:\n\n**A3:** Thank you! Following your comments, we are now testing our approach to Atari games. Considering the achievements of Dreamer  (v2/v3) in similar situations and the flexibility of our method to work with different representation types and training techniques, we believe that combining IFactor with Dreamer (v2/v3)  holds the potential. Our visualizations in Figures 4 and 5 of the main paper, and Figures 1 and 6 in the appendix, offer insights into how latent features are organized. \n\n**Q4:**  Inclusion of Related Work [7] and performance of IFactor in scenarios where different rewards are associated with different objects: \n\n**A4:** Thank you for bringing the related work [7] to our attention. We will make sure to include the proper citation in our revised version. Moreover, When dealing with multiple rewards, our proposed framework actually can be directly extended to identify the state representation or object that is associated with each reward signal. To illustrate, suppose there are two reward signals r1 and r2 that are associated with object 1 and object 2, respectively. Then with our framework, we can recover the representation $s^{r1}$ that is relevant to reward r1 and $s^{r2}$ that is relevant to reward r2.\n\n**Q5:** Functionality in Multi-task or Meta-learning:\n\n**A5:** Yes! In fact, our framework can be easily extended to cover heterogeneous/nonstationary environments, with the change of reward function, observation function, or transition dynamics. Basically, we can introduce a (latent) low dimensional change factor to characterize those changes, a similar strategy as AdaRL. This is the direction we plan to take in our next steps.\n\n*Reference:* *[AdaRL] Huang, Feng, Lu, Magliacane, Zhang. AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning. ICLR, 2022.*\n\n**Q6** Application in Unsupervised World Models:\n\n**A6:**  Our framework is versatile when it comes to different supervision signals. Utilizing the reward and action signals in reinforcement learning systems is merely one specific application of its capabilities. In purely unsupervised world models without any supervision signals, we are unable to distinguish different categories of state representations. However, if intrinsic rewards are present, we can learn the state representation relevant to these intrinsic rewards.\n\n**Limitations:**\n\nIn-depth discussion: While recent studies suggest that causal variables can be reconstructed from temporal sequences of observations, assuming no instantaneous causal relations, practical challenges arise. Specifically, if our measurement or frame rate lags behind the speed of causal effects, it can inadvertently introduce “instantaneous” effects, compromising prior identifiability conclusions. Our work's foundational limitation rests on the assumption that latent processes lack instantaneous causal relations. This assumption is evident in Figure 1(c) where no edges exist between concurrent latent states. For a deeper dive, we direct readers to two pivotal papers: [8] introduces iCITRIS, a method allowing instantaneous effects in intervened temporal sequences when intervention targets are observable. Meanwhile, [9] articulates that neglecting instantaneous dependence can lead to subpar policy learning in MBRL. We commit to elaborating on these limitations in the camera-ready edition of our paper.\n\nThank you once again for your rigorous review and constructive feedback. \n\nThe reference index can be found at the end of the general response."}]

Author rebuttal:
Dear Reviewers,

We sincerely thank the reviewers for their effort and helpful comments regarding our paper. We have carefully revised the manuscript according to your comments.  A summary of the primary changes we've made is outlined below:

1. **Identifiability Scores**: We've incorporated Identifiability Scores ($R^2$) for the compared models in the synthetic dataset experiments, as well as Identifiability Scores for IFactor in the modified Cartpole environment. Both can be referenced in the newly uploaded pdf file.
2. **Structural Changes**: Propositions from lines 98 to 113 have been organized under a singular heading, followed by an enumeration for improved clarity.
3. **Theoretical Clarification**: An intuitive elaboration for Assumption 2 within Theorem 1 has been added.
4. **Literature Citations**: Proper attributions to related works, namely InfoPower [1], Iso-Dream [2], DreamerPro [4], VSG [5], TPC [6] and C-SWMs [7] have been integrated. We ensure a thorough comparison with these in the revised manuscript.
5. **Experimental Enhancements**:
    - DreamerPro is now incorporated as an additional baseline in the DMC variants experiments. The results can be found in the tables at the end of this response and in Figure 4 of the newly uploaded pdf file.
    - The middle panel of Figure 5 now displays error bars for baselines.
    - For a comprehensive understanding, three tables elucidating the mean and standard deviation values of the performance of the converged policy in DMC variants experiments have been appended.
6. **Discussion Expansion**: A more in-depth discussion on the limitations between lines 346-350 is now available.

To get a thorough look at these changes, we direct your attention to the newly uploaded pdf file and the specific responses to the reviewers.

We hope our revisions and detailed responses satisfactorily address the concerns raised. Once again, thank you for your generous contribution of time and expertise to the community.

## Evaluation of Policy Performance in DMC Variants
**cheetah**

|  | IFactor (Ours) | DreamerPro | Denoised MDP | Dreamer | TIA | DBC | CURL | PI-SAC |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Noiseless | **874.0±39.0** | 802.5±75.0 | 770.7±59.9 | 713.7±347.8 | 765.9±29.3 | 181.8±49.7 | 171.2±34.2 | 369.4±27.7 |
| Video Background | **572.5±192.5** | 366.0±96.4 | 430.6±111.0 | 171.0±44.9 | 387.8±297.3 | 247.9±71.3 | 70.3±47.6 | 111.9±31.0 |
| Noisy Sensor | **455.5±87.9** | 134.3±41.1 | 400.3±190.0 | 166.0±41.7 | 226.8±13.2 | 140.6±13.0 | 198.9±7.1 | 150.9±14.1 |
| Camera Jittering | **417.7±22.4** | 150.0±104.4 | 293.9±99.6 | 160.1±32.2 | 202.1±93.2 | 141.3±35.4 | 168.6±21.6 | 155.8±19.6 |

**walker**

|  | IFactor (Ours) | DreamerPro | Denoised MDP | Dreamer | TIA | DBC | CURL | PI-SAC |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Noiseless | **966.0±4.5** | 941.4±16.4 | 947.4±13.3 | 954.9±6.0 | 955.3±5.2 | 613.9±110.7 | 416.5±295.8 | 202.9±92.0 |
| Video Background | **916.8±51.7** | 909.1±48.5 | 790.2±113.3 | 247.2±134.6 | 685.4±336.6 | 198.6±67.3 | 607.8±99.7 | 200.3±18.2 |
| Noisy Sensor | **700.5±174.3** | 242.1±64.9 | 660.7±120.1 | 269.6±145.0 | 424.7±281.2 | 95.1±53.7 | 338.4±91.7 | 221.9±20.8 |
| Camera Jittering | **523.7±194.2** | 367.8±300.9 | 290.5±103.9 | 105.9±22.1 | 229.5±331.9 | 62.0±17.3 | 447.5±69.5 | 115.6±5.9 |

**Reacher**

|  | IFactor (Ours) | DreamerPro | Denoised MDP | Dreamer | TIA | DBC | CURL | PI-SAC |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Noiseless | 923.7±37.2 | **924.2±61.2** | 685.7±216.4 | 876.1±57.1 | 587.1±256.3 | 95.2±57.5 | 663.4±220.5 | 166.2±235.4 |
| Video Background | **962.7±9.5** | 555.1±91.5 | 543.9±121.3 | 252.8±127.0 | 123.2±21.3 | 101.6±57.7 | 751.1±188.9 | 76.2±34.6 |
| Noisy Sensor | **839.3±50.9** | 675.1±137.2 | 561.1±182.4 | 201.8±81.6 | 263.7±279.8 | 96.8±38.6 | 606.8±259.7 | 84.7±4.5 |
| Camera Jittering | **735.7±52.6** | 674.5±81.6 | 213.1±105.9 | 108.9±18.9 | 89.3±25.7 | 86.7±50.6 | 631.9±96.0 | 84.3±13.1 |

[1] Bharadhwaj, Homanga, et al. "Information prioritization through empowerment in visual model-based rl.” ICLR 2022.

[2] Pan, Minting, et al. "Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models." *Advances in Neural Information Processing Systems* 35 (2022): 23178-23191.

[3] Wang, Tongzhou, et al. "Denoised mdps: Learning world models better than the world itself." ICML 2022: 22591-22612

[4] Deng, Fei, Ingook Jang, and Sungjin Ahn. "Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations." International Conference on Machine Learning. PMLR, 2022.

[5] Jain, Arnav Kumar, et al. "Learning robust dynamics through variational sparse gating." Advances in Neural Information Processing Systems 35 (2022): 1612-1626.

[6] Nguyen, Tung D., et al. "Temporal predictive coding for model-based planning in latent space." International Conference on Machine Learning. PMLR, 2021.

[7]  Kipf, Thomas, Elise Van der Pol, and Max Welling. "Contrastive learning of structured world models." arXiv preprint arXiv:1911.12247 (2019).

[8] Lippe, Phillip, et al. "Causal representation learning for instantaneous and temporal effects in interactive systems." *The Eleventh International Conference on Learning Representations*. 2022.

[9] Zhu, Zhengmao, et al. "Beware of Instantaneous Dependence in Reinforcement Learning." *arXiv preprint arXiv:2303.05458*(2023).

Meta-review: The authors proposed to learn world models with additional structural assumptions on the causal graph. Specifically, the authors assumed that state variables can be decomposed into four parts with the reward-relevancy and controllability, and discussed when and how these four parts can be identified from the data. Experimental results demonstrate the benefits of such structural assumptions in DMC tasks. Reviewers all agree that the proposed methods are well-motivated and properly justified. I encourage the authors to incorporate the suggestions from reviewers into the final version.
Decision: Accept (poster)