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 https://chongyi-zheng.github.io/onestep-fb.

• The paper presents rigorous theoretical limits of FB representation learning, showing that low-dimensional representations cannot universally optimize all rewards.
• It analyzes structural constraints, rank requirements, and the lack of contraction in FB methods, quantifying error lower bounds in optimal Q-value approximations.
• The introduction of one-step FB breaks circular dependencies, achieving up to 10^5 times lower Q prediction errors and enhanced sample efficiency.

Detailed Essay on "Can We Really Learn One Representation to Optimize All Rewards?" (2602.11399)

Problem Statement and Theoretical Motivation

The paper interrogates the foundational premise of zero-shot reinforcement learning (RL) via universal representation learning: given a controlled Markov process (CMP) and a yet-unknown reward function, is it possible to learn a single representation during unsupervised pre-training that enables optimal policy adaptation for arbitrary downstream rewards, without further fine-tuning? This is embodied in the forward-backward (FB) representation learning paradigm, which factorizes the successor measure as a bilinear product of state/action representations, presuming that such factorization allows direct and optimal policy adaptation for any reward.

A central theme is the theoretical demystification of FB’s objective: under what structural and rank constraints can a universal representation support all reward functions, and does the practical optimization of FB fulfill its theoretical promises in either discrete or continuous settings?

Figure 1: FB representation learning seeks to factorize successor measures for universal policy adaptation, but faces optimization challenges that the paper analyzes and addresses.

Analysis of FB Representation Learning: Structural Constraints and Objective

The paper formalizes the underlying mathematical structures of FB, introducing the notion of ground-truth forward and backward representations ($F^\star$, $B^\star$) and their ability to express occupancy measures as ratios over a marginal distribution. Through rigorous linear algebra, it is shown that:

• Dimensionality Constraint: The representation dimension must be at least the size of the state-action space ($d \geq |\mathcal{S} \times \mathcal{A}|$); otherwise, no universal representation can encode the full successor measure.
• Rank Matching: Both representations must satisfy specific rank conditions; full-rank backward representations are necessary, and forward representations must span the state-action space.
• Error Lower Bound: When using finite-dimensional representations ($d < |\mathcal{S} \times \mathcal{A}|$), there exist reward functions for which the induced Q-values incur arbitrarily large errors.

These constraints directly challenge previous bold claims that FB representations can approximate optimal Q-values for all rewards within bounded error.

FB’s Optimization Objective: LSIF and Lack of Contraction

FB’s representation learning is reinterpreted as a temporal-difference variant of least-squares importance fitting (LSIF) applied to estimating the occupancy ratio. The paper proves that FB’s Bellman updates do not constitute a $\gamma$-contraction, due to an inherent circular dependency between representations and latent-conditioned policy. Thus, classical fixed-point guarantees (Banach fixed-point theorem) do not apply and convergence may not be achievable.

Empirical Evaluation: Failure Modes and Convergence

The authors empirically test FB in analytically solvable CMPs (e.g., three-state, five-state circular MDPs). After extensive training (up to $10^5$ steps), FB fails to converge to the ground-truth occupancy ratio or Q-values, exhibiting high prediction error, policy KL, and violation of affine equivariance properties required of universal value functions.

Figure 2: FB representation learning in a three-state CMP confirms the theoretical prediction that FB fails to fully recover successor measures and ground-truth Q values.

Figure 3: One-step FB representations in a five-state circular CMP show exact fitting within $4 \times 10^4$ steps, validating the stronger convergence of the simplified approach.

One-Step FB: Breaking Circular Dependency for Robust Pre-Training

Motivated by the above limitations, the paper introduces "one-step forward-backward (FB)" representation learning: instead of fitting for latent-conditioned policies, the representations are learned for a fixed behavioral policy. This breaks the circular dependency, enabling stable learning via LSIF/Bellman updates. Policy adaptation proceeds via one-step improvement over behavioral Q-values—not full optimality, but empirically sufficient.

One-step FB is empirically shown to give zero-shot Q prediction errors up to $10^5$ times smaller than FB, and consistently achieves affine equivariance required of universal Q-value functions.

Benchmark Evaluation: Performance and Adaptivity

The simplified one-step FB is benchmarked against FB and other unsupervised representation RL methods (HILP, Laplacian, BYOL-$\gamma$, ICVF) on ExORL (state-based) and OGBench (state/image-based) domains.

Figure 4: The ExORL and OGBench domains used for large-scale evaluation.

One-step FB matches or surpasses state-of-the-art unsupervised RL methods in 6/10 domains, providing +24% improvement in zero-shot performance and +1.4x over FB on average. Critically, the method enables sample-efficient initialization for online fine-tuning, competitive across both dense and sparse reward regimes.

Figure 5: One-step FB delivers consistently higher efficiency during fine-tuning compared to FB (+40% average sample efficiency).

Ablations and Domain Sensitivity

Extensive ablation experiments confirm that performance depends on appropriate hyperparameter choices: behavioral cloning regularization, backward orthonormalization, reward weighting temperature, and representation dimension, with domain-specific tuning yielding optimal results. The method is robust across MDP structures and reward sparsity.

Figure 6: Learning rate ablations show no explanation for FB’s failure mode; only one-step FB enjoys convergence.

Figure 7: Policy temperature ablation highlights the sensitivity of FB’s optimization to temperature, reaffirming the robustness of one-step FB.

Implications and Future Directions

The theoretical analysis establishes limits on the universality of representation learning in RL: unless representations are full-rank and match the state-action space, universal optimal control for all rewards is impossible. One-step FB trades universal optimality for tractability and convergence, but empirical results show the convergence property often translates to higher practical performance—especially in offline RL and foundational RL models.

These findings have broad implications for behavioral foundation models, prompting further research into representations for continuous and infinite spaces, as well as alternative fixed-point analysis tools such as Lefschetz theory and Lyapunov stability. The method opens new possibilities for sample-efficient transfer, offline-to-online adaptation, and scalable unsupervised RL pre-training.

Conclusion

The paper provides a rigorous theoretical and empirical investigation into FB representation learning for zero-shot RL, demonstrating structural impossibilities in the original FB paradigm and motivating a tractable, convergent variant (one-step FB). One-step FB achieves superior numerical results, stable convergence, and competitive performance on benchmark domains, positioning it as a robust practical method for behavioral foundation model pre-training, albeit with limitations in universal reward optimality. The work prompts further exploration into representation learning for continuous spaces and alternative approaches to universal RL adaptation.