Probabilistic World Modeling with Asymmetric Distance Measure
This work addresses the challenge of representation learning for planning in stochastic worlds, which is incremental as it builds on contrastive learning by introducing asymmetry for improved reasoning.
The paper tackles the problem of learning representations for planning in stochastic environments by proposing an asymmetric distance function that reflects state reachability, enabling the discovery of geometrically salient states as subgoals to break down long-horizon tasks, and demonstrates effectiveness in gridworld environments.
Representation learning is a fundamental task in machine learning, aiming at uncovering structures from data to facilitate subsequent tasks. However, what is a good representation for planning and reasoning in a stochastic world remains an open problem. In this work, we posit that learning a distance function is essential to allow planning and reasoning in the representation space. We show that a geometric abstraction of the probabilistic world dynamics can be embedded into the representation space through asymmetric contrastive learning. Unlike previous approaches that focus on learning mutual similarity or compatibility measures, we instead learn an asymmetric similarity function that reflects the state reachability and allows multi-way probabilistic inference. Moreover, by conditioning on a common reference state (e.g. the observer's current state), the learned representation space allows us to discover the geometrically salient states that only a handful of paths can lead through. These states can naturally serve as subgoals to break down long-horizon planning tasks. We evaluate our method in gridworld environments with various layouts and demonstrate its effectiveness in discovering the subgoals.