Factorized World Models for Learning Causal Relationships

Аннотация

World models serve as a powerful framework for model-based reinforcement learn-ing, 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 achallenge.  In this work, we propose CEMA, a structured world model with factorizedlatent state capable of modeling sparse interaction, with non-zero components corre-sponding to events of interest.  This is possible due to a separate state and dynamicsof three components:  the actor, the object of manipulation, the latent influence factorbetween these two states.  In multitask setting, we analyze the mutual information ofthe hierarchical latent states to show how the model can represent sparse updates anddirectly model the causal influence of the robot on the object.