---
title: One-Step FB Representation Learning
url: https://www.emergentmind.com/topics/one-step-forward-backward-representation-learning-one-step-fb
type: topic
---

# One-Step FB Representation Learning

One-Step Forward-Backward Representation Learning, usually abbreviated **one-step FB**, is a fixed-policy variant of forward-backward representation learning for reinforcement learning that learns a bilinear factorization of the **behavioral successor-measure ratio** and uses that factorization for **zero-shot policy improvement** on downstream rewards without updating network parameters [2602.11399]. In the formulation studied in an infinite-horizon discounted controlled Markov process (CMP) with state space $\mathcal S$, action space $\mathcal A$, initial distribution $p_0$, transition kernel $p(s' \mid s,a)$, and discount $\gamma \in [0,1)$, one-step FB is trained from offline data collected by a behavioral policy $\pi_\beta(a \mid s)$ and is designed to recover the family of behavioral value functions $\{Q_r^{\pi_\beta}\}_r$ rather than globally optimal policies for arbitrary rewards [2602.11399]. Its central claim is therefore narrower than the original forward-backward (FB) program: it performs **one step of greedy policy improvement** over the behavioral policy, not universal optimal control.

## 1. Historical origin and conceptual reframing

The immediate precursor to one-step FB is the FB framework introduced in “Learning One Representation to Optimize All Rewards” [2103.07945]. In that formulation, a reward-free MDP is represented through a forward map $F$, a backward map $B$, and a latent-indexed policy family $(\pi_z)_{z\in \mathbb R^d}$ such that the discounted successor measure under $\pi_z$ factorizes bilinearly with respect to a reference measure $\rho$. For a reward $r$, one defines a reward embedding
\[
z_R=\mathbb{E}_{(s,a)\sim \rho}[r(s,a)B(s,a)],
\]
and then obtains a policy directly by greedy action selection with respect to $F(s,a,z_R)^\top z_R$, with no planning [2103.07945]. The 2021 formulation states that, if training is perfect, the induced policy is provably optimal for any reward function [2103.07945].

One-step FB emerged from a critique of that universal-control interpretation. The 2026 analysis argues that the original FB objective is **circular**: $F$ determines the policy through an $\arg\max$ over $F(s,a,z)^\top z$, but the successor measure being fitted also depends on that induced policy [2602.11399]. In that sense, FB is learning a representation of a target whose dynamics depend on the representation itself. The resulting reformulation is explicit: once the limitations of universal optimal-control representations are taken seriously, the appropriate object to learn is not an optimal policy family for all rewards, but a representation of the **fixed behavioral successor measure** from which one can perform one step of policy improvement [2602.11399].

This reframing changes the meaning of “one-step.” In one-step FB, the method is not “one-step” because it uses a single gradient update; it is “one-step” because, after learning a reward-independent representation of $Q_r^{\pi_\beta}$ for all rewards $r$, it extracts a downstream policy by acting greedily with respect to that behavioral $Q$-function estimate [2602.11399]. A common misconception is therefore to read one-step FB as a weakened approximation to universal optimal control. The more precise interpretation is that it solves a different problem: **universal fixed-policy evaluation plus one-step policy improvement**.

## 2. Mathematical formulation

The mathematical backbone of both FB and one-step FB is the **successor measure**
\[
M^\pi(s_f,a_f \mid s,a)
= (1-\gamma)\,\delta(s_f,a_f \mid s,a)
+ \gamma \,\mathbb E_{\substack{s' \sim p(s' \mid s,a)\\ a' \sim \pi(a \mid s')}}\!\left[M^\pi(s_f,a_f \mid s',a')\right],
\]
which is the discounted occupancy of future state-action pairs under policy $\pi$ [2602.11399]. This object linearly induces action-values:
\[
Q_r^\pi(s,a)
=
\mathbb E_{(s_f,a_f)\sim M^\pi(s_f,a_f \mid s,a)}[r(s_f,a_f)].
\]
In the original FB construction, one attempts to factorize the successor-measure ratio of a latent-conditioned policy:
\[
\frac{M^\pi(s_f,a_f \mid s,a,z)}{\rho(s_f,a_f)}
=
F^\star(s,a,z)^\top B^\star(s_f,a_f),
\]
with policy extraction defined by
\[
\pi(a \mid s,z)
=
\delta\!\left(a \,\middle|\, \arg\max_{a \in \mathcal A} F^\star(s,a,z)^\top z \right).
\]
A downstream reward is then embedded as
\[
z_r
=
\mathbb E_{(s_f,a_f)\sim \rho}\!\left[B^\star(s_f,a_f)\, r(s_f,a_f)\right],
\]
and the ambition is that $\pi(\cdot \mid s,z_r)$ is optimal [2602.11399].

One-step FB fixes the policy target to the behavioral policy $\pi_\beta$ and removes the latent variable from the forward encoder input. At the ideal solution, it learns forward and backward representations $F_\beta$ and $B_\beta$ such that
\[
F_\beta^\star(s,a)^\top B_\beta^\star(s_f,a_f)
=
\frac{M^{\pi_\beta}(s_f,a_f \mid s,a)}{\rho(s_f,a_f)}.
\]
For a downstream reward $r$, one defines the reward embedding
\[
z_r^\beta
=
\mathbb E_{(s_f,a_f)\sim \rho}\!\left[B_\beta^\star(s_f,a_f)\, r(s_f,a_f)\right],
\]
which yields the behavioral action-value family
\[
Q_r^{\pi_\beta}(s,a)=F_\beta^\star(s,a)^\top z_r^\beta.
\]
The induced latent-conditioned policy is then
\[
\pi(a \mid s,z_r^\beta)
=
\arg\max_{a \in \mathcal A} F_\beta^\star(s,a)^\top z_r^\beta
=
\arg\max_{a \in \mathcal A} Q_r^{\pi_\beta}(s,a).
\]
This equality is the formal content of the method’s name: the policy is **exactly one step of greedy policy improvement over the behavioral $Q$-function $Q_r^{\pi_\beta}$** [2602.11399].

The distinction from the original FB program can be summarized as follows.

| Aspect | FB | One-step FB |
|---|---|---|
| Forward representation | $F(s,a,z)$ | $F_\beta(s,a)$ |
| Target successor measure | Latent-conditioned policy induced by the current representation | Fixed behavioral policy $\pi_\beta$ |
| Downstream interpretation | Claimed optimal control for arbitrary rewards | One step of greedy policy improvement over $Q_r^{\pi_\beta}$ |

## 3. Objective, optimization, and zero-shot adaptation

The one-step FB objective is derived by simplifying the original FB Bellman-residual objective. The 2026 analysis traces the FB loss to **least-squares importance fitting (LSIF)** for a density ratio $g(x)\approx p(x)/q(x)$, then instantiates the ratio bilinearly as
\[
g(s,a,z,s_f,a_f)=F(s,a,z)^\top B(s_f,a_f)
\]
with the successor measure as target and $\rho$ as anchor [2602.11399]. In the original formulation, bootstrapping with target networks yields a temporal-difference objective whose next action is sampled from the latent-conditioned policy induced by the current representation [2602.11399].

One-step FB keeps the same density-ratio perspective but freezes the next-action sampling distribution to the behavioral policy and removes the latent from the forward representation. Its temporal-difference objective is
\[
\begin{aligned}
\mathcal L_{\text{TD one-step FB}}(F_\beta,B_\beta)
&=
\frac12 \mathbb E \Big[
\big(
F_\beta(s,a)^\top B_\beta(s_f,a_f)
-
\gamma \bar F_\beta(s',a')^\top \bar B_\beta(s_f,a_f)
\big)^2
\Big] \\
&\quad
-
(1-\gamma)\,\mathbb E \left[ F_\beta(s,a)^\top B_\beta(s,a) \right],
\end{aligned}
\]
where expectation is taken over $(s,a,s',a')$ sampled from the offline dataset and $a'$ is sampled from $\pi_\beta$ [2602.11399]. The complete objective adds orthonormal regularization on the backward representation,
\[
\mathcal L_{\text{ortho}}(B_\beta),
\]
with coefficient $\lambda_{\text{ortho}}$ [2602.11399]. The paper states the total representation objective as a sum of the TD one-step FB term and the orthonormal regularizer [2602.11399].

Policy learning is decoupled from representation learning. In discrete action spaces, the latent-conditioned policy can be written analytically as a softmax over $F_\beta(s,a)^\top z$ with temperature $\tau_{\text{policy}}$ [2602.11399]. In continuous action spaces, the paper trains an explicit actor with reparameterized policy gradients and behavior-cloning regularization:
\[
\mathcal L(\pi)
=
-
\mathbb E\!\left[
F_\beta(s,a)^\top z + \lambda_{\text{BC}} \log \pi(a_\beta \mid s,z)
\right],
\]
using sampled latents $z \sim p_{\mathcal Z}$ during unsupervised pretraining [2602.11399]. Crucially, policy learning does not alter the fixed-policy target of the representation learner.

At test time, adaptation is fully zero-shot in the sense that **no network weights are updated**. One computes
\[
z_r^\beta
=
\mathbb E_{(s_f,a_f)\sim \rho} \left[B_\beta(s_f,a_f)\,r(s_f,a_f)\right],
\]
or an empirical estimate from a reward-labeled inference dataset, optionally with reward softmax reweighting
\[
\tilde r(s_i,a_i)=w_i r(s_i,a_i),
\qquad
w_i=\frac{\exp(\tau_{\text{reward}} r(s_i,a_i))}{\sum_j \exp(\tau_{\text{reward}} r(s_j,a_j))},
\]
and then executes $\pi(a \mid s, z_r^\beta)$ [2602.11399]. Relative to FB, the implementation change is minimal but consequential: remove $z$ from the forward encoder input, and in the representation loss sample $a'$ from the dataset or behavior policy rather than the current latent-conditioned policy [2602.11399].

## 4. Interpretive frameworks and theoretical guarantees

One-step FB is presented through four mutually reinforcing interpretations [2602.11399]. First, it is a **Bellman-style density-ratio learner**: rather than predicting scalar $Q$-values directly, it learns a reward-independent factorization of the behavioral successor-measure ratio. Second, it is a **fitted Q-evaluation analogue** for successor-measure ratios. Standard FQE learns $Q^{\pi_\beta}$ by applying a Bellman evaluation operator for a fixed policy; one-step FB learns a richer object from which $Q_r^{\pi_\beta}$ can later be recovered for any reward $r$ [2602.11399]. Third, it admits a **rank-matching or matrix-factorization view**. Once $\pi_\beta$ is fixed, the target matrix
\[
M^{\pi_\beta}\operatorname{diag}(\rho)^{-1}
\]
is fixed, and the learning problem becomes low-rank bilinear factorization. In the appendix, the paper gives an SVD interpretation: if
\[
M^{\pi_\beta}\operatorname{diag}(\rho)^{-1}=U_\beta \Sigma_\beta V_\beta^\top,
\]
then one exact choice is
\[
F_\beta^\star=U_\beta \Sigma_\beta,
\qquad
B_\beta^\star=V_\beta^\top,
\]
assuming sufficient dimension $d = |\mathcal S \times \mathcal A|$ [2602.11399].

The fourth interpretation concerns **contraction mapping**. For the original FB formulation, the paper defines an FB Bellman operator on the bilinear score $f(s,a,z)^\top b(s_f,a_f)$ and proves that it is **not** a $\gamma$-contraction under $L^p$ norms because the policy inside the operator depends on the same representation being updated [2602.11399]. As stated explicitly,
\[
\mathcal T_{\text{FB}} \text{ is not a } \gamma\text{-contraction.}
\]
This blocks direct Banach fixed-point arguments for convergence. One-step FB breaks that circular dependence by sampling the next action from the fixed behavior policy $\pi_\beta$, turning the learning problem into a fixed-policy evaluation and factorization problem [2602.11399].

The same analysis also sharpens the expressivity question that motivated the reformulation. In a discrete CMP with finite $\mathcal S,\mathcal A$, if exact “ground-truth” FB representations encode all latent-conditioned successor measures exactly and enable optimal adaptation for every reward, then necessarily
\[
d \ge |\mathcal S \times \mathcal A|,
\]
with
\[
\operatorname{rank}(F^\star_{\mathcal Z}) \ge |\mathcal S \times \mathcal A|,
\qquad
\operatorname{rank}(B^\star)=|\mathcal S \times \mathcal A|.
\]
For continuous problems, this implies that finite-dimensional exact FB cannot realize the original universal-control theory:
\[
|\mathcal S \times \mathcal A| \to \infty
\quad \Longrightarrow \quad
\text{no exact ground-truth FB with finite } d
\]
[2602.11399]. The paper further states that when $d < |\mathcal S \times \mathcal A|$, there exists a reward for which the FB-predicted optimal $Q$-value can be arbitrarily wrong [2602.11399]. This diagnosis is the formal basis for replacing universal optimal-control claims with the more limited but firmer one-step FB objective.

## 5. Empirical behavior and benchmark results

The most controlled empirical evidence comes from didactic discrete CMPs. In a three-state CMP and an additional five-state circular CMP, the paper reports that FB fails to converge to its claimed fixed point, while one-step FB converges essentially exactly to its own fixed point [2602.11399]. After $10^5$ gradient steps in the three-state example, FB still has successor-measure-ratio error $\epsilon_{\text{SMR}} = 4 \times 10^{-2}$, policy KL around $10^{-2}$, and equivariance error $\epsilon_{\text{equiv}} = 10^{-4}$, violating the affine reward-equivariance property that true $Q$-functions should satisfy [2602.11399]. By contrast, one-step FB drives all analogous errors below $10^{-7}$ within $4\times 10^4$ steps, with equivariance error
\[
\epsilon_{\text{equiv}} = 5 \times 10^{-9},
\]
which is the basis for the claim that one-step FB converges to errors $10^5$ smaller [2602.11399].

The large-scale evaluation covers **10 continuous-control domains**: $8$ state-based and $2$ image-based [2602.11399]. These include ExORL locomotion domains—walker, cheetah, quadruped, jaco—and OGBench domains—antmaze large navigate, antmaze teleport navigate, cube single play, scene play, visual cube single play, and visual scene play [2602.11399]. The abstract reports that one-step FB improves zero-shot performance by **$+24\%$ on average** [2602.11399]. The main table reports an average improvement of **$1.4\times$ over FB** [2602.11399].

| Domain | One-step FB | FB |
|---|---:|---:|
| walker | 379 | 400 |
| cheetah | 378 | 271 |
| quadruped | 645 | 246 |
| jaco | 22 | 10 |
| antmaze large navigate | 30 | 25 |
| antmaze teleport navigate | 11 | 16 |
| cube single play | 3 | 2 |
| scene play | 8 | 6 |
| visual cube single play | 14 | 12 |
| visual scene play | 16 | 13 |

These results support two distinct empirical claims [2602.11399]. First, one-step FB is best or near-best on **$6/10$ domains** and is especially strong in cheetah, quadruped, jaco, and image-based OGBench tasks. Second, the method does not dominate uniformly: walker and antmaze teleport navigate are counterexamples, and the paper notes that state-based OGBench tasks with sparse goal-indicator rewards are harder for one-step FB because such rewards may effectively induce only a single backward representation, favoring methods that explicitly learn goal distances like HILP or ICVF [2602.11399]. In additional online fine-tuning experiments, one-step FB yields **$+40\%$ higher sample efficiency than FB on average** and about **$2.25\times$ over training TD3 from scratch** [2602.11399].

A later spectral analysis gives a complementary explanation of why practical FB-style methods can remain difficult in continuous domains even after the one-step reformulation. “Spectral Alignment in Forward-Backward Representations via Temporal Abstraction” studies the low-rank bottleneck of FB representations and argues that continuous environments often induce successor representations with slow spectral decay, making low-rank bootstrapped learning unstable [2603.20103]. That work proposes **temporal abstraction via action repetition** as a low-pass spectral filter, with repeated-action abstraction shrinking subdominant singular modes according to
\[
\sigma_i(P^k)\le \sigma_i(P)^k,
\]
thereby reducing effective rank and stabilizing bootstrapped FB learning [2603.20103]. A plausible implication is that one-step FB addresses the policy-coupling pathology diagnosed in [2602.11399], while temporal abstraction addresses an additional **spectral mismatch** that persists in continuous deep function approximation [2603.20103].

## 6. Scope, limitations, and current interpretation

One-step FB explicitly trades away “optimal” zero-shot adaptation [2602.11399]. One step of policy improvement is not repeated policy improvement to optimality, and the method’s quality depends on the coverage of the offline behavioral policy. If $\pi_\beta$ never visits the regions required for good downstream behavior, then both $Q_r^{\pi_\beta}$ and its greedy improvement may be poor [2602.11399]. This dependence on behavioral coverage distinguishes one-step FB from the original universal-control aspiration of FB and is central to interpreting its zero-shot results.

The method also remains a **finite-dimensional low-rank factorization**. Its quality depends on the representation dimension $d$, the reward class, and environment structure [2602.11399]. The appendix reports sensitivity to $d$, $\lambda_{\text{ortho}}$, $\lambda_{\text{BC}}$, and the reward-weighting temperature $\tau_{\text{reward}}$ [2602.11399]. The theoretical statements are most explicit in discrete settings; continuous domains rely on neural parameterization and approximation [2602.11399]. This suggests that the theoretical simplification is substantial but not complete: the reformulated objective is better grounded than FB, yet still inherits the usual approximation issues of deep RL.

The current consensus suggested by the 2026 literature is correspondingly precise. One-step FB is best understood as a method that learns
\[
\frac{M^{\pi_\beta}(s_f,a_f \mid s,a)}{\rho(s_f,a_f)} \approx F_\beta(s,a)^\top B_\beta(s_f,a_f),
\]
so that for any reward $r$ it can recover
\[
Q_r^{\pi_\beta}(s,a)=F_\beta(s,a)^\top z_r^\beta
\]
and extract a downstream policy by greedy improvement with respect to that behavioral $Q$-function [2602.11399]. It is therefore not a universal optimal-control representation in the sense asserted for exact FB in [2103.07945], but a **stable form of universal fixed-policy evaluation plus one-step policy improvement** [2602.11399]. That narrower interpretation is the one supported simultaneously by the expressivity analysis, the contraction-mapping diagnosis, and the empirical comparisons now associated with the method [2602.11399].

Source: https://www.emergentmind.com/topics/one-step-forward-backward-representation-learning-one-step-fb