Title: Learning World Models with Identifiable Factorization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Model-based Reinforcement Learning; Causal Representation Learning;
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

Abstract

Extracting a stable and compact representation of the environment is crucial for e ffi cient reinforcement learning in high-dimensional, noisy, and non-stationary environments. Di ff erent categories of information coexist in such environments how to e ff ectively 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 su ffi cient 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.

1 Introduction

Humans excel at extracting various categories of information from complex environments [1, 2]. By e ff ectively distinguishing between task-relevant information and noise, humans can learn e ffi ciently 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.

World models [3, 4] tackle this challenge by learning compact representations from images and modeling dynamics using low-dimensional features. Recent research has demonstrated that learning policies through latent imagination in world models significantly enhances sample e ffi ciency [5, 6, 7, 8]. However, these approaches often treat all information as an undi ff erentiated whole, leaving policies susceptible to irrelevant distractions [9] and lacking transparency in decision-making [10].

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

(a) Car-driving task.

Reward-

Relevant

Reward-

Irrelevant

Controllable

Uncontrollable

Surrounding

vehicles

Remote

scenery

Music and

air conditioner

Speed, position,

and direction

(b) Four types of states.

s

ar

t

-1

s

¯

ar

t

-1

s

a

¯

r

t

-1

s

¯

a

¯

r

t

-1

Ot

-1

s

ar

t

s

¯

ar

t

s

a

¯

r

t

s

¯

a

¯

r

t

Ot

a

t

-1

s

ar

t

+1

s

¯

ar

t

+1

s

a

¯

r

t

+1

s

¯

a

¯

r

t

+1

Ot

+1

r

t

a

t

r

t

+1

Rt

(c) IFactor (Ours).

x

t

-1

y

t

-1

z

t

-1

Ot

-1

x

t

y

t

z

t

Ot

a

t

-1

r

x

t

r

y

t

+

r

t

(d) Denoised MDP.

comprehensive classification system for di ff erent 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 [11] partition the state space into reward-relevant and reward-irrelevant 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 [9] further decomposes reward-relevant 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 [13], the question of identifiability remains unexplored for general nonlinear cases.

In this paper, we present IFactor, a general framework to model four distinct categories of latent state variables within the RL system. These variables capture di ff erent 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) reward-irrelevant but uncontrollable parts, (iii) controllable but reward-irrelevant parts, and (iv) unrelated noise (see Section 2). Diverging from prior methods, our approach employs a 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 reward-relevant variables are required for policy optimization. We emphasize the necessity of considering states that directly or indirectly influence the reward during the decision-making process, irrespective of their 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.

2 Four Categories of (Latent) State Variables in RL

For generality, we consider tasks in the form of Partially Observed Markov Decision Process (POMDP) [15], which is described as M ≜ ( S , A , Ω , R , T , O , γ ), where S is the latent state space, A is the action space, Ω is the observation space, R : S × A → R defines the reward function, T : S × A → ∆ ( S ) is the transition dynamics, O : S × A → ∆ ( Ω ) is the observation function, γ ∈ [0 , 1] is the discount factor. We use ∆ ( S ) to denote the set of all distributions over S . The agent can interact with the environment to get sequences of observations {⟨ ot , at , rt ⟩} T t = 1 . 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-dimensional observations is crucial. Considering that action and reward information is fundamental to decisionmaking across various scenarios, we disentangle the latent variables in the environment into four distinct categories, represented as st = { s ar t , s ¯ ar t , s a ¯ r t , s ¯ a ¯ r t } , based on their relationships with action and reward (see Figure 1(c) as a graphical illustration):

Type 1: s ar t has an incident edge from at -1, and there is a directed path from s ar t to rt .

Type 3: s a ¯ r t has an incident edge from at -1, and there is no directed path from s a ¯ r t to rt .

Type 2: s ¯ ar t has no incident edge from at -1, and there is a directed path from s ¯ ar t to rt .

Type 4: s ¯ a ¯ r t has no incident edge from at -1, and there is no directed path from s ¯ a ¯ r t to rt .

s ar t represents controllable and reward-relevant state variables that are essential in various scenarios. Taking the example of a car driving context, s ar t 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 as steering and braking directly influence these states, which, in turn, have a direct impact on the reward received.

s ¯ ar t refers to reward-relevant state variables that are beyond our control. Despite not being directly controllable, these variables are still necessary for policy learning. In the context of car driving, s ¯ ar t 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 s a ¯ r t consist of controllable but reward-irrelevant factors. Examples of s a ¯ r t could be the choice of music being played or the positioning of ornaments within the car. On the other hand, s ¯ a ¯ r t represents uncontrollable and reward-irrelevant latent variables such as the remote scenery. Both s a ¯ r t and s ¯ a ¯ r t 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 di ff erent types of state variables can be characterized from the data. To simplify symbol notation, we define s r t : = ( s ar t , s ¯ ar t ) , s ¯ r t : = ( s a ¯ r t , s ¯ a ¯ r t ) , s a t : = ( s a ¯ r t , s a ¯ r t ) and s ¯ a t : = ( s ¯ ar t , s ¯ a ¯ r t ). Specifically, the following proposition shows that s r t that has directed paths to rt + τ (for τ > 0), is minimally su ffi cient for policy learning that aims to maximize the future reward and can be characterized by conditional dependence with the cumulative reward variable Rt = P t ′ = t γ t ′ -t rt ′ , which has also been shown in [13].

Proposition 1. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data, s r t ⊆ st is a minimal subset of state dimensions that are su ffi cient for policy learning, and si , t ∈ s r t if and only if si , t ̸ ⊥ ⊥ Rt | at -1: t , s r t -1 .

Moreover, the proposition below shows that s a t , that has a directed edge from at -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 a t ⊆ st is a minimal subset of state dimensions that are su ffi cient for direct control, and si , t ∈ s a t if and only if si , t ̸ ⊥ ⊥ at -1 | st -1 .

Furthermore, based on Proposition 1 and Proposition 2, we can further di ff erentiate s ar t , s ¯ ar t , s a ¯ r t from s r t and s a t , which is given in the following proposition.

Proposition 3. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful 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) si , t ∈ s ar t if and only if si , t ̸ ⊥ ⊥ Rt | at -1: t , s r t -1 and si , t ̸ ⊥ ⊥ at -1 | st -1 , (2) si , t ∈ s ¯ ar t if and only if si , t ̸ ⊥ ⊥ Rt | at -1: t , s r t -1 and si , t ⊥ ⊥ at -1 | st -1 , (3) si , t ∈ s a ¯ r t if and only if si , t ⊥ ⊥ Rt | at -1: t , s r t -1 and si , t ̸ ⊥ ⊥ at -1 | st -1 , and (4) si , t ∈ s ¯ a ¯ r t if and only if si , t ⊥ ⊥ Rt | at -1: t , s r t -1 and si , t ⊥ ⊥ at -1 | st -1 .

By identifying these four categories of latent state variables, we can achieve interpretable input for policies, including (i)reward-relevant and directly controllable parts, (ii) reward-relevant but not controllable parts, (iii) controllable but non-reward-relevant parts, and (iv) unrelated noise. Furthermore, policy training will become more sample e ffi cient and robust to task-irrelevant changes by utilizing only s ar t and s ¯ ar t for policy learning. We show the block-wise identifiability of the four categories of variables in the next section.

3 Identifiability Theory

Identifying causally-related latent variables from observations is particularly challenging, as latent variables are generally not uniquely recoverable [16, 17]. Previous work in causal representation learning and nonlinear-ICA has established identifiability conditions for non-parametric temporally latent causal processes [18, 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-wise identifiability, our focus lies on the block-wise identifiability of the four categories of latent state variables. From a policy optimization standpoint, this is su ffi cient because we do not need the latent variables to be recovered up to permutation and component-wise 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 to the causal process in the RL system (as described in Eq.1 in [13]), we can build the following mapping from latent state variables s t to observed variables ot and future cumulative reward Rt :

where

Here, note that to recover s r t , it is essential to take into account all future rewards rt : T , because any state dimension si , t ∈ s t that has a directed path to the future reward rt + τ , for τ > 0, is involved in s r t . Hence, we consider the mapping from s r t to the future cumulative reward Rt , and η t represents residuals, except s r t , that have an e ff ect to Rt .

Below, we first provide the definition of blockwise identifiability and relevant notations, related to [20, 21, 22, 23].

Definition 1 (Blockwise Identifiability) . A latent variable st is blockwise identifiable if there exists a one-to-one mapping h ( · ) between st and the estimated ˆ st , i.e., ˆ st = h ( st ) .

Notations. We denote by ˜ s t : = ( s r t , s ¯ r t , η t ) and by | s | the dimension of a variable s . We further denote dsr : = | s r t | , ds ¯ r : = | s ¯ r t | , d ˜ s : = | ˜ s t | , do : = | ot | , and dR : = | Rt | . We denote by F the support of Jacobian J f ( s t ), by ˆ F the support of J ˆ f (ˆ s t ), and by T the support of T ( s t ) with J ˆ f (ˆ s t ) = J f ( s t ) T ( s ). We also denote T as a matrix with the same support as T .

In addition, given a subset D⊆{ 1 , · · · , d ˜ s } , the subspace R d ˜ s D is defined as R d ˜ s D : = { z ∈ R d ˜ s | i < D = ⇒ zi = 0 } . In other words, R d ˜ s D refers to the subspace of R d ˜ s indicated by an index set D .

We next show that the di ff erent types of states s ar t , s ¯ ar t , s a ¯ r t , and s ¯ a ¯ r t are blockwise identifiable from observed image variable ot , reward variable rt , and action variable at , 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 1(c). Under the following assumptions

A1. The mapping f in Eq. 1 is smooth and invertible with smooth inverse.

A2. For all i ∈ { 1 , . . . , do + dR } and j ∈ F i , : , there exist { ˜ s ( l ) t } |F i , : | l = 1 , so that span { J f (˜ s ( l ) t ) i , : } |F i , : | l = 1 = R d ˜ s F i , : , and there exists a matrix T with its support identical to that of J -1 ˆ f ( ˆ ˜ s t ) J f (˜ s t ) , so that GLYPH<2> J f (˜ s ( l ) t ) T GLYPH<3> j , : ∈ R d ˜ s ˆ F i , : .

Then, reward-relevant and controllable states s ar t , reward-relevant but not controllable states s ¯ ar t , reward-irrelevant but controllable states s a ¯ r t , and noise s ¯ a ¯ r t , are blockwise identifiable.

In the theorem presented above, Assumption A 1 only assumes the invertibility of function f , while functions f 1 and f 2 are considered general and not necessarily invertible, as that in [23]. Since the function f 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 f 2 since usually, the reward function is not invertible. Intuitively, Assumption A 2 requires that the Jacobian varies 'enough' so that it cannot be contained in a proper subspace of R d ˜ s F i , : . 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. Some proof techniques of Theorem 1 follow from [20, 22, 23], with the detailed proof given in Appendix A.4.

4 World Model with Disentangled Latent Dynamics

Based on the four categories of latent state variables, we formulate a world model with disentangled latent dynamics. Each component of the world model is described as follows:

Specifically, the world model includes three generative models: an observation model, a reward function, and a transition model (prior), and a representation model (posterior). The observation model and reward model are parameterized by θ . The transition model is parameterized by γ and the representation model is parameterized by ϕ . We assume noises are i.i.d for all models. The action at -1 has a direct e ff ect on the latent states st , but not on the perceived signals ot . The perceived signals, ot , are generated from the underlying states st , while signals rt are generated only from reward-relevant latent variables s r t .

According to the graphical model depicted in Figure 1(c), the transition model and the representation model can be further decomposed into four distinct sub-models, according to the four categories of latent state variables, as shown in equation 4.

Specifically, we have p γ = p γ 1 · p γ 2 · p γ 3 · p γ 4 and q ϕ = q ϕ 1 · q ϕ 2 · q ϕ 3 · q ϕ 4 . Note this factorization di ff ers from previous works [11, 9] that assume independent latent processes. Instead, we only assume conditional independence among the four categories of latent variables given st -1 , which provides a more general factorization. In particular, the dynamics of s ar t and s ¯ ar t are dependent on s r t -1 , ensuring that they are una ff ected by any reward-irrelevant latent variables present in s ¯ r t -1 . On the other hand, s ¯ r t may be influenced by all the latent variables from the previous time step. Note that the connections between st -1 and s ¯ r t can be adjusted based on the concrete problem. Since s ar t and s a ¯ r t are controllable variables, their determination also relies on at -1 .

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 [25, 26]. The bound encompasses reconstruction terms for both observations and rewards, along with a KL regularizer:

Further, the KL regularizer J t D 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:

Additionally, we introduce two supplementary objectives to explicitly capture the distinctive characteristics of the four distinct representation categories. Specifically, we characterize the reward-

̂

̂

h ar

t

h ar

t

h ar

t

h ar

t

o

t

h ar

t

+1

t

t

s

ar

t

s

ar

t

s

t

+1

s

t

+1

o

t

+1

a

t

a

t

-1

min

KL

s

ar

t

s

ar

t

s

ar

t

s

ar

t

s

ar

t

s

ar

t

h ar

t

+1

h ar

t

+1

h ar

t

+1

min

KL

r

t

+1

Vt

+1

a

t

s

ar

t

+1

s

ar

t

+1

h ar

t

+1

h ar

t

+1

r

t

Vt

a

t

-1

s

ar

t

s

ar

t

o

t

h ar

t

h ar

t

r

t

+2

Vt

+2

a

t

+1

s

ar

t

+2

s

ar

t

+2

h ar

t

+2

h ar

t

+2

a

t

+2

h ar

t

+3

h ar

t

+3

̂

̂

̂

(a) Learning world model from experiences.

̂

̂

̂

̂

(b) Learning policy by imagination.

relevant representations by measuring the dependence between s r t and Rt , given at -1: t and s r t -1 , that is I ( s r t , Rt | at -1: t , s r t -1 ). To ensure that s r t are minimally su ffi cient for policy training, we maximize I ( s r t , Rt | at -1: t , s r t -1 ) while minimizing I ( s ¯ r t , Rt | at -1: t , s r t -1 ) to discourage the inclusion of redundant information in s ¯ r t concerning the rewards:

The conditional mutual information can be expressed as the disparity between two mutual information.

After removing the common term, we leverage mutual information neural estimation [27] to approximate the value of mutual information. Thus, we reframe the objective in Eq.7 as follows:

We employ additional neural networks, parameterized by α , to estimate the mutual information. To incorporate the conditions from the original objective, we apply the stop_gradient operation sg to the variable s r t -1 . Similarly, to ensure that the representations s a t are directly controllable by actions, while s ¯ a t are not, we maximize the following objective:

Intuitively, these two objective functions ensure that s r t is predictive of the reward, while s ¯ r t is not; similarly, s a t can be predicted by the action, whereas s ¯ a t cannot. The total objective for learning the world model is summarized as:

The expectation is computed over the dataset and the representation model. Throughout the model learning process, the objectives for estimating mutual information and learning the world model are alternately optimized. A thorough derivation and discussion of the objective function are given in Appendix B.

It is important to note that our approach to learning world models with identifiable factorization is independent of the specific policy optimization algorithm employed. In this work, we choose to build upon the state-of-the-art method Dreamer [6, 7], which iteratively performs exploration, model-fitting, and policy optimization. As illustrated in Figure 2, the learning process for the dynamics involves the joint training of the four categories of latent variables. However, only the variables of s r t are utilized for policy learning. We provide the pseudocode for the IFactor algorithm in Appendix C.

r

o

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 environment that includes distractors. Further, we assess the advantages of our factored world model in policy learning by conducting experiments on variants of Robodesk and DeepMind Control Suite. The reported results are aggregated over 5 runs for the Robodesk experiments and over 3 runs for others. A detailed introduction to each environment is provided in Appendix D, while experiment details and hyperparameters can be found in Appendix E. The source code is available at https: // github.com / AlexLiuyuren / IFactor

5.1 Latent State Identification Evaluation

5.1.1 Synthetic environments

Evaluation Metrics and Baselines We generate synthetic datasets that satisfy the identifiability conditions outlined in the theorems (see Appendix D.1). To evaluate the block-wise identifiability of the four categories of representations, we follow the experimental methodology of [28] and compute the coe ffi cient of determination R 2 from s ∗ t to ˆ s ∗ t , as well as from ˆ s ∗ t to s ∗ t , for ∗ ∈ { ar , ¯ ar , a ¯ r , ¯ a ¯ r } . R 2 can be viewed as the identifiability score. R 2 = 1 in both directions suggests a one-to-one mapping between the true latent variables and the recovered ones. We compare our latent representation learning method with Denoised MDP [9]. We also include baselines along the line of nonlinear ICA: BetaVAE [29] and FactorVAE [30] which do not consider temporal dependencies; SlowVAE [31] and PCL[32] which leverage temporal constraints but assume independent sources; and TDRL [19] which incorporates temporal constraints and non-stationary noise but does not utilize the action and reward information in the Markov Decision Process.

Results Figure 9 demonstrates the e ff ectiveness of our method in accurately recovering the four categories of latent state variables, as evidenced by high R 2 values ( > 0 . 9). In contrast, the baseline methods exhibit distortions in the identification results. Particularly, Denoised MDP assumes independent latent process for xt , yt , and the existence of instantaneous causal e ff ect from xt , yt to zt (see Figure 1(d)), leading to unidentifiable latent variables without further intervention. BetaV AE, FactorVAE, SlowVAE, and PCL, which assume independent sources, show subpar performance. Although TDRL allows for causally-related processes under conditional independence assumptions, it fails to explicitly leverage the distinct transition structures of the four variable categories. More results on the identifiability scores of baselines are provided in Appendix E.1.

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.98

0.15

0.01

0.03

0.10

1.00

0.08

0.03

0.03

0.07

0.97

0.03

0.03

0.02

0.03

0.99

R

2 =0.9839

(a)

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.99

0.13

0.03

0.06

0.25

1.00

0.07

0.05

0.02

0.08

0.97

0.04

0.02

0.02

0.12

0.99

R

2 =0.9850

(b)

0

1

2

3

4

5

6

7

8

Training Steps

1e4

0.0

0.2

0.4

0.6

0.8

1.0

Value

Average

R

2

:

st

to

st

IFactor (ours)

DenoisedMDP

BetaVAE

FactorVAE

PCL

SlowVAE

TDRL

(c)

0

1

2

3

4

5

6

7

8

Training Steps

1e4

0.0

0.2

0.4

0.6

0.8

1.0

Value

Average

R

2

:

st

to

st

(d)

5.1.2 Modified Cartpole with distractors

We have introduced a variant of the original Cartpole environment by incorporating two distractors. The first distractor is an uncontrollable Cartpole located in the upper portion of the image, which is irrelevant to the rewards. The second distractor is a controllable but reward-irrelevant green light positioned below the reward-relevant Cartpole in the lower part of the image.

After estimating the representation model, we utilize latent traversal to visualize the four categories of representations in the modified Cartpole environment (see Figure 4). Specifically, the recovered s ar t corresponds to the position of the cart, while s ¯ ar t corresponds to the angle of the pole. This demonstrates the ability of our method to automatically factorize reward-relevant representations based on (one-step) controllability: the cart experiences a force at the current time step, whereas the pole angle is influenced in the subsequent time step. Note that during the latent traversal of s ¯ ar t , the position of the cart remains fixed, even though the pole originally moves with the cart in the video. Additionally, s a ¯ r t corresponds to greenness of the light, and s ¯ a ¯ r t corresponds to the uncontrollable and reward-irrelevant cartpole. These findings highlight the capability of our representation learning method to disentangle and accurately recover the ground-true latent variables from videos. The identifiability scores of four categories of representations are provided in Appendix E.2.

Cart position

Pole angle

Greeness

Distractor

Obs

s

ar

t

s

¯

ar

s

a

¯

r

t

s

¯

a

¯

r

t

1

-------一-

TTTTT

T

TTTTT

TTTT

T

5.2 Policy Optimization Evaluation

To assess the e ff ectiveness of our method in enhancing policy learning by utilizing minimal yet su ffi cient representations for control, we conducted experiments on more complex control tasks. One of these tasks is a variant of Robodesk [34], which includes realistic noise elements such as flickering lights, shaky cameras, and a dynamic video background. In this task, the objective for the agent is to change the hue of a TV screen to green using a button press. We also consider variants of DeepMind Control Suite (DMC) [35, 9], where distractors such as a dynamic video background, noisy sensor readings, and a jittering camera are introduced to the original DMC environment. These additional elements aim to challenge the agent's ability to focus on the relevant aspects of the task while filtering out irrelevant distractions. Baseline results except for DreamerPro [36] are derived from the Denoised MDP paper. We have omitted the standard deviation of their performance for clarity.

Evaluation Metrics and Baselines We evaluate the performance of the policy at intervals of every 10,000 environment steps and compare our method with both model-based and model-free approaches. Among the model-based methods, we include Denoised MDP [9], which is the state-of-the-art method for variants of Robodesk and DMC. We also include DreamerPro [36], TIA [11] and Dreamer [6] as additional model-based baselines. For the model-free methods, we include DBC [37], CURL[38], and PI-SAC [39]. To ensure a fair comparison, we have aligned all common hyperparameters and neural network structures with those used in the Denoised MDP [9].

5.2.1 RoboDesk with Various Noise Distractors

In Figure 5, the left image demonstrates that our model IFactor achieves comparable performance to Denoised MDP while outperforming other baselines. The results of baselines except for DreamerPro are directly copied from the paper of Denoised MDP [9] (its error bar cannot be directly copied), and the replication of Denoised MDP results is given in Appendix E.3. Furthermore, to investigate whether s r t serves as minimal and su ffi cient representations for policy learning, we retrain policies using the Soft Actor-Critic algorithm [40] with di ff erent combinations of the four learned categories as input. Remarkably, the policy trained using s r t exhibits the best performance, while the performance of other policies degrades due to insu ffi cient information (e.g., s ar t and s ¯ r t ) or redundant information

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps

1e6

300

400

500

600

700

Robodesk with Various Distractors

IFactor (Ours)

DreamerPro

DenoisedMDP

TIA

CURL

DBC

PI-SAC

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps

1e6

400

500

600

700

Using Representations Learned by IFactor

s r

t

s ar

t

s r

t

+

s ar

t

st

s r

t

Env.

Rollout

s

r

t

s

a

¯

r

t

s

¯

r

t

Greeness of

the Screen

Greeness +

Flickering Light

(e.g., s r t + s a ¯ r t and st ). 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 a ¯ r t , which may initially appear counter-intuitive. However, this phenomenon arises from the robot arm pressing the green button, resulting in the screen turning green. Consequently, the screen turning green becomes independent of rewards conditioned on the robot arm pressing the green button, aligning with the definition of s a ¯ r t . This empirically confirms the minimality of s r t for policy optimization.

5.2.2 DeepMind Control Suite (DMC)

Figure 6 demonstrates the consistent superiority of our method over baselines in the Deepmind Control Suite (DMC) across various distractor scenarios. Our approach exhibits robust performance in both noiseless and noisy environments, demonstrating the e ff ectiveness of our factored world model in eliminating distractors while preserving essential information for e ff ective and e ffi cient control in policy learning. The ablation study of IFactor's objective function terms is presented in Appendix B.1 and Appendix E.3, which shows that the inclusion of mutual information constraints is essential to promote disentanglement and enhance policy performance. The policy performance that includes the standard error is provided in Appendix E.4. Visualization of the learned representations is also provided in Appendix E.5.

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 [5, 6, 7, 8] adopt a world model-based learning approach, where the policy is learned solely through dreaming in the world model. However, the agents trained by these techniques tend to underperform in noisy environments [9]. To address this issue, numerous approaches have been introduced to enhance robustness to distractions. Task Informed Abstractions (TIA) [11] explicitly partition the latent state space into reward-relevant and reward-irrelevant features. Denoised MDP [9] takes a step further and decomposes the reward-relevant states into controllable and uncontrollable components. InfoPower [48] prioritizes information that is correlated with action based on mutual information. Iso-Dream [12] learns controllable and noncontrollable sources of spatiotemporal changes on isolated state transition branches. Our method extends these work by accommodating a more general factorization with block-wise identifiablity. Recent research also explores reconstruction-free representation learning methods [36, 49, 50, 51]. DreamerPro [36] combines Dreamer with prototypes, which distills temporal structures from past observations and actions. Temporal Predictive Coding [50] encodes elements in the environment that can be predicted across time. Contrastively-trained Structured World Models [51] utilize graph neural networks and a contrastive approach for representation learning in environments with compositional structure. Latent states learned by these methods are not identifiable due to the reconstruction-free property, where the mixing function is not invertible anymore. A detailed comparison between IFactor and related work is also given in Appendix F.

Identifiability in Causal Representation Learning. Temporal structure and nonstationarities were recently used to achieve identifiability in causal representation learning [52]. Methods such as TCL

Noiseless

Camera Jittering

Cheetah

Run

Walker

Walk

Reacher

Easy

[53] and PCL [32] leverage nonstationarity in temporal data through contrastive learning to identify independent sources. CITRIS [54] and iCITRIS [55] takes advantage of observed intervention targets to identify causal factors. LEAP [18] introduces identifiability conditions for both parametric and nonparametric latent processes. TDRL [19] explores identifiability of latent processes under stationary environments and distribution shifts. ASRs [13] establish the identifiability of the representations in the linear-gaussian case in the RL setting. Our work extends the theoretical results in ASRs to enable block-wise identifiability of four categories of latent variables in general nonlinear cases.

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 blockwise identifiability of these latent categories in general nonlinear cases, under weak and realistic assumptions. Accordingly, we propose IFactor to extract four types of representations from raw observations and use reward-relevant representations for policy optimization. Experiments verify our theoretical results and show that our method achieves state-of-the-art performance in variants of the DeepMind Control Suite and RoboDesk. The basic limitation of this work is that the underlying latent processes are assumed to have no instantaneous causal relations but only time-delayed influences. This assumption does not hold true if the resolution of the time series is much lower than the causal frequency. We leave the identifiability of the latent dynamics in the presence of instantaneous causal and the extension of our method to heterogeneous environments to future work.

References

Le Chang and Doris Y Tsao. The code for facial identity in the primate brain. Cell , 169(6):10131028, 2017.

R Quian Quiroga, Leila Reddy, Gabriel Kreiman, Christof Koch, and Itzhak Fried. Invariant visual representation by single neurons in the human brain. Nature , 435(7045):1102-1107, 2005.

Jürgen Schmidhuber. On learning to think: Algorithmic information theory for novel combinations of reinforcement learning controllers and recurrent neural world models. CoRR , abs / 1511.09249, 2015.

David Ha and Jürgen Schmidhuber. World models. CoRR , abs / 1803.10122, 2018.

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.

Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In Proceedings of the 8th International Conference on Learning Representations , Addis Ababa, Ethiopia, 2020.

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.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse domains through world models. CoRR , abs / 2301.04104, 2023.

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 on Machine Learning , pages 22591-22612, Baltimore, Maryland, 2022.

Claire Glanois, Paul Weng, Matthieu Zimmer, Dong Li, Tianpei Yang, Jianye Hao, and Wulong Liu. A survey on interpretable reinforcement learning. CoRR , abs / 2112.13112, 2021.

Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi S. Jaakkola. Learning task informed abstractions. In Proceedings of the 38th International Conference on Machine Learning , volume 139, pages 3480-3491, Virtual Event, 2021.

Minting Pan, Xiangming Zhu, Yunbo Wang, and Xiaokang Yang. Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models. In Advances in neural information processing systems , volume 35, pages 23178-23191, New Orleans, Louisiana, 2022.

Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-su ffi cient state representation learning for control with structural constraints. In Proceedings of the 39th International Conference on Machine Learning , volume 162, pages 9260-9279, Baltimore, Maryland, 2022.

Aapo Hyvärinen, Shohei Shimizu, and Patrik O. Hoyer. Causal modelling combining instantaneous and lagged e ff ects: an identifiable model based on non-gaussianity. In Proceedings of the 25th International Conference on Machine Learning , volume 307, pages 424-431, Helsinki, Finland, 2008.

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.

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.

Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural networks , 12(3):429-439, 1999.

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.

Weiran Yao, Guangyi Chen, and Kun Zhang. Temporally disentangled representation learning. In Advances in Neural Information Processing Systems 35 , New Orleans, Louisiana, 2022.

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.

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.

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.

Lingjing Kong, Biwei Huang, Feng Xie, Eric Xing, Yuejie Chi, and Kun Zhang. Identification of nonlinear latent hierarchical models. CoRR , abs / 2306.07916, 2023.

Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning , 37:183-233, 1999.

Naftali Tishby, Fernando C. N. Pereira, and William Bialek. The information bottleneck method. CoRR , physics / 0004057, 2000.

Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. In Proceedings of the 5th International Conference on Learning Representations , Toulon, France, 2017.

Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, R. Devon Hjelm, and Aaron C. Courville. Mutual information neural estimation. In Proceedings of the 35th International Conference on Machine Learning , volume 80, pages 530-539, Stockholm, Sweden, 2018.

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.

Irina Higgins, Loïc Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In Proceedings of the 5th International Conference on Learning Representations , Toulon, France, 2017.

Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Proceedings of the 35th International Conference on Machine Learning , volume 80, pages 2654-2663, Stockholm, Sweden, 2018.

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.

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.

Vladimir Vovk. Kernel ridge regression. In Empirical Inference - Festschrift in Honor of Vladimir N. Vapnik , pages 105-116, 2013.

Tongzhou Wang. Robodesk with a diverse set of distractors. https://github.com/SsnL/ robodesk , 2022.

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.00690, 2018.

Fei Deng, Ingook Jang, and Sungjin Ahn. Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations. In Proceedings of the 39th International Conference on Machine Learning , volume 162, pages 4956-4975, Baltimore, Maryland, 2022.

Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In Proceedings of the 9th International Conference on Learning Representations , Virtual Event, Austria, 2021.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised representations for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning , volume 119, pages 5639-5650, Virtual Event, 2020.

Kuang-Huei Lee, Ian Fischer, Anthony Z. Liu, Yijie Guo, Honglak Lee, John F. Canny, and Sergio Guadarrama. Predictive information accelerates learning in RL. In Advances in Neural Information Processing Systems 33 , Virtual Event, 2020.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: O ff -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.

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.

Niklas Wahlström, Thomas B. Schön, and Marc Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. CoRR , abs / 1502.02251, 2015.

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 on Robotics and Automation , pages 1134-1141, Brisbane, Australia, 2018.

Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR , abs / 1807.03748, 2018.

Wilson Yan, Ashwin Vangipuram, Pieter Abbeel, and Lerrel Pinto. Learning predictive representations for deformable objects using contrastive estimation. In Proceedings of the 5th Conference on Robot Learning , pages 564-574, 2021.

Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforcement and infomax learning. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33 , pages 3686-3698, Virtual Event, 2020.

Zheng-Mao Zhu, Shengyi Jiang, Yu-Ren Liu, Yang Yu, and Kun Zhang. Invariant action e ff ect model for reinforcement learning. In Proceedings of the 36th AAAI Conference on Artificial Intelligence , pages 9260-9268, Virtual Event, 2022.

Homanga Bharadhwaj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information prioritization through empowerment in visual model-based RL. In Proceedings of the Tenth International Conference on Learning Representations , Virtual Event, 2022.

Arnav Kumar Jain, Shivakanth Sujit, Shruti Joshi, Vincent Michalski, Danijar Hafner, and Samira Ebrahimi Kahou. Learning robust dynamics through variational sparse gating. In Advances in neural information processing systems , volume 35, pages 1612-1626, New Orleans, Louisiana, 2022.

Tung D. Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139, pages 8130-8139, Virtual Event, 2021.

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.

Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proc. IEEE , 109(5):612-634, 2021.

Aapo Hyvärinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ICA. In Advances in Neural Information Processing Systems 29 , pages 3765-3773, Barcelona, Spain, 2016.

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.

Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. Causal representation learning for instantaneous and temporal e ff ects in interactive systems. In Proceedings of the 11th International Conference on Learning Representations , 2023.

P. Spirtes, C. Glymour, and R. Scheines. Causation, Prediction, and Search . Spring-Verlag Lectures in Statistics, 1993.

J. Pearl. Causality: Models, Reasoning, and Inference . Cambridge University Press, Cambridge, 2000.

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.

Antonin Ra ffi n, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22:268:1-268:8, 2021.

A Theoretical Proofs

A.1 Proof of Proposition 1

Proposition 1 shows that s r t , which has directed paths to rt + τ (for τ ≥ 0), is minimally su ffi cient for policy learning that aims to maximize the future reward and can be characterized by conditional dependence with the cumulative reward variable Rt .

Proposition 1. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data, s r t ⊆ st is a minimal subset of state dimensions that are su ffi cient for policy learning, and si , t ∈ s r t if and only if si , t ̸ ⊥ ⊥ Rt | at -1: t , s r t -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 p over a set of variables V satisfies the global Markov property on graph G if for any partition ( A , B , C ) such that if B d-separates A from C, then p ( A , C | B ) = p ( A | B ) p ( C | B ) .

Definition 2 (Faithfulness Assumption) . There are no independencies between variables that are not entailed by the Markov Condition in the graph.

Below, we give the proof of Proposition 1.

Proof. The proof contains the following three steps.

In step 1, we show that a state dimension si , t is in s r t , that is, it has a directed path to rt + τ , if and only if si , t ̸ ⊥ ⊥ Rt | at -1: t , s t -1 .

In step 2, we show that for si , t with si , t ̸ ⊥ ⊥ Rt | at -1: t , s t -1 , if and only if si , t ̸ ⊥ ⊥ Rt | at -1: t , s r t -1 .

In step 3, we show that s r t are minimally su ffi cient for policy learning.

Step 1: We first show that if a state dimension si , t is in s r t , then si , t ̸ ⊥ ⊥ Rt | at -1: t , s t -1 .

We prove it by contradiction. Suppose that si , t is independent of Rt given at -1: t and s t -1 . Then according to the faithfulness assumption, we can see from the graph that si , t does not have a directed path to rt + τ , which contradicts the assumption, because, otherwise, at -1: t and s t -1 cannot break the paths between si , t and Rt which leads to the dependence.

We next show that if si , t ̸ ⊥ ⊥ Rt | at -1: t , s t -1, then si , t ∈ s r t .

Similarly, by contradiction suppose that si , t does not have a directed path to rt + τ . From the graph, it is easy to see that at -1: t and s t -1 must d-separate the path between si , t and Rt . According to the Markov assumption, si , t is independent of Rt given at -1: t and s t -1, which contradicts to the assumption. Since we have a contradiction, it must be that si , t has a directed path to rt + τ , i.e. si , t ∈ s r t .

Step 2: In step 1, we have shown that si , t ̸ ⊥ ⊥ Rt | at -1: t , s t -1, if and only if it has a directed path to rt + τ . From the graph, it is easy to see that for those state dimensions which have a directed path to rt + τ , at -1: t and s t -1 cannot break the path between si , t and Rt . Moreover, for those state dimensions which do not have a directed path to rt + τ , at -1: t and s r t -1 are enough to break the path between si , t and Rt .

Therefore, for si , t , si , t ̸ ⊥ ⊥ Rt | at -1: t , s t -1 , if and only if si , t ̸ ⊥ ⊥ Rt + 1 | at -1: t , s r t -1 .

Step 3: In the previous steps, it has been shown that if a state dimension si , t is in s r t , then si , t ̸ ⊥ ⊥ Rt | at -1: t , s r t -1 , and if a state dimension si , t is not in s r t , then si , t ⊥ ⊥ Rt | at -1: t , s r t -1 . This implies that s r t are minimally su ffi cient for policy learning to maximize the future reward. □

A.2 Proof of Proposition 2

Moreover, the proposition below shows that s a t , which receives an edge from at -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 a t ⊆ st is a minimal subset of state dimensions that are su ffi cient for direct control, and si , t ∈ s a t if and only if si , t ̸ ⊥ ⊥ at -1 | st -1 .

Below, we give the proof of Proposition 2.

Proof. The proof contains the following two steps.

In step 1, we show that a state dimension si , t is in s a t , that is, it receives an edge from at -1 , if and only if si , t ̸ ⊥ ⊥ at -1 | st -1 .

In step 2, we show that s a t contains a minimally su ffi cient subset of state dimensions that can be directly controlled by actions.

Step 1: We first show that if a state dimension si , t is in s a t , then si , t ̸ ⊥ ⊥ at -1 | st -1 .

We prove it by contradiction. Suppose that si , t is independent of at -1 given s t -1. Then according to the faithfulness assumption, we can see from the graph that si , t does not receive an edge from at -1 , which contradicts the assumption, because, otherwise, s t -1 cannot break the paths between si , t and at -1 which leads to the dependence.

We next show that if si , t ̸ ⊥ ⊥ at -1 | st -1 , then si , t ∈ s a t .

Similarly, by contradiction suppose that si , t does not receive an edge from at -1 . From the graph, it is easy to see that s t -1 must break the path between si , t and at -1 . According to the Markov assumption, si , t is independent of at -1 given s t -1, which contradicts to the assumption. Since we have a contradiction, it must be that si , t has an edge from at -1 .

Step 2: In the previous steps, it has been shown that if a state dimension si , t is in s a t , then si , t ̸⊥ ⊥ at -1 | st -1 , and if a state dimension si , t is not in s a t , then si , t ⊥ ⊥ at -1 | st -1 . This implies that s a t is minimally su ffi cient for one-step direct control. □

A.3 Proof of Proposition 3

Furthermore, based on Proposition 1 and Proposition 2, we can further di ff erentiate s ar t , s ¯ ar t , s a ¯ r t from s r t and s a t , which is given in the following proposition.

Proposition 3. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful 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) si , t ∈ s ar t if and only if si , t ̸ ⊥ ⊥ Rt | at -1 , s r t -1 and si , t ̸ ⊥ ⊥ at -1 | st -1 , (2) si , t ∈ s ¯ ar t if and only if si , t ̸ ⊥ ⊥ Rt | at -1 , s r t -1 and si , t ⊥ ⊥ at -1 | st -1 , (3) si , t ∈ s a ¯ r t if and only if si , t ⊥ ⊥ Rt | at -1 , s r t -1 and si , t ̸ ⊥ ⊥ at -1 | st -1 , and (4) si , t ∈ s ¯ a ¯ r t if and only if si , t ⊥ ⊥ Rt | at -1 , s r t -1 and si , t ⊥ ⊥ at -1 | st -1 .

Proof. This proposition can be easily proved by levering the results from Propositions 1 and 2. □

A.4 Proof of Theorem 1

According to the causal process in the RL system (as described in Eq.1 in [13]), we can build the following mapping from latent state variables s t to observed variables ot and future cumulative reward Rt :

where

Here, note that to recover s r t , it is essential to take into account all future rewards rt : T , because any state dimension si , t ∈ s t that has a directed path to the future reward rt + τ , for τ > 0, is involved in s r t . Hence, we consider the mapping from s r t to the future cumulative reward Rt , and η t represents residuals, except s r t , that have an e ff ect to Rt .

The following theorem shows that the di ff erent types of states s ar t , s ¯ ar t , s a ¯ r t , and s ¯ a ¯ r t are blockwise identifiable from observed image variable ot , reward variable rt , and action variable at , 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 2 and illustrated in Figure 1(c). Under the following assumptions

A1. The mapping f in Eq. 12 is smooth and invertible with smooth inverse.

A2. For all i ∈ { 1 , . . . , do + dR } and j ∈ F i , : , there exist { ˜ s ( l ) t } |F i , : | l = 1 , so that span { J f (˜ s ( l ) t ) i , : } |F i , : | l = 1 = R d ˜ s F i , : , and there exists a matrix T with its support identical to that of J -1 ˆ f ( ˆ ˜ s t ) J f (˜ s t ) , so that GLYPH<2> J f (˜ s ( l ) t ) T GLYPH<3> j , : ∈ R d ˜ s ˆ F i , : .

Then, reward-relevant and controllable states s ar t , reward-relevant but not controllable states s ¯ ar t , reward-irrelevant but controllable states s a ¯ r t , and noise s ¯ a ¯ r t , are blockwise identifiable.

In the theorem presented above, Assumption A 1 only assumes the invertibility of function f , while functions f 1 and f 2 are considered general and not necessarily invertible, as that in [23]. Since the function f 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 f 2 since usually, the reward function is not invertible. Assumption A 2, which is also given in [20, 22, 23], aims to establish a more generic condition that rules out certain sets of parameters to prevent ill-posed conditions. Specifically, it ensures that the Jacobian is not partially constant. This condition is typically satisfied asymptotically, and it is necessary to avoid undesirable situations where the problem becomes ill-posed.

Proof. The proof consists of four steps.

In step 1, we show that s a t = s ar t ∪ s a ¯ r t is blockwise identifiable, by using the characterization that the action variable at only directly influences s ar t and s a ¯ r t .

In step 2, we show that s r t = s ar t ∪ s ¯ ar t is blockwise identifiable, by using the characterization that the future cumulative reward Rt is only influenced by s ar t and s ¯ ar t .

In step 3, we show that s ar t is blockwise identifiable, by using the identifiability of s ar t ∪ s a ¯ r t and s ar t ∪ s ¯ ar t .

In step 4, we further show the blockwise identifiability of s ¯ ar t , s a ¯ r t , and s ¯ a ¯ r t .

Step 1: prove the block identifiability of s a t .

For simplicity of notation, below, we omit the subscript t .

where h = f -1 ◦ ˆ f is the transformation between the true latent variable and the estimated one, and ˆ f : S → X denotes the estimated invertible generating function. Note that as both f -1 and ˆ f are smooth and invertible, h and h -1 is smooth and invertible.

Since h ( · ) is smooth over S , its Jocobian can be written as follows:

The invertibility of h -1 implies that Jh -1 is full rank. Since s a has changing distributions over the action variable a while s ¯ a has invariant distributions over di ff erent values of a , we can derive that C = 0. Furthermore, because Jh -1 is full rank and C is a zero matrix, D must be of full rank, which implies h ′ -1 a is invertible, where h ′ -1 a denotes the first derivative of h -1 a . Therefore, s a is blockwise identifiable up to invertible transformations.

Step 2: prove the blockwise identifiability of s r t .

Recall that we have the following mapping:

where

Note that here to recover s r t , we need to take into account all future rewards, because s r t contains all those state dimensions that have a directed path to future rewards rt + 1: T . η t represents all other factors, except s r t , that influence Rt at time instance t . Further note here we assume the invertibility of f , while f 1 and f 2 are general functions not necessarily invertible.

We denote by ˜ s = ( s r , s ¯ r , η ). We further denote by the dimension of s r by ds r , the dimension of s ¯ r by ds ¯ r , the dimension of ˜ s by d ˜ s , the dimension of o by do , and the dimension of Rt by dR .

We denote by F the support of J f ( s ), by ˆ F the support of J ˆ f (ˆ s ), and by T the support of 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, 22].

Since h : = ˆ f -1 ◦ f , we have ˆ f = f ◦ h -1 (˜ s ). By applying the chain rule repeatedly, we have

With Assumption A2, for any i ∈ { 1 , . . . , do + dR } , there exists { ˜ s ( l ) } |F i , : | l = 1 , s.t. span( { J f (˜ s ( l ) ) i , : } |F i , : | l = 1 ) = R d ˜ s F i , : .

Since { J f (˜ s ( l ) ) i , : } |F i , : | l = 1 forms a basis of R d ˜ s F i , : , for any j 0 ∈ F i , :, we can write canonical basis vector ej 0 ∈ R d ˜ s F i , : as:

where α l ∈ R is a coe ffi cient.

Then, following Assumption A2, there exists a deterministic matrix T such that

where ∈ is due to that each element in the summation belongs to R d ˜ s ˆ F i , : .

Therefore,

Equivalently, we have:

,

We would like to show that ˆ s r does not depend on s ¯ r and η , that is, Ti , j = 0 for i ∈ { 1 , . . . , ds r } and j ∈ { ds r + 1 , . . . , d ˜ s } .

We prove it by contradiction. Suppose that ˆ s r had dependence on s ¯ r , that is, ∃ ( j s r , j s ¯ r ) ∈ T with j s r ∈ { 1 , . . . , ds r } and j s ¯ r ∈ { ds r + 1 , . . . , ds r + ds ¯ r } .

Hence, there must exist i r ∈ { do + 1 , . . . , do + dR } , such that, ( i r , j s ¯ r ) ∈ F .

It follows from Equation 19 that:

However, due to the structure of ˆ f 2 , [ J ˆ f 2 ] i r , j s ¯ r = 0, which results in a contradiction. Therefore, such ( i r , j s ¯ r ) does not exist and ˆ s r does not depend on s ¯ r . The same reasoning implies that ˆ s r does not dependent on η . Thus, ˆ sr does not depend on ( s ¯ r , η ). In conclusion, ˆ s r does not contain extra information beyond s r .

Similarly, we can show that (ˆ s ¯ r , ˆ η ) does not contain information of sr .

Therefore, there is a one-to-one mapping between s r and ˆ s r .

Step 3: prove the blockwise identifiability of s ar t .

In Step 1 and Step 2, we have shown that both s a and s r are blockwise identifiable. That is,

where ha and hr are invertible functions.

According to the invariance relation of s ar , We have the following relations:

It remains to show that both ˜ hr : = hr ( · ) 1: ds ar and ˜ ha : = ha ( · ) 1: ds ar do not depend on s ¯ ar and s a ¯ r in their arguments.

We will prove this by contradiction, following the proof technique in [58]. Without loss of generality, we suppose ∃ l ∈ { 1 , · · · , ds ¯ ar } , s r ∗ ∈ S r , s.t., ∂ ˜ hr ∂ s ¯ ar l ( s r ∗ ) , 0. As h is smooth, it has continuous partial derivatives. Thus, ∂ ˜ hr ∂ s ¯ ar l , 0 holds true in a neighbourhood of s r ∗ , i.e.,

where s ¯ ar -l denotes variable s ¯ ar excluding the dimension l .

We further define an auxiliary function ψ : S ar × S ¯ ar × S a ¯ r → R ≥ 0 as follows:

To obtain the contradiction to the invariance, it remains to show that ψ > 0 with a probability greater than zero w.r.t. the true generating process.

There are two situations at ( s ar ∗ , s ¯ ar ∗ , s a ¯ r ∗ ) where s ¯ ar ∗ is an arbitrary point in S ¯ ar :

situation 1: ψ ( s ar ∗ , s ¯ ar ∗ , s a ¯ r ∗ ) > 0;

situation 2: ψ ( s ar ∗ , s ¯ ar ∗ , s a ¯ r ∗ ) = 0.

In situation 1, we have identified a specific point ψ ( s ar ∗ , s ¯ ar ∗ , s a ¯ r ∗ ) that makes ψ > 0.

In situation 2, Eq. 23 implies that ∀ s ¯ ar l ∈ ( s ar l ∗ , s ar l ∗ + η ), ψ ( s ar ∗ , ( s ¯ ar -l ∗ , s ¯ ar l ) , s a ¯ r ∗ ) > 0.

Thus, in both situations, we can locate a point ( s ar ∗ , s ¯ ar ∗ ′ , s a ¯ r ∗ ) such that ψ ( s ar ∗ , s ¯ ar ∗ ′ , s a ¯ r ∗ ) > 0, where s ¯ ar ∗ ′ = s ¯ ar ∗ in situation 1 and s ¯ ar l ∗ ′ ∈ ( s ar l ∗ , s ar l ∗ + η ) , s ar -l ∗ ′ = s ar -l ∗ in situation 2.

Since ψ is a composition of continuous functions, it is continuous. As pre-image of open sets are always open for continuous functions, the open set R > 0 has an open set U ∈ S ar × S ¯ ar × S a ¯ r as its preimage. Due to ( s ar ∗ , s ¯ ar ∗ ′ , s a ¯ r ∗ ) ∈ U , U is nonempty. Further, as U is nonempty and open, U has a Lebesgue measure of greater than zero.

As we assume that ps ar , s ¯ ar , s a ¯ r is fully supported over the entire domain S ar ×S ¯ ar ×S a ¯ r , we can deduce that P p [ U ] > 0. That is, ψ > 0 with a probability greater than zero, which contradicts the invariance condition, Therefore, we can show that ˆ hr ( s r ) does not depend on s ¯ ar .

Similarly, we can show that ˆ ha ( s a ) does not depend on s a ¯ r .

Finally, the smoothness and invertibility of ˆ hr and ˆ ha follow from the smoothness and invertibility of hr and ha over the entire domain.

Therefore, hr ( ha ) is a smooth invertible mapping between s ar and ˆ s ar . That is, s ar is blockwise invertible.

Step 4: prove the blockwise identifiability of s ¯ ar t , s a ¯ r t , and s ¯ a ¯ r t .

We can use the same technique in Step 3 to show the identifiability of s ¯ ar and s a ¯ r . Specifically, since sr and s ar are identifiable, we can show that s ¯ ar is identifiable. Similarly, since s a and s ar are identifiable, we can show that s a ¯ r is identifiable. Furthermore, since s ar , s ¯ ar , and s a ¯ r are identifiable, we can show that s ¯ a ¯ r is identifiable □

B Derivation of the Objective Function

We start by defining the components of the world mode as follows:

The latent dynamics can be disentangled into four catogories:

We follow the derivation framework in Dreamer [6] and define the information bottleneck objective for latent dynamics models [25]

where β is scalar and i t are dataset indices that determine the observations p ( ot | i t ) = δ ( ot -¯ ot ) as in [26]. Maximizing the objective leads to 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:

Thus, we obtain the objective function:

For the second term in Equation 27, we use the non-negativity of the KL divergence to obtain an upper bound,

According to equation 26, we have p γ = p γ 1 · p γ 2 · p γ 3 · p γ 4 and q ϕ = q ϕ 1 · q ϕ 2 · q ϕ 3 · q ϕ 4 .

We introduce additional hyperparameters to regulate the amount of information contained within each category of variables:

Additionally, we introduce two supplementary objectives to explicitly capture the distinctive characteristics of the four distinct representation categories. Specifically, we characterize the rewardrelevant representations by measuring the dependence between s r t and Rt , given at -1: t and s r t -1 , that is I ( s r t , Rt | at -1: t , s r t -1 ) (see Figure 15(a). Note that if the action at is not dependent on s r t , such as when it is randomly chosed. at does not need to be conditioned. In this case, the mutual information turns into I ( s r t , Rt | at -1 , s r t -1 ). To ensure that s r t are minimally su ffi cient for policy training, we maximize I ( s r t , Rt | at -1: t , s r t -1 ) while minimizing I ( s ¯ r t , Rt | at -1: t , s r t -1 ) to discourage the inclusion of redundant information in s ¯ r t concerning the rewards:

The conditional mutual information can be expressed as the disparity between two mutual information values.

Combining the above two equations, we eliminated the identical terms, ultimately yielding the following formula

We use the Donsker-Varadhan representation to express mutual information as a supremum over functions,

We employ mutual information neural estimation [27] to approximate the mutual information value. We represent the function T using a neural network that accepts variables ( x , y ) as inputs and is parameterized by α . The neural network is optimized through stochastic gradient ascent to find the supremum. Substituting x and y with variables defined in Equation 35, our objective is reformulated as follows:

To incorporate the conditions from the original objective, we apply the stop_gradient operation to the variable s r t -1 . Similarly, to ensure that the representations s a t are directly controllable by actions, while s ¯ a t are not, we maximize the following objective:

By splitting the conditional mutual information and eliminating identical terms, we obtain the following objective function:

where α 1 , α 2 , α 3 , α 4 can be obtained by maximizing Equation 36. Intuitively, these two objective functions ensure that s r t is predictive of the reward, while s ¯ r t is not; similarly, s a t can be predicted by the action, whereas s ¯ a t cannot.

Combine the equation 29, equation 32, equation 37 and equation 39, the total objective function is:

The expectation is computed over the dataset and the representation model. Throughout the model learning process, the objectives for estimating mutual information and learning the world model are alternately optimized.

B.1 Discussions

In this subsection, we examine the mutual information constraints in equation 37 and equation 39 and their relationship with other objectives. Our findings reveal that while other objectives partially fulfill the desired functionality of the mutual information constraints, incorporating both mutual information objectives is essential for certain environments.

The objective functions can be summarized as follows::

and the KL divergence term can be further decomposed into 4 components:

Specifically, maximizing I α 1 in J t RS enhances the predictability of Rt based on the current state s r t -1 conditioning on ( s r t -1 , at -1: t ). However, notice that this objective can be partially accomplished by optimizing J t R . When learning the world model, both the transition function and the reward function are trained: the reward function predicts the current reward rt using s r t , while the transition model predicts the next state. These combined predictions contribute to the overall prediction of Rt .

Minimizing I α 2 in J t RS eliminates extraneous reward-related information present in s ¯ r t . According to our formulation, s ¯ r t can still be predictive of Rt as long as it does not introduce additional predictability beyond what is already captured by ( s r t -1 , at -1: t ). This is because we only assume that s ¯ r t is conditionally independent from Rt when conditioning on ( s r t -1 , at -1: t ). If we don't condition on ( s r t -1 , at -1: t ), it introduces s r t -1 as the confounding factor between s ¯ r t and Rt , establishing association between s ¯ r t and Rt (refer to Figure 15(a)). Note that the KL divergence constraints govern the information amount within each state category. By amplifying the weight of the KL constraints on s ¯ r t , the value of I α 2 can indirectly be diminished.

By maximizing I α 3 and minimizing I α 4 in J t AS , we ensure that s a t can be predicted based on ( at -1 , st -1 ) while s ¯ a t cannot. The KL constraints on s ar t and s a ¯ r t incorporate the action at -1 into the prior and posterior, implicitly requiring that s a t should be predictable given at -1. Conversely, the KL constraints on s ¯ ar t and s ¯ a ¯ r t do not include the action at -1 in the prior and posterior, implicitly requiring that s a t should not be predictable based on at -1 . However, relying solely on indirect constraints can sometimes be ine ff ective, as it may lead to entangled representations that negatively impact policy performance (see Figure 13).

Ablation of the mutual information constraints. The inclusion of both J t RS and J t AS is essential in certain environments to promote disentanglement and enhance policy performance, despite sharing

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps

1e6

0

200

400

600

800

Reacher+Video+CameraJittering

IFactor

IFactor w/o MI

DenoisedMDP

some common objectives. We have observed improved training stability in the variant of Robodesk environment (see Figure 12) and significant performance gains in the Reacher environment with video background and camera jittering (see Figure 7). When two mutual information objectives are removed, we notice that entangled representations emerged in these environments, as depicted in Figure 13. We assign values of 0.1 to λ 1 and λ 2 in the environment of modified Cartpole, variant of Robodesk and Reacher with video background and jittering camera. Empirically, a value of 0.1 has been found to be preferable for both λ 1 and λ 2. Using a higher value for regularization might negatively impact the learning of representation and transition model. In other DMC environments, the ELBO loss alone has proven e ff ective due to the inherent structure of our disentangled latent dynamics. The choice of hyperparameters ( β 1 , β 2 , β 3 , β 4) depends on the specific goals of representation learning and the extent of noise interference in the task. If the objective is to accurately recover the true latent variables for understanding the environment, it is often e ff ective to assign equal weights to the four KL divergence terms (for experiments on synthetic data and modified cartpole). When the aim is to enhance policy training stability by mitigating noise, it is recommended to set the values of β 1 and β 2 higher than β 3 and β 4 (for experiments on variants of Robodesk and DMC). Moreover, in environments with higher levels of noise, it is advisable to increase the discrepancy in values between the hyperparameters.

C Algorithm

Algorithm 1: IFactor

Input:

Representation model: q ϕ ( st | st -1 , at -1 , ot ) = q ϕ 1 · q ϕ 2 · q ϕ 3 · q ϕ 4 ;

Transition model: p γ ( st | st -1 , at -1 ) = p γ 1 · p γ 2 · p γ 3 · p γ 4 ;

Observation Model: p θ ( ot | st ); Reward model: p θ ( rt | s r t );

Policy Model: πψ ( at | s r t ); Value model: v ψ ( s r t );

Mutual information estimator: I α 1 , I α 2 , I α 3 , I α 4 ;

Policy optimization algorithm P i -O pt , which is in default the same as that used in Dreamer.

Output:

Transition Model p γ . Representation Model q ϕ . Policy Model πψ .

1: while training do

2: // Exploration

3: Collect new trajectories D ′ = { ( ot , at , rt ) } t with π acting on q ϕ encoded outputs

4: Add experience to the replay bu ff er D = D ∪ D ′

5: // Model learning

6: Sample a batch of data sequences { ( ot , at , rt ) } t from the reply bu ff er D

7: Obtain st by the representation model and estimate mutual information terms I α 1 , I α 2 , I α 3 , I α 4

8: for i = 1 to n do

9: Train mutual information neural estimators by maximizing I α 1 , I α 2 , I α 3 , I α 4

10: end for

11: Freeze the parameters in I α 1 , I α 2 , I α 3 , I α 4 and Train q ϕ , p γ and p θ with Equation 11

12: // Policy learning by dreaming

13: Imagine trajectories of s r t using the learned world model.

14: Train πψ and v ψ by running P i -O pt

15: end while

D Environment Descriptions

D.1 Synthetic Data

For the sake of simplicity, we consider one lag for the latent processes in Section 4. Our identifiability proof can actually be applied for arbitrary lags directly because the identifiability does not rely on the number of previous states. We extend the latent dynamics of the synthetic environment to incorporate a general time-delayed causal e ff ect with τ ≥ 1 in the synthetic environment. When τ = 1, it reduces to a common MDP. The ground-truth generative model of the environment is as follows::

Data Generation We generate synthetic datasets with 100, 000 data points according to the generating process in Equation 43, which satisfies the identifiability conditions stated in Theorem 1. The latent variables st have 8 dimensions, where s ar t = s ¯ ar t = s a ¯ r t = s ¯ a ¯ r t = 2. At each timestep, a one-hot action of dimension 5, denoted as at , is taken. The lag number of the process is set to τ = 2. The observation model p θ ( ot , | st ) is implemented using a random three-layer MLPs with LeakyReLU units. The reward model p θ GLYPH<0> rt , | s r t GLYPH<1> is represented by a random one-layer MLP. It's worth noting that the reward model is not invertible due to the scalar nature of rt . Four distinct transition functions, namely p γ 1 , p γ 2 , p γ 3 , and p γ 4 , are employed and modeled using random one-layer MLP with LeakyReLU units. The process noise is sampled from an i.i.d. Gaussian distribution with a standard deviation of σ = 0 . 1. To simulate nonstationary noise for various latent variables in RL, the process noise terms are coupled with the historical information by multiplying them with the average value of all the time-lagged latent variables, as suggested in [19].

D.2 Modified Cartpole

Wehave 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 a ff ect the rewards. The second distractor is a controllable green light positioned below the reward-relevant Cartpole in the lower part of the image, but it is not associated with any rewards. The task-irrelevant cartpole undergoes random actions at each time step and stops moving when its angle exceeds 45 degrees or goes beyond the screen boundaries. The action space consists of three independent degrees of freedom: direction (left or right), force magnitude (10N or 20N), and green light intensity (lighter or darker). This results in an 8-dimensional one-hot vector. The objective of this variant is to maintain balance for the reward-relevant cartpole by applying suitable forces.

D.3 Variant of Robodesk

The RoboDesk environment with noise distractors [9] is a control task designed to simulate realistic sources of noise, such as flickering lights and shaky cameras. Within the environment, there is a large TV that displays natural RGB videos. On the desk, there is a green button that controls both the hue of the TV and a light. The agent's objective is to manipulate this button in order to change the TV's hue to green. The agent's reward is determined based on the greenness of the TV image. In this environment, all four types of information are present (see Table 1).

Reacher Easy

Walker Walk

Modified Cartpole

Robodesk

Table 1: Categorization of various types of information in the environments we evaluated. We use the red color to emphasize the categorization di ff erence between IFactor and Denoised MDP. Unlike Denoised MDP that assumes independent latent processes, IFactor allows for causally-related processes. Therefore, in this paper, the term "controllable" refers specifically to one-step controllability, while "reward-relevant" is characterized by the conditional dependence between s ∗ t and the cumulative reward variable Rt when conditioning on ( st -1 , at -1: t ). Following this categorization, certain agent information can be classified as (one-step) uncontrollable (including indirectly controllable and uncontrollable factors) but reward-relevant factors, such as some position information determined by the position and velocity in the previous time-step rather than the action. On the other hand, the green hue of TV in Robodesk is classified as controllable but reward-irrelevant factors, as they are independent of the reward given the state of the robot arm and green button, aligning with the definition of s a ¯ r t .

D.4 Variants of DeepMind Control Suite

Four variants [9] are introduced for each DMC task:

Noiseless : Original environment without distractors.

Video Background : Replacing noiseless background with natural videos [37] ( Ctrl + Rew ).

Video Background + Sensor Noise : Imperfect sensors sensitive to intensity of a background patch ( 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 are replaced. Camera jittering is introduced through a smooth random walk shift using Gaussian-perturbing acceleration, velocity decay, and pulling force. Sensor noise is added by perturbing a specific sensor based on the intensity of a patch in the background video. The perturbation involves adding the average patch value minus 0.5. Di ff erent sensors are perturbed for di ff erent environments. These sensor values undergo non-linear transformations, primarily piece-wise linear, to compute rewards. While the additive reward noise model may not capture sensor behavior perfectly, it is generally su ffi cient as long as the values remain within moderate ranges and stay within one linear region. (Note: the variants of Robodesk and DMC are not the contributions of this paper. We kindly refer readers to the paper of Denoised MDP [9] for a more detailed introduction.)

E Experimental Details

Computing Hardware Weused a machine with the following CPU specifications: Intel(R) Xeon(R) Silver 4110 CPU @ 2.10GHz; 32 CPUs, eight physical cores per CPU, a total of 256 logical CPU units. The machine has two GeForce RTX 2080 Ti GPUs with 11GB 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 a similar experimental setup to TDRL [19], while extending it by decomposing the dynamics into four causally related latent processes proposed in this paper (refer to Equation 43). For all experiments, we assign β 1 = β 2 = β 3 = β 4 = 0 . 003 as the weights for the KL divergence terms. In this particular experiment, we set λ 1 and λ 2 to 0 because the utilization of the ELBO loss alone has e ff ectively maximized J t RS and J t AS , as illustrated in Figure 10. Here, J t RS represents I α 1 -I α 2 , and J t AS represents I α 3 -I α 4 . 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 AdamW optimizer for 100 epochs. A 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.1 The identifiability socres for baselines

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.34

0.12

0.43

0.12

0.20

0.39

0.19

0.15

0.29

0.26

0.19

0.09

0.17

0.19

0.18

0.45

R

2 =0.3430

(a) DenoisedMDP

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.32

0.22

0.38

0.32

0.12

0.46

0.31

0.32

0.35

0.24

0.23

0.27

0.17

0.21

0.13

0.47

R

2 =0.3700

(b) DenoisedMDP

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.04

0.16

0.36

0.14

0.32

0.28

0.28

0.07

0.14

0.21

0.35

0.23

0.27

0.18

0.35

0.24

R

2 =0.2274

(c) BetaVAE

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.05

0.34

0.14

0.25

0.14

0.25

0.21

0.30

0.38

0.32

0.45

0.43

0.22

0.12

0.28

0.31

R

2 =0.2643

(d) BetaVAE

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.20

0.27

0.36

0.15

0.28

0.11

0.13

0.17

0.07

0.14

0.21

0.25

0.23

0.32

0.33

0.11

R

2 =0.1599

(e) FactorVAE

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.22

0.22

0.17

0.31

0.17

0.13

0.15

0.34

0.30

0.26

0.27

0.39

0.23

0.26

0.35

0.16

R

2 =0.1947

(f) FactorVAE

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.08

0.16

0.50

0.09

0.07

0.12

0.46

0.10

0.07

0.15

0.43

0.14

0.07

0.31

0.28

0.13

R

2 =0.1930

(g) PCL

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.14

0.17

0.15

0.09

0.34

0.24

0.26

0.28

0.43

0.44

0.35

0.28

0.12

0.19

0.15

0.12

R

2 =0.2122

(h) PCL

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.41

0.34

0.17

0.15

0.14

0.06

0.13

0.54

0.15

0.13

0.06

0.55

0.18

0.14

0.02

0.33

R

2 =0.2136

(i) SlowVAE

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.36

0.28

0.13

0.32

0.28

0.09

0.32

0.21

0.26

0.17

0.23

0.07

0.28

0.33

0.43

0.32

R

2 =0.2500

(j) SlowVAE

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.36

0.62

0.21

0.05

0.12

0.11

0.37

0.52

0.08

0.09

0.30

0.40

0.24

0.22

0.15

0.22

R

2 =0.2459

(k) TDRL

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.39

0.17

0.12

0.23

0.71

0.12

0.11

0.28

0.22

0.45

0.27

0.13

0.05

0.59

0.39

0.22

R

2 =0.2508

(l) TDRL

E.1.2 Extra Results.

During the training process, we record the estimation value of four mutual information (MI) terms. The corresponding results are presented in Figure 10. Despite not being explicitly incorporated into the objective function, the terms I α 1 -I α 2 and I α 3 -I α 4 exhibit significant maximization. Furthermore, the estimation values of I α 2 and I α 4 are found to be close to 0. These findings indicate that the state variable s ¯ r t contains little information about the reward, and the predictability of s ¯ a t by the action is also low.

E.2 Modified Cartpole

In the modified Cartpole environment, we configure the values as follows: β 1 = β 2 = β 3 = β 4 = 0 . 1 and λ 1 = λ 2 = 0 . 1. Recurrent State Space Model (RSSM) uses a deterministic part and a stochastic part to represent latent variables. The deterministic state size for four dynamics 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 3 and Table 4 (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

0

2

4

6

8

Training Steps

1e4

1

0

1

2

3

4

Estimation Value

Estimated MI Value for

I

1

0

2

4

6

8

Training Steps

1e4

0.1

0.0

0.1

0.2

0.3

Estimated MI Value for

I

2

0

2

4

6

8

Training Steps

1e4

0.05

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Estimated MI Value for

I

3

0

2

4

6

8

Training Steps

1e4

0.04

0.02

0.00

0.02

0.04

0.06

0.08

Estimated MI Value for

I

4

0

2

4

6

8

Training Steps

1e4

0

1

2

3

4

Estimation Value

I

1

I

2

0

2

4

6

8

Training Steps

1e4

0.0

0.1

0.2

0.3

I

3

I

4

.

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 ¯ ar t as the pole's angle, s a ¯ r t as the light's greenness, and s ¯ a ¯ r t as the distractor Cartpole's state (including cart position and pole angle). The disentanglement scores for individual s ar t and s ¯ ar t , as well as the combined s r t , are shown in Figure 11. We can obviously see that the true latent variables can be clearly identified.

s ar

t

s ar

t

s ar

t

s ar

t

True latents

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

0.74

0.53

-0.03

0.05

0.55

0.80

0.00

0.05

-0.00

-0.01

0.79

-0.00

0.14

0.12

0.00

0.85

R

2 =0.7953

(a) Cartpole

s ar

t

s ar

t

s ar

t

s ar

t

Estimated latents

s ar

t

s ar

t

s ar

t

s ar

t

True latents

0.91

0.36

-0.04

0.29

0.28

0.80

-0.02

0.47

-0.02

0.04

0.79

-0.00

-0.03

-0.09

-0.00

0.93

R

2 =0.8579

(b) Cartpole

s r

t

s ar

t

s ar

t

True latents

s r

t

s ar

t

s ar

t

Estimated latents

0.93

-0.01

0.04

-0.01

0.79

-0.00

0.13

0.00

0.85

R

2 =0.8544

(c) Cartpole

s r

t

s ar

t

s ar

t

Estimated latents

s r

t

s ar

t

s ar

t

True latents

0.97

-0.03

0.38

0.18

0.79

-0.00

-0.02

-0.00

0.93

R

2 =0.8960

(d) Cartpole

E.3 Variant of Robodesk

In the variant of Robodesk, we conduct experiments with the following hyperparameter settings: β 1 = β 2 = 2, β 3 = β 4 = 0 . 25, and λ 1 = λ 2 = 0 . 1. For the four dynamics, we set the deterministic state sizes to (120, 40, 40, 40), and the stochastic state sizes to (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 a hidden size of 128. To ensure a 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 by 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 12. The constraints J t RS and J t AS are observed to stabilize the training process of IFactor. The results of IFactor are areaveraged over 5 runs, while the results of Denoised MDP and IFactor without MI are averaged over three runs.

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps

1e6

300

400

500

600

700

800

Robodesk with Various Distractors

IFactor

IFactor w/o MI

Denoised MDP

Policy learning based on the learned representations by IFactor We retrain policies using the Soft Actor-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 to form the feature. For instance, when referring to s r t , we use both the deterministic states and stochastic states of s r t . The implementation of SAC algorithm is based on Stableb-Baselines3[59], with a learning rate of 0.0002. Both the policy network and Q network consist of 4-layer MLPs with a hidden size of 256. We use the default hyperparameter settings in Stable-Baselines3 for other parameters.

E.4 Variants of Deep Mind Control Suite

In the noiseless DMC environments, we set β 1 = β 2 = β 3 = β 4 = 1. For the DMC environments with video background, we set β 1 = β 2 = 1 and β 3 = β 4 = 0 . 25. In the DMC environments with video background and noisy sensor, we set β 1 = β 2 = 2 and β 3 = β 4 = 0 . 25. Lastly, for the DMC environments with video background and jittering camera, we set β 1 = β 2 = 1 and β 3 = β 4 = 0 . 25. Regarding the Reacher environment with video background and jittering camera, we set λ 1 = λ 2 = 0 . 1 for our experiments. For the other environments, we set λ 1 = λ 2 = 0. The deterministic state sizes for the four dynamics are set to (120, 120, 60, 60), while the stochastic state sizes are set to (20, 20, 10, 10). The four mutual information neural estimators utilize a 4-layer MLPs with a hidden size of 128. We align the other hyperparameters and network structure with those used in the Denoised MDP for a fair comparison.

E.4.1 Policy optimization results on variants of DMC

We present detailed evaluation results in Table 6, showcasing both the mean values and the corresponding 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 J t RS and J t AS . To further investigate how they a ff ects 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 I α 1 -I α 2 and I α 3 -I α 4 are maximized for both IFactor and IFactor without MI. However, IFactor exhibits a significantly higher rate of maximizing I α 3 -I α 4 compared to IFactor without MI. This increased maximization leads to greater predictability of s a t 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, s a t is the controllable and reward relevant representation. s r t = ( s ar t , s ¯ ar t ) is the reward-relevant representation. s ¯ a ¯ r t is the controllable but reward-irrelevant representation. s ¯ a ¯ r t is the uncontrollable and reward-irrelevant representation (noise). s ¯ r t = ( s a ¯ r t , s ¯ a ¯ r t ) is the reward-irrelevant representation. Only representations of s r t 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 of representation contains, surpassing the limitations of the original decoder that is used in latent traversal. The visualization results are shown in Figure 14. It can be observed that s ar t captures the movement of the agent partially but not well enough; s r t captures the movement of the agent precisely but s ¯ r t fails (Reacher and Cheetah) or captures extra information of the background (Walker). This finding

suggests that s r t contains su ffi cient information within the original noisy observation for e ff ective control, while e ff ectively excluding other sources of noise.

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps

1e6

0

2

4

6

Estimation Value

1e

2

Estimated MI Value for

I

1

IFactor

IFactor w/o MI

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps

1e6

0

1

2

3

4

5

6

1e

2

Estimated MI Value for

I

2

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps

1e6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

Estimated MI Value for

I

3

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps

1e6

0.00

0.05

0.10

0.15

0.20

0.25

Estimated MI Value for

I

4

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps

1e6

1.0

0.5

0.0

0.5

1.0

Estimation Value

1e

2

I

1

I

2

IFactor

IFactor w/o MI

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps

1e6

0.00

0.25

0.50

0.75

1.00

1.25

I

3

I

4

.

Env.

Rollout

s

ar

t

s

r

t

s

a

¯

r

t

s

¯

a

¯

r

t

s

¯

r

t

Reacher Easy

Env.

Rollout

s

ar

t

s

r

t

s

a

¯

r

t

s

¯

a

¯

r

t

s

¯

r

t

Env.

Rollout

s

ar

t

s

r

t

s

a

¯

r

t

s

¯

a

¯

r

t

s

¯

r

t

Walker Walk

F Comparison between IFactor and Related Work

F.1 Comparison between IFactor and Denoised MDP

While both IFactor and Denoised MDP share the common aspect of factorizing latent variables based on controllability and reward relevance, it is crucial to recognize the numerous fundamental distinctions between them.

s

ar

t

-1

s

¯

ar

t

-1

s

a

¯

r

t

-1

s

¯

a

¯

r

t

-1

Ot

-1

s

ar

t

s

¯

ar

t

s

a

¯

r

t

s

¯

a

¯

r

t

Ot

a

t

-1

s

ar

t

+1

s

¯

ar

t

+1

s

a

¯

r

t

+1

s

¯

a

¯

r

t

+1

Ot

+1

r

t

a

t

r

t

+1

Rt

(a) IFactor (ours).

x

t

-1

y

t

-1

z

t

-1

Ot

-1

x

t

y

t

z

t

Ot

a

t

-1

r

x

t

r

y

t

+

r

t

(b) Denoised MDP.

First and foremost, Denoised MDP adopts a rather stringent assumption by solely considering three types of latent variables and assuming independent latent processes for xt (controllable and reward-relevant) and yt (uncontrollable but reward-relevant). However, this strict assumption may not 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 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 ryt from rxt , disregarding the possibility of non-additive e ff ects in many uncontrollable yet reward-relevant factors. In contrast, our method embraces the inclusion of non-additive e ff ects of s ¯ ar t on the reward, which is more general.

Thirdly, Denoised MDP uses only controllable and reward-relevant latent variables xt for policy optimization, which we show in the theoretical analysis that it is generally insu ffi cient. 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 e ff ect from xt to zt and from yt to zt , which is inherently unidentifiable without further intervention. It is worth noting that imposing interventions on the latent states is unrealistic in 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 instantaneous causal e ff ect 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 su ffi cient reward-relevant variables for policy optimization, and (2) it provides a 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 di ff ers from that of Denoised MDP. The structure of prior and posterior of IFactor is shown as follows:

While Denoised MDP has the following prior and posterior:

Anotable 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 xt and yt , whereas IFactor only incorporates conditional independence, utilizing s r t -1 as input for the transition of both s ar t and s ¯ ar t . Moreover, the posterior of yt receives at -1 as input, potentially implying controllability. Similarly, the posterior of xt incorporates zt -1 as input, which may introduce noise from zt -1 into xt . 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 two supplementary mutual information constraints, namely J t RS and J t AS , to promote disentanglement and improve policy performance.

F.2 Comparison between IFactor, InfoPower and IsoDream

IFactor learns di ff erent categories of state representations according to their relation with action and reward, which is di ff erent from InfoPower [48] and IsoDream [12], 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 ¯ ar t and s ¯ a ¯ r t . IsoDream uses three branches for latent dynamics, distinguishing controllable, noncontrollable, and static parts.

Other di ff erences with InfoPower:

Reconstruction Basis : InfoPower is reconstruction-free, while IFactor is reconstructionbased.

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.

Other di ff erences 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 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLearning 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe paper briefly talks about a few limitations between lines 346-350 and needs to elaborate more on the discussion and why they are challenging.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nNegative societal impact is not discussed", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nNo serious weakness. I have some comments regarding how section 2 and 3 are written but otherwise it is a solid paper.\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\n1. 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}