Regularized Latent Dynamics Prediction is a Strong Baseline For Behavioral Foundation Models

TL;DR

This paper introduces RLDP, a simple regularization method for latent state prediction that maintains feature diversity and outperforms complex methods in zero-shot reinforcement learning, especially in low-coverage scenarios.

Contribution

RLDP demonstrates that a straightforward orthogonality regularization can match or surpass complex representation learning methods for zero-shot RL.

Findings

RLDP matches or exceeds state-of-the-art in zero-shot RL.

RLDP performs well in low-coverage data scenarios.

Complex objectives are not necessary for effective latent state prediction.

Abstract

Behavioral Foundation Models (BFMs) produce agents with the capability to adapt to any unknown reward or task. These methods, however, are only able to produce near-optimal policies for the reward functions that are in the span of some pre-existing state features, making the choice of state features crucial to the expressivity of the BFM. As a result, BFMs are trained using a variety of complex objectives and require sufficient dataset coverage, to train task-useful spanning features. In this work, we examine the question: are these complex representation learning objectives necessary for zero-shot RL? Specifically, we revisit the objective of self-supervised next-state prediction in latent space for state feature learning, but observe that such an objective alone is prone to increasing state-feature similarity, and subsequently reducing span. We propose an approach, Regularized Latent Dynamics Prediction (RLDP), that adds a simple orthogonality regularization to maintain feature diversity and can match or surpass state-of-the-art complex representation learning methods for zero-shot RL. Furthermore, we empirically show that prior approaches perform poorly in low-coverage scenarios where RLDP still succeeds.