Simplifying Latent Dynamics with Softly State-Invariant World Models

  • 2024-11-01 16:55:34
  • Tankred Saanum, Peter Dayan, Eric Schulz
  • 0

Abstract

To solve control problems via model-based reasoning or planning, an agentneeds to know how its actions affect the state of the world. The actions anagent has at its disposal often change the state of the environment insystematic ways. However, existing techniques for world modelling do notguarantee that the effect of actions are represented in such systematic ways.We introduce the Parsimonious Latent Space Model (PLSM), a world model thatregularizes the latent dynamics to make the effect of the agent's actions morepredictable. Our approach minimizes the mutual information between latentstates and the change that an action produces in the agent's latent state, inturn minimizing the dependence the state has on the dynamics. This makes theworld model softly state-invariant. We combine PLSM with different modelclasses used for i) future latent state prediction, ii) planning, and iii)model-free reinforcement learning. We find that our regularization improvesaccuracy, generalization, and performance in downstream tasks, highlighting theimportance of systematic treatment of actions in world models.

 

Quick Read (beta)

loading the full paper ...