The Impact of Negative Sampling on Contrastive Structured World Models

TL;DR

This paper investigates how different negative sampling strategies in contrastive learning significantly affect the performance of structured world models, demonstrating improvements through leveraging temporal correlations and dataset diversity.

Contribution

It reveals the critical impact of negative sampling choices on contrastive world models and introduces methods to enhance performance by exploiting temporal correlations.

Findings

Leveraging time step correlations doubles model performance.

Negative sampling strategies drastically influence contrastive learning outcomes.

Diverse datasets enable more robust contrastive world models.

Abstract

World models trained by contrastive learning are a compelling alternative to autoencoder-based world models, which learn by reconstructing pixel states. In this paper, we describe three cases where small changes in how we sample negative states in the contrastive loss lead to drastic changes in model performance. In previously studied Atari datasets, we show that leveraging time step correlations can double the performance of the Contrastive Structured World Model. We also collect a full version of the datasets to study contrastive learning under a more diverse set of experiences.