ReCoRe: Regularized Contrastive Representation Learning of World Model

TL;DR

The paper tackles sample inefficiency and poor out-of-distribution generalization in visual navigation with RL by introducing ReCoRe, a world-model framework that learns invariant representations through contrastive learning while enforcing invariance via an intervention-invariant regularizer implemented as an auxiliary task such as depth prediction. It decouples world-model learning from controller optimization, optimizing a joint loss $L_{WM} = E_p \left( \sum_t ( L^q_t + L^d_t + L^r_t + \beta L^{KL}_t ) \right)$ to encourage stable latent dynamics and robust feature representations. Empirical results on iGibson OoD generalization, sim-to-real transfer on Gibson, and DMControl demonstrate that ReCoRe outperforms state-of-the-art model-based and model-free baselines, with the depth-based regularizer proving crucial to prevent feature collapse under augmentation. The work shows significant practical impact for robust, sample-efficient navigation in real robots and suggests a flexible framework for incorporating other invariant auxiliary tasks across tasks and environments.

Abstract

While recent model-free Reinforcement Learning (RL) methods have demonstrated human-level effectiveness in gaming environments, their success in everyday tasks like visual navigation has been limited, particularly under significant appearance variations. This limitation arises from (i) poor sample efficiency and (ii) over-fitting to training scenarios. To address these challenges, we present a world model that learns invariant features using (i) contrastive unsupervised learning and (ii) an intervention-invariant regularizer. Learning an explicit representation of the world dynamics i.e. a world model, improves sample efficiency while contrastive learning implicitly enforces learning of invariant features, which improves generalization. However, the naïve integration of contrastive loss to world models is not good enough, as world-model-based RL methods independently optimize representation learning and agent policy. To overcome this issue, we propose an intervention-invariant regularizer in the form of an auxiliary task such as depth prediction, image denoising, image segmentation, etc., that explicitly enforces invariance to style interventions. Our method outperforms current state-of-the-art model-based and model-free RL methods and significantly improves on out-of-distribution point navigation tasks evaluated on the iGibson benchmark. With only visual observations, we further demonstrate that our approach outperforms recent language-guided foundation models for point navigation, which is essential for deployment on robots with limited computation capabilities. Finally, we demonstrate that our proposed model excels at the sim-to-real transfer of its perception module on the Gibson benchmark.

Figures (3)

• Figure 1: Intervention invariant regularizer is applied in addition to the contrastive loss in world-model-based RL. ReCoRe learns robust representations that are invariant to out-of-distribution appearance variations which help in the generalization of downstream tasks (such as navigation). Notice the consistent depth predictions despite texture variations of an iGibson evaluation scene.
• Figure 1: The out-of-distribution generalization tests of proposed ReCoRe on held-out scenes and visual textures from iGibson 1.0. Green circle is a random PointGoal, blue circle is a random starting point and blue line represents the travel path of the Turtlebot robot.
• Figure 2: Flow diagram of proposed Regularized Contrastive Representation learning (ReCoRe) of World Model. It consists of four components: (i) invariant representation learning module, (ii) intervention-invariant regularizer, (iii) latent dynamics model, and (iv) actor-critic controller. The invariant representation learning module utilizes data augmentation and contrastive learning for invariant features abstraction ($\tilde{s}_t$) from image observations ($o_t$). The latent dynamics model employs a recurrent neural network with deterministic hidden states ($h_t$) to predict the stochastic latent prior states ($\widehat{s}_t$), and corresponding rewards ($\widehat{r}_t$) from the posterior ($s_t$). Intervention invariant regularizer considers an auxiliary task (here depth prediction i.e. $p_{\theta}(\widehat{d}_t \vert s_t,h_t)$) invariant to data augmentation (here texture variations) which prevents feature collapse in training the world-model with contrastive learning. The controller maximizes the expected rewards of the action using an actor critic approach. In addition to being sample efficient, the proposed approach is more robust to out-of-distribution and sim-to-real generalization, since the controller is learned separately using invariant states of the environment.