Factorized World Models for Learning Causal Relationships

Artem Zholus, Yaroslav Ivchenkov, Aleksandr Panov

Keywords: reinforcement learning, world models, object-centric RL, causal representation learning

TL;DR: We present a structured world model that learns to represent the cause-effect relationship between the actuator and the object of manipulation.

Abstract: World models serve as a powerful framework for model-based reinforcement learning, and they can greatly benefit from the shared structure of the world environments. However, learning the high-level causal influence of objects on each other remains a challenge. In this work, we propose CEMA, a structured world model with factorized latent state capable of modeling sparse interaction. This is possible due to a separate state and dynamics of three components: the actor, the object of manipulation, and the latent influence factor between these two states. In multitask setting, we analyze the mutual information of the hierarchical latent states to show how the model can represent sparse updates and directly model the causal influence of the robot on the object.