One-Step FB Representation Learning

• One-Step FB is a representation learning approach that factorizes the behavioral successor-measure ratio for reinforcement learning to enable fixed-policy evaluation plus one-step policy improvement.
• It decouples representation learning from policy updating by using a bilinear factorization, thereby stabilizing zero-shot adaptation without updating network parameters.
• Empirical studies show that one-step FB converges with significantly reduced errors and improves zero-shot performance compared to the original forward-backward framework.

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 (Zheng et al., 11 Feb 2026). 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 (Zheng et al., 11 Feb 2026). 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” (Touati et al., 2021). 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 $\mathcal A$0 factorizes bilinearly with respect to a reference measure $\mathcal A$1. For a reward $\mathcal A$2, one defines a reward embedding

$\mathcal A$3

and then obtains a policy directly by greedy action selection with respect to $\mathcal A$4, with no planning (Touati et al., 2021). The 2021 formulation states that, if training is perfect, the induced policy is provably optimal for any reward function (Touati et al., 2021).

One-step FB emerged from a critique of that universal-control interpretation. The 2026 analysis argues that the original FB objective is circular: $\mathcal A$5 determines the policy through an $\mathcal A$6 over $\mathcal A$7, but the successor measure being fitted also depends on that induced policy (Zheng et al., 11 Feb 2026). 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 (Zheng et al., 11 Feb 2026).

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 $\mathcal A$8 for all rewards $\mathcal A$9, it extracts a downstream policy by acting greedily with respect to that behavioral $p_0$0-function estimate (Zheng et al., 11 Feb 2026). 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

$p_0$1

which is the discounted occupancy of future state-action pairs under policy $p_0$2 (Zheng et al., 11 Feb 2026). This object linearly induces action-values: $p_0$3 In the original FB construction, one attempts to factorize the successor-measure ratio of a latent-conditioned policy: $p_0$4 with policy extraction defined by

$p_0$5

A downstream reward is then embedded as

$p_0$6

and the ambition is that $p_0$7 is optimal (Zheng et al., 11 Feb 2026).

One-step FB fixes the policy target to the behavioral policy $p_0$8 and removes the latent variable from the forward encoder input. At the ideal solution, it learns forward and backward representations $p_0$9 and $p(s' \mid s,a)$0 such that

$p(s' \mid s,a)$1

For a downstream reward $p(s' \mid s,a)$2, one defines the reward embedding

$p(s' \mid s,a)$3

which yields the behavioral action-value family

$p(s' \mid s,a)$4

The induced latent-conditioned policy is then

$p(s' \mid s,a)$5

This equality is the formal content of the method’s name: the policy is exactly one step of greedy policy improvement over the behavioral $p(s' \mid s,a)$6-function $p(s' \mid s,a)$7 (Zheng et al., 11 Feb 2026).

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

Aspect	FB	One-step FB
Forward representation	$p(s' \mid s,a)$8	$p(s' \mid s,a)$9
Target successor measure	Latent-conditioned policy induced by the current representation	Fixed behavioral policy $\gamma \in [0,1)$0
Downstream interpretation	Claimed optimal control for arbitrary rewards	One step of greedy policy improvement over $\gamma \in [0,1)$1

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 $\gamma \in [0,1)$2, then instantiates the ratio bilinearly as

$\gamma \in [0,1)$3

with the successor measure as target and $\gamma \in [0,1)$4 as anchor (Zheng et al., 11 Feb 2026). 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 (Zheng et al., 11 Feb 2026).

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

$\gamma \in [0,1)$5

where expectation is taken over $\gamma \in [0,1)$6 sampled from the offline dataset and $\gamma \in [0,1)$7 is sampled from $\gamma \in [0,1)$8 (Zheng et al., 11 Feb 2026). The complete objective adds orthonormal regularization on the backward representation,

$\gamma \in [0,1)$9

with coefficient $\pi_\beta(a \mid s)$0 (Zheng et al., 11 Feb 2026). The paper states the total representation objective as a sum of the TD one-step FB term and the orthonormal regularizer (Zheng et al., 11 Feb 2026).

Policy learning is decoupled from representation learning. In discrete action spaces, the latent-conditioned policy can be written analytically as a softmax over $\pi_\beta(a \mid s)$1 with temperature $\pi_\beta(a \mid s)$2 (Zheng et al., 11 Feb 2026). In continuous action spaces, the paper trains an explicit actor with reparameterized policy gradients and behavior-cloning regularization: $\pi_\beta(a \mid s)$3 using sampled latents $\pi_\beta(a \mid s)$4 during unsupervised pretraining (Zheng et al., 11 Feb 2026). 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

$\pi_\beta(a \mid s)$5

or an empirical estimate from a reward-labeled inference dataset, optionally with reward softmax reweighting

$\pi_\beta(a \mid s)$6

and then executes $\pi_\beta(a \mid s)$7 (Zheng et al., 11 Feb 2026). Relative to FB, the implementation change is minimal but consequential: remove $\pi_\beta(a \mid s)$8 from the forward encoder input, and in the representation loss sample $\pi_\beta(a \mid s)$9 from the dataset or behavior policy rather than the current latent-conditioned policy (Zheng et al., 11 Feb 2026).

4. Interpretive frameworks and theoretical guarantees

One-step FB is presented through four mutually reinforcing interpretations (Zheng et al., 11 Feb 2026). First, it is a Bellman-style density-ratio learner: rather than predicting scalar $\{Q_r^{\pi_\beta}\}_r$0-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_r^{\pi_\beta}\}_r$1 by applying a Bellman evaluation operator for a fixed policy; one-step FB learns a richer object from which $\{Q_r^{\pi_\beta}\}_r$2 can later be recovered for any reward $\{Q_r^{\pi_\beta}\}_r$3 (Zheng et al., 11 Feb 2026). Third, it admits a rank-matching or matrix-factorization view. Once $\{Q_r^{\pi_\beta}\}_r$4 is fixed, the target matrix

$\{Q_r^{\pi_\beta}\}_r$5

is fixed, and the learning problem becomes low-rank bilinear factorization. In the appendix, the paper gives an SVD interpretation: if

$\{Q_r^{\pi_\beta}\}_r$6

then one exact choice is

$\{Q_r^{\pi_\beta}\}_r$7

assuming sufficient dimension $\{Q_r^{\pi_\beta}\}_r$8 (Zheng et al., 11 Feb 2026).

The fourth interpretation concerns contraction mapping. For the original FB formulation, the paper defines an FB Bellman operator on the bilinear score $\{Q_r^{\pi_\beta}\}_r$9 and proves that it is not a $F$0-contraction under $F$1 norms because the policy inside the operator depends on the same representation being updated (Zheng et al., 11 Feb 2026). As stated explicitly,

$F$2

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 $F$3, turning the learning problem into a fixed-policy evaluation and factorization problem (Zheng et al., 11 Feb 2026).

The same analysis also sharpens the expressivity question that motivated the reformulation. In a discrete CMP with finite $F$4, if exact “ground-truth” FB representations encode all latent-conditioned successor measures exactly and enable optimal adaptation for every reward, then necessarily

$F$5

with

$F$6

For continuous problems, this implies that finite-dimensional exact FB cannot realize the original universal-control theory: $F$7 (Zheng et al., 11 Feb 2026). The paper further states that when $F$8, there exists a reward for which the FB-predicted optimal $F$9-value can be arbitrarily wrong (Zheng et al., 11 Feb 2026). 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 (Zheng et al., 11 Feb 2026). After $B$0 gradient steps in the three-state example, FB still has successor-measure-ratio error $B$1, policy KL around $B$2, and equivariance error $B$3, violating the affine reward-equivariance property that true $B$4-functions should satisfy (Zheng et al., 11 Feb 2026). By contrast, one-step FB drives all analogous errors below $B$5 within $B$6 steps, with equivariance error

$B$7

which is the basis for the claim that one-step FB converges to errors $B$8 smaller (Zheng et al., 11 Feb 2026).

The large-scale evaluation covers 10 continuous-control domains: $B$9 state-based and $(\pi_z)_{z\in \mathbb R^d}$0 image-based (Zheng et al., 11 Feb 2026). 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 (Zheng et al., 11 Feb 2026). The abstract reports that one-step FB improves zero-shot performance by $(\pi_z)_{z\in \mathbb R^d}$1 on average (Zheng et al., 11 Feb 2026). The main table reports an average improvement of $(\pi_z)_{z\in \mathbb R^d}$2 over FB (Zheng et al., 11 Feb 2026).

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 (Zheng et al., 11 Feb 2026). First, one-step FB is best or near-best on $(\pi_z)_{z\in \mathbb R^d}$3 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 (Zheng et al., 11 Feb 2026). In additional online fine-tuning experiments, one-step FB yields $(\pi_z)_{z\in \mathbb R^d}$4 higher sample efficiency than FB on average and about $(\pi_z)_{z\in \mathbb R^d}$5 over training TD3 from scratch (Zheng et al., 11 Feb 2026).

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 (Azad et al., 20 Mar 2026). That work proposes temporal abstraction via action repetition as a low-pass spectral filter, with repeated-action abstraction shrinking subdominant singular modes according to

$(\pi_z)_{z\in \mathbb R^d}$6

thereby reducing effective rank and stabilizing bootstrapped FB learning (Azad et al., 20 Mar 2026). A plausible implication is that one-step FB addresses the policy-coupling pathology diagnosed in (Zheng et al., 11 Feb 2026), while temporal abstraction addresses an additional spectral mismatch that persists in continuous deep function approximation (Azad et al., 20 Mar 2026).

6. Scope, limitations, and current interpretation

One-step FB explicitly trades away “optimal” zero-shot adaptation (Zheng et al., 11 Feb 2026). 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_z)_{z\in \mathbb R^d}$7 never visits the regions required for good downstream behavior, then both $(\pi_z)_{z\in \mathbb R^d}$8 and its greedy improvement may be poor (Zheng et al., 11 Feb 2026). 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 $(\pi_z)_{z\in \mathbb R^d}$9, the reward class, and environment structure (Zheng et al., 11 Feb 2026). The appendix reports sensitivity to $\mathcal A$00, $\mathcal A$01, $\mathcal A$02, and the reward-weighting temperature $\mathcal A$03 (Zheng et al., 11 Feb 2026). The theoretical statements are most explicit in discrete settings; continuous domains rely on neural parameterization and approximation (Zheng et al., 11 Feb 2026). 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

$\mathcal A$04

so that for any reward $\mathcal A$05 it can recover

$\mathcal A$06

and extract a downstream policy by greedy improvement with respect to that behavioral $\mathcal A$07-function (Zheng et al., 11 Feb 2026). It is therefore not a universal optimal-control representation in the sense asserted for exact FB in (Touati et al., 2021), but a stable form of universal fixed-policy evaluation plus one-step policy improvement (Zheng et al., 11 Feb 2026). 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 (Zheng et al., 11 Feb 2026).