Simplifying Latent Dynamics with Softly State-Invariant World Models
This addresses the challenge of systematic action representation in world models for control and reinforcement learning, though it appears incremental as it builds on existing model classes.
The paper tackled the problem of world models not systematically representing action effects by introducing the Parsimonious Latent Space Model (PLSM), which regularizes latent dynamics to make actions more predictable, resulting in improved accuracy, generalization, and performance in downstream tasks.
To solve control problems via model-based reasoning or planning, an agent needs to know how its actions affect the state of the world. The actions an agent has at its disposal often change the state of the environment in systematic ways. However, existing techniques for world modelling do not guarantee that the effect of actions are represented in such systematic ways. We introduce the Parsimonious Latent Space Model (PLSM), a world model that regularizes the latent dynamics to make the effect of the agent's actions more predictable. Our approach minimizes the mutual information between latent states and the change that an action produces in the agent's latent state, in turn minimizing the dependence the state has on the dynamics. This makes the world model softly state-invariant. We combine PLSM with different model classes used for i) future latent state prediction, ii) planning, and iii) model-free reinforcement learning. We find that our regularization improves accuracy, generalization, and performance in downstream tasks, highlighting the importance of systematic treatment of actions in world models.