# Learning One Representation to Optimize All Rewards

arXiv: https://arxiv.org/abs/2103.07945 · HTML: https://mlsanity.com/paper/2103.07945/

## Key facts

- Cited by 115 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 42 references that are in the map.
- First posted on arXiv 14 Mar 2021.

## Builds on (most cited first, 13 of 13)

- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [Multi-Goal Reinforcement Learning: Challenging Robotics Environments and Request for Research](https://mlsanity.com/paper/1802.09464/) (2018)
- [Reward-Free Exploration for Reinforcement Learning](https://mlsanity.com/paper/2002.02794/) (2020)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [Fast Task Inference with Variational Intrinsic Successor Features](https://mlsanity.com/paper/1906.05030/) (2019)
- [Universal Successor Features Approximators](https://mlsanity.com/paper/1812.07626/) (2018)
- [C-Learning: Learning to Achieve Goals via Recursive Classification](https://mlsanity.com/paper/2011.08909/) (2020)
- [Learning Successor States and Goal-Dependent Values: A Mathematical Viewpoint](https://mlsanity.com/paper/2101.07123/) (2021)
- [Deep Reinforcement Learning with Successor Features for Navigation across Similar Environments](https://mlsanity.com/paper/1612.05533/) (2016)
- [Self-Correcting Models for Model-Based Reinforcement Learning](https://mlsanity.com/paper/1612.06018/) (2016)
- [Hindsight policy gradients](https://mlsanity.com/paper/1711.06006/) (2017)
- [Disentangled Cumulants Help Successor Representations Transfer to New Tasks](https://mlsanity.com/paper/1911.10866/) (2019)

## Built on by (most cited first, 30 of 115)

- [Foundation Policies with Hilbert Representations](https://mlsanity.com/paper/2402.15567/) (2024)
- [Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models](https://mlsanity.com/paper/2502.14819/) (2025)
- [METRA: Scalable Unsupervised RL with Metric-Aware Abstraction](https://mlsanity.com/paper/2310.08887/) (2023)
- [Does Zero-Shot Reinforcement Learning Exist?](https://mlsanity.com/paper/2209.14935/) (2022)
- [Reinforcement Learning from Passive Data via Latent Intentions](https://mlsanity.com/paper/2304.04782/) (2023)
- [Diffusion Guidance Is a Controllable Policy Improvement Operator](https://mlsanity.com/paper/2505.23458/) (2025)
- [Learning Temporal Distances: Contrastive Successor Features Can Provide a Metric Structure for Decision-Making](https://mlsanity.com/paper/2406.17098/) (2024)
- [Zero-Shot Whole-Body Humanoid Control via Behavioral Foundation Models](https://mlsanity.com/paper/2504.11054/) (2025)
- [Controllability-Aware Unsupervised Skill Discovery](https://mlsanity.com/paper/2302.05103/) (2023)
- [Understanding Self-Predictive Learning for Reinforcement Learning](https://mlsanity.com/paper/2212.03319/) (2022)
- [Proto-Value Networks: Scaling Representation Learning with Auxiliary Tasks](https://mlsanity.com/paper/2304.12567/) (2023)
- [Provable Benefits of Representational Transfer in Reinforcement Learning](https://mlsanity.com/paper/2205.14571/) (2022)
- [Spectral Decomposition Representation for Reinforcement Learning](https://mlsanity.com/paper/2208.09515/) (2022)
- [Contrastive Difference Predictive Coding](https://mlsanity.com/paper/2310.20141/) (2023)
- [Multi-agent Reinforcement Learning: A Comprehensive Survey](https://mlsanity.com/paper/2312.10256/) (2023)
- [On the Importance of Exploration for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2306.05483/) (2023)
- [Unsupervised Zero-Shot Reinforcement Learning via Functional Reward Encodings](https://mlsanity.com/paper/2402.17135/) (2024)
- [The Generalization Gap in Offline Reinforcement Learning](https://mlsanity.com/paper/2312.05742/) (2023)
- [Temporal Difference Flows](https://mlsanity.com/paper/2503.09817/) (2025)
- [Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching from Offline Data](https://mlsanity.com/paper/2306.03346/) (2023)
- [Maximum State Entropy Exploration using Predecessor and Successor Representations](https://mlsanity.com/paper/2306.14808/) (2023)
- [TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning](https://mlsanity.com/paper/2510.00739/) (2025)
- [Zero-Shot Reinforcement Learning from Low Quality Data](https://mlsanity.com/paper/2309.15178/) (2023)
- [Fast Adaptation with Behavioral Foundation Models](https://mlsanity.com/paper/2504.07896/) (2025)
- [Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL](https://mlsanity.com/paper/2106.02193/) (2021)
- [Proto Successor Measure: Representing the Behavior Space of an RL Agent](https://mlsanity.com/paper/2411.19418/) (2024)
- [Zero-Shot Reinforcement Learning via Function Encoders](https://mlsanity.com/paper/2401.17173/) (2024)
- [Improving Zero-shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions](https://mlsanity.com/paper/2111.14629/) (2021)
- [A Distributional Analogue to the Successor Representation](https://mlsanity.com/paper/2402.08530/) (2024)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)

## Abstract (arXiv)

We introduce the forward-backward (FB) representation of the dynamics of a reward-free Markov decision process. It provides explicit near-optimal policies for any reward specified a posteriori. During an unsupervised phase, we use reward-free interactions with the environment to learn two representations via off-the-shelf deep learning methods and temporal difference (TD) learning. In the test phase, a reward representation is estimated either from observations or an explicit reward description (e.g., a target state). The optimal policy for that reward is directly obtained from these representations, with no planning. We assume access to an exploration scheme or replay buffer for the first phase. The corresponding unsupervised loss is well-principled: if training is perfect, the policies obtained are provably optimal for any reward function. With imperfect training, the sub-optimality is proportional to the unsupervised approximation error. The FB representation learns long-range relationships between states and actions, via a predictive occupancy map, without having to synthesize states as in model-based approaches. This is a step towards learning controllable agents in arbitrary black-box stochastic environments. This approach compares well to goal-oriented RL algorithms on discrete and continuous mazes, pixel-based MsPacman, and the FetchReach virtual robot arm. We also illustrate how the agent can immediately adapt to new tasks beyond goal-oriented RL.