Learning Robust Dynamics through Variational Sparse Gating

TL;DR

This paper introduces Variational Sparse Gating (VSG) and its simplified version SVSG, which improve world models' ability to handle environments with many objects by sparsely updating latent features, leading to better performance in complex scenarios.

Contribution

The paper proposes VSG and SVSG, novel stochastic latent dynamics models that incorporate sparse interactions, enhancing world models' scalability to environments with numerous objects.

Findings

VSG and SVSG outperform prior models in BBS environment.

Models effectively handle large numbers of objects with sparse interactions.

Fully stochastic transition function improves robustness and scalability.

Abstract

Learning world models from their sensory inputs enables agents to plan for actions by imagining their future outcomes. World models have previously been shown to improve sample-efficiency in simulated environments with few objects, but have not yet been applied successfully to environments with many objects. In environments with many objects, often only a small number of them are moving or interacting at the same time. In this paper, we investigate integrating this inductive bias of sparse interactions into the latent dynamics of world models trained from pixels. First, we introduce Variational Sparse Gating (VSG), a latent dynamics model that updates its feature dimensions sparsely through stochastic binary gates. Moreover, we propose a simplified architecture Simple Variational Sparse Gating (SVSG) that removes the deterministic pathway of previous models, resulting in a fully stochastic transition function that leverages the VSG mechanism. We evaluate the two model architectures in the BringBackShapes (BBS) environment that features a large number of moving objects and partial observability, demonstrating clear improvements over prior models.