Learning World Models with Identifiable Factorization

Summary

Learning efficient world models requires architectural priors to learn better representations that can capture different aspects of the environment. However, existing methods like Dreamer do not focus on learning disentangled representations and lack the ability to separate noise from the reward-relevant information in the observations. In this work, a new method is proposed where the latent state is divided into 4 blocks based on the dependence of action and reward function. The paper shows that using the blocks of states that affect the reward function can be used to learn policies. Experiments show the efficacy of the proposed algorithm on Robodesk and DMControl tasks with distractions and noise in the form of background or based on a sensor/camera. Moreover, ablation studies show that the model learns to disentangle the factors in the environment. Lastly, the paper also provides a theoretical justification for the proposed architecture.

Strengths

1. The idea of disentangling the learned representations to different factors based on actions and rewards is interesting and nicely explained in the paper. 2. 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. 3. The paper presents a theoretical derivation of the identifiability of the latent state representations.

Weaknesses

1. World Models have been extended to scenarios with noisy distractions in [1][2][3]. They have shown improvements over the Dreamer method for such tasks. However, these methods are not used as baselines in the experiments. 2. Some ablations and experiments are required to further understand the scenarios where the proposed method works (More on this in Questions below)

Questions

1. For learning efficient latent representations, recent works like DreamerPro have used auxiliary losses to extract relevant concepts from the observation [1]. How does the proposed method compare with reconstruction-free model-based methods? 2. 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? 3. 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? 4. 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? 5. 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? 6. 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. ### References [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.\ [2] Jain, Arnav Kumar, et al. "Learning robust dynamics through variational sparse gating." Advances in Neural Information Processing Systems 35 (2022): 1612-1626.\ [3] Nguyen, Tung D., et al. "Temporal predictive coding for model-based planning in latent space." International Conference on Machine Learning. PMLR, 2021.\ [4] Kipf, Thomas, Elise Van der Pol, and Max Welling. "Contrastive learning of structured world models." arXiv preprint arXiv:1911.12247 (2019).

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The paper briefly talks about a few limitations between lines 346-350 and needs to elaborate more on the discussion and why they are challenging.

Summary

- The paper introduces a framework called IFactor for modeling latent state variables in reinforcement learning (RL) systems. The framework categorizes these variables into four distinct types based on their interactions with actions and rewards. The paper further establishes block-wise identifiability of these latent variables, which provides a stable and compact representation and discloses that all reward-relevant factors are significant for policy learning. Overall, the paper contributes a comprehensive framework, IFactor, that provides representations of different aspects of information in RL systems and highlights the importance of considering states that influence the reward during decision-making.

Strengths

- IFactor models different latent state variables in based on their interactions with actions and rewards in RL systems. These categories include reward-relevant and controllable parts, reward-irrelevant but uncontrollable parts, controllable but reward-irrelevant parts, and unrelated noise. - 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. - The paper includes extensive experiments and ablation studies to validate the effectiveness of the proposed framework and to analyze the impact of different components. Meanwhile, the paper includes a thorough evaluation of the proposed framework on several benchmark tasks.

Weaknesses

- Missing related work. The separate modeling of controllable and uncontrollable components has been explored in various previous works, such as InfoPower[1] and Iso-Dream[2], which are missing from related work. It would be better if the authors can discuss the differences between the proposed method and this type of work. - [1] INFOrmation Prioritization through EmPOWERment in Visual Model-Based RL. ICLR 2022. - [2] Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models. NeurIPS 2022. - The error bars of the baselines are not included. - The performance of TIA on variants of DMC sometimes encounters a sudden drop during training, which is worth a further check for the reason. - 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.

Questions

- I understand that blockwise identification refers to the existence of a one to one mapping between a late variable and its estimated value. However, I lacked a thorough understanding of how decoupling the four latent components is related to the concept of block-wise identifiability. Could you please explain further? Or what are the benefits of defining these components blockwise identifiable? - 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Negative societal impact is not discussed

Summary

The authors propose an alternative way to create world models in an RL system by separating them into blocks dependant on their casual effects on future observations, states and rewards. They theoretically show that under some assumptions it is possible to identify such classes of variables, even if specific instances cannot be identified. They go on to show how their approach can increase robustness and disentanglement in several control tasks.

Strengths

1. The proposed method is well motivated. I especially liked that instead of focusing on variable level disentanglement the authors have opted for block level disentanglement. This is a good way to relax the assumptions of the model without being completely unconstrained. 2. The theoretical motivation seems good (but see below). It is always a bonus when researchers can provide some theoretical justification for their work. 3. Experimental evidence is convincing. I especially liked the analysis of the disentanglement of structured datasets and the results look extremely promising.

Weaknesses

No serious weakness. I have some comments regarding how section 2 and 3 are written but otherwise it is a solid paper.

Questions

In general my questions revolve around the quality of the plots and some results that I think are missing. 1. 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. 2. 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

1. Section 2 could use a slight rewrite. Each of the propositions in lines 98 to 113 start the same way. Seems like it should be possible to just use one heading and enumerate the propositions afterward. 2. 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.