MIMEx: Intrinsic Rewards from Masked Input Modeling

TL;DR

MIMEx introduces a flexible framework for intrinsic rewards based on masked input modeling, improving exploration in high-dimensional environments by tuning mask distributions for better novelty estimation.

Contribution

The paper unifies existing intrinsic reward methods under a conditional prediction perspective and proposes MIMEx, a novel approach with tunable mask distributions for enhanced exploration.

Findings

MIMEx outperforms baselines on sparse-reward visuomotor tasks.

Flexible mask tuning improves exploration efficiency.

Unified view of intrinsic rewards as pseudo-likelihood estimation.

Abstract

Exploring in environments with high-dimensional observations is hard. One promising approach for exploration is to use intrinsic rewards, which often boils down to estimating "novelty" of states, transitions, or trajectories with deep networks. Prior works have shown that conditional prediction objectives such as masked autoencoding can be seen as stochastic estimation of pseudo-likelihood. We show how this perspective naturally leads to a unified view on existing intrinsic reward approaches: they are special cases of conditional prediction, where the estimation of novelty can be seen as pseudo-likelihood estimation with different mask distributions. From this view, we propose a general framework for deriving intrinsic rewards -- Masked Input Modeling for Exploration (MIMEx) -- where the mask distribution can be flexibly tuned to control the difficulty of the underlying conditional prediction task. We demonstrate that MIMEx can achieve superior results when compared against competitive baselines on a suite of challenging sparse-reward visuomotor tasks.