[2602.11399] Can We Really Learn One Representation to Optimize All Rewards?
Abstract:As machine learning has moved towards leveraging large models as priors for downstream tasks, the community has debated the right form of prior for solving reinforcement learning (RL) problems. If one were to try to prefetch as much computation as possible, they would attempt to learn a prior over the policies for some yet-to-be-determined reward function. Recent work (forward-backward (FB) representation learning) has tried this, arguing that an unsupervised representation learning procedure can enable optimal control over arbitrary rewards without further fine-tuning. However, FB's training objective and learning behavior remain mysterious. In this paper, we demystify FB by clarifying when such representations can exist, what its objective optimizes, and how it converges in practice. We draw connections with rank matching, fitted Q-evaluation, and contraction mapping. Our analysis suggests a simplified unsupervised pre-training method for RL that, instead of enabling optimal control, performs one step of policy improvement. We call our proposed method $\textbf{one-step forward-backward representation learning (one-step FB)}$. Experiments in didactic settings, as well as in $10$ state-based and image-based continuous control domains, demonstrate that one-step FB converges to errors $10^5$ smaller and improves zero-shot performance by $+24\%$ on average. Our project website is available at this https URL.
Can We Really Learn One Representation to Optimize All Rewards? Chongyi Zheng∗1 Royina Karegoudra Jayanth∗1 Benjamin Eysenbach1 1 Princeton University chongyiz@princeton.edu rj5498@princeton.edu…
saved by
related reading
- [2507.13181] Spectral Bellman Method: Unifying Representation and Exploration in RLarxiv.org
- [2606.05555] Representation Learning Enables Scalable Multitask Deep Reinforcement Learningarxiv.org
- Progressive Point Matchingprestonfu.com
- A Gallery of Methods Beyond RL — Part I: Sampling Methodsshengyu-feng.github.io
- RLHF | John Lambertjohnwlambert.github.io
- State of RL for reasoning LLMs | A. Weersaweers.de
- [2305.18290] Direct Preference Optimization: Your Language Model is Secretly a Reward Modelarxiv.org
- 2103.04551arxiv.org
- A Free Lunch from the Noise:Provable and Practical Exploration for Representation Learningarxiv.org
- Contrastive Learning as Goal-Conditioned Reinforcement Learningarxiv.org
- RL's Razor: Why Online Reinforcement Learning Forgets Lessarxiv.org
- RL is even more information inefficient than you thoughtdwarkesh.com