Timezone: »

Learning Robust Dynamics through Variational Sparse Gating
Arnav Kumar Jain · Shivakanth Sujit · Shruti Joshi · Vincent Michalski · Danijar Hafner · Samira Ebrahimi Kahou

Tue Nov 29 09:00 AM -- 11:00 AM (PST) @ Hall J #409

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.

Author Information

Arnav Kumar Jain (University de Montreal)
Shivakanth Sujit (École de technologie supérieure)
Shruti Joshi (Montreal Institute for Learning Algorithms, University of Montreal, Université de Montréal)
Vincent Michalski (Université de Montréal)
Danijar Hafner (Google)
Samira Ebrahimi Kahou (McGill University)

More from the Same Authors