PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning One Representation to Optimize All Rewards

Ahmed Touati, Yann Ollivier

arXiv:2103.07945v3cs.LGcs.AImath.OC

TL;DR

The paper asks whether a reward-free MDP can be summarized so that policies for later-specified rewards can be obtained without planning. It learns forward-backward representations of long-term reachability and reports comparisons with goal-oriented methods across several environments, alongside qualitative results for more general rewards.

  • Problem

    The paper asks whether a compact object learned without reward information can provide policies for arbitrary later-specified rewards, while goal-oriented methods are limited to rewards specified in advance.

    Source: §1 · p. 1

  • Method

    The method learns forward-backward representations from reward-free interactions, then maps a specified reward into the representation space to obtain a policy without planning.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    The method compares well with goal-oriented methods for arbitrary-goal tasks across mazes, pixel-based MsPacman, and a virtual robotic arm, and it illustrates adaptation to more general rewards.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The learned representation summarizes reward-free MDP dynamics and provides near-optimal policies for later-specified rewards without planning.

    Source: §7 · p. 10

    The method assumes access to a good exploration strategy; variance may also be an issue in larger environments.

    Source: §4 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

The paper asks whether a compact object learned without reward information can provide policies for rewards specified later, without planning. It introduces forward-backward representations and reports results across maze, pixel-based, and robotic environments.

1 Introduction

  • Goal-oriented methods handle rewards specified in advance, while learned world models still require planning and accurate long-horizon state trajectories are difficult to synthesize.The paper contrasts its state-occupancy representation with finite user-supplied feature functions, which can require infinitely many features for continuous goal spaces.Source: §1 · pp. 1–2
  • The method is evaluated from mazes to pixel-based MsPacman and a virtual robotic arm, with quantitative comparisons for reaching arbitrary states and qualitative examples for more general rewards.The authors note it is not a substitute for HER and could be combined with HER to improve replay-buffer management.Source: §1 · p. 2
  • The paper proposes a learnable summary of a reward-free MDP that gives explicit policies for later rewards, with optimality under perfect training and sub-optimality proportional to training error.The summary uses two representations to approximate long-term reachability without synthesizing state trajectories.Source: §1 · p. 2

2 Problem and Notation

This section sets up reward-free Markov decision processes and asks whether a learnable object can yield near-optimal policies once a reward is specified, without planning.

2 Problem and Notation

  • The MDP is described by states, actions, transition probabilities, and a discount factor, with state-action sequences generated by the transition law and a policy.The transition probabilities may be represented as a matrix in finite spaces or as a probability measure more generally.Source: §2 · p. 2
  • The successor measure tracks expected discounted time spent in state-action sets, while a policy's Q-function sums discounted rewards along its state-action sequences.The measure formulation applies to both discrete and continuous spaces.Source: §2 · p. 2
  • The problem is to derive a near-optimal policy from a reward-free MDP representation and a later-specified reward, without planning.The setup allows discrete or continuous state spaces and bounded reward functions.Source: §2 · pp. 2–3

3 Encoding All Optimal Policies via the Forward-Backward Representation

The forward-backward representation encodes long-term state-action reachability for a family of policies. A reward is mapped into the representation space, where it specifies an optimal policy; approximation quality controls the resulting guarantee.

3 Encoding All Optimal Policies via the Forward-Backward Representation

  • A reward's Q-function is linear in the reward for a fixed policy, enabling the policy-specific occupancy representation to recover the optimal policy through a greedy choice.The argument specializes the parameter to the reward-weighted backward representation and defines the policy as greedy with respect to the resulting Q-function.Source: §3 · p. 3
  • The unsupervised criterion is a fixed point because the representations depend on policies that are themselves defined using the forward representation.The passage notes that exact solutions exist for sufficiently large dimension, while smaller dimensions provide lower-rank approximations.Source: §3 · p. 3
  • The product of the forward and backward representations approximates the long-term probability of reaching a state-action pair under the parameterized policy, without synthesizing state trajectories.The forward representation encodes the future of a state-action pair, while the backward representation encodes ways to reach a state-action pair.Source: §3 · p. 3
  • For any bounded reward, the reward-derived representation parameter selects an optimal policy, whose Q-function is given by the forward representation dotted with that parameter.The reward-derived parameter is obtained by integrating reward values against the backward representation under the measure ρ.Source: §3 · pp. 3–4
  • Exact representations exist in finite spaces when dimension exceeds the number of state-action pairs, while approximate representations have optimality gaps proportional to their occupancy-approximation error.An average error bound of ε yields a 3ε||r||∞/(1 − γ)-optimal policy under the condition stated in the passage.Source: §3 · p. 4

4 Learning and Using Forward-Backward Representations

The method learns forward-backward representations from reward-free transitions, then uses a reward representation to obtain a policy without planning. Its scope depends on exploration and representation capacity, and the section describes ways to incorporate prior reward information and adapt the method.

4 Learning and Using Forward-Backward Representations

  • Once rewards are available, the method estimates z_R by weighting B with rewards and directly uses the resulting policy.For black-box rewards, estimation requires reward observations under the state-action distribution used to train F and B.Source: §4 · pp. 4–5
  • F and B are trained off-policy without rewards so their product approximates the successor measure for each policy parameter z.The target is the successor-state density relative to the training-data distribution.Source: §4 · p. 4
  • When the reward is known algebraically, z_R can be computed from a replay buffer, and known target or forbidden state-actions can be encoded directly.A reward model trained from observations is another option.Source: §4 · p. 5
  • Prior knowledge can restrict B to relevant state features while F still uses the full state, preserving the stated guarantee for arbitrary rewards of those features.The paper gives MsPacman position as an example feature representation.Source: §4 · p. 5
  • The method assumes a good exploration strategy; the paper uses policies with random z, while also allowing an existing off-policy dataset.Source: §4 · pp. 5–6
  • A task-agnostic representation may not give the best performance on a particular reward; proposed mitigations are larger d, reward priors, or test-time Q-function fine-tuning.Source: §4 · p. 6
  • The d-dimensional reward representation necessarily loses information, and low-rank approximation favors long-range behavior while potentially blurring local reward details.The paper reports spatial blurring with small d.Source: §4 · p. 6
  • Variance may become an issue in larger environments; a proposed mitigation combines FB with related-goal selection, while expectation averaging entails handling full d × d matrices.The authors report that variance did not arise in their experiments.Source: §4 · p. 6

5 Experiments

Experiments compare FB with goal-oriented baselines, test rewards specified after training, and visualize what its learned embeddings encode. FB matches goal-oriented performance on the reported goal-reaching tasks and also succeeds on more complex reward tasks.

5.1 Environments and Experimental Setup

  • Experiments span four-room and continuous mazes, FetchReach, and pixel-based Ms. Pacman, with sparse goal rewards and success-rate evaluation outside Discrete Maze.The environments differ in state representation, dynamics, and stochasticity; training runs for 800 epochs across three random seeds.Source: §5.1 · pp. 6–7, §5.2 · p. 7

5.2 Goal-Oriented Setting: Quantitative Comparisons

  • FB’s best dimension matches the goal-oriented baseline, with no sample-efficiency loss and faster learning than DQN+HER in Ms. Pacman.Performance increases consistently with dimension; DQN alone fails to learn in Ms. Pacman.Source: Figure 1 · p. 7, Figure 2 · p. 7, §5.2 · p. 7
  • With d=25, spatial smoothing blurs exact-goal rewards in Discrete Maze, often stopping the agent next to its goal; the heatmaps contrast this with d=75.The drop occurs because reward is nonzero only at the exact goal.Source: §5.2 · p. 7, Figure 3 · p. 8

5.3 More Complex Rewards: Qualitative Results

  • FB policies successfully reach targets while avoiding forbidden positions and select the nearer of two equally rewarding targets, tasks not solvable by an already-trained goal-oriented model.The paper also reports similar results for a nearby small distracting reward and a larger distant reward.Source: §5.3 · pp. 7–8, Figure 5 · p. 8, Figure 6 · p. 8

5.4 Embedding Visualizations

  • In Continuous Maze, FB embeddings separate states across walls and group states on the same side; dents correspond to the steps needed to pass a wall.The forward embedding visualization uses z=0, corresponding to a uniform policy.Source: §5.4 · pp. 8–9, Figure 7 · p. 8

6 Related work

The related work contrasts FB with approaches that require advance task features or policy learning after a reward is known. FB instead uses successor-state occupancy representations to support policies for rewards specified later.

Successor features and successor states

  • [BBQ+18] require user-provided features and cannot represent generic or goal rewards in continuous spaces; FB uses successor-state probabilities and two representations to cover rewards beyond linear feature combinations.The backward representation B is learned alongside F and can represent rewards that depend linearly or nonlinearly on supplied features.Source: §6 · p. 9

Successor representations

  • [BTO21] introduced successor-state learning and simpler F and B forms for a fixed policy, whereas this paper addresses the every-optimal-policy setting.The authors borrow a successor-state learning algorithm from [BTO21].Source: §6 · p. 9

Goal-oriented reinforcement learning

  • Goal-oriented methods such as HER train for a reward family specified beforehand, so they cannot adapt to a posteriori tasks such as avoiding forbidden states or reaching the nearest of several goals.The passage also cites goal-dependent value-function work,, and.Source: §6 · p. 9

Other multi-task and reward-free methods

  • Options still require policy learning after a task is known, and finite-state reward-free summaries likewise require later policy learning for each reward.[JKSY20] proves that policies computed on its training set are ε-optimal for any reward.Source: §6 · p. 9

7 Conclusion

The conclusion presents FB as a learned summary of reward-free MDP dynamics that yields near-optimal policies for later-specified rewards without planning. Experiments find performance comparable to goal-oriented methods for arbitrary goals, alongside real-time adaptation to more complex rewards.

7 Conclusion

  • FB summarizes reward-free MDP dynamics and provides near-optimal policies without planning; experimentally, it performs comparably to goal-oriented methods for arbitrary goals and handles more complex rewards in real time.The learned representations encode MDP dynamics and may have broader interest.Source: §7 · p. 10

Checklist

The checklist records the authors’ reporting statements on scope, theory, reproducibility, and assets. It notes both included materials and several items marked no or not applicable.

Research scope and impact

  • The authors identify the work as theoretical research on generic reinforcement-learning methods, report discussing its limitations, and mark potential negative societal impacts as not discussed.The checklist says the potential impact is the same as reinforcement learning in general.Source: Checklist · p. 12

Theoretical results

  • The authors report stating theoretical assumptions and including complete proofs in Appendix C.Source: Checklist · p. 12

Experimental reporting

  • The authors report providing code, data, and reproduction instructions in supplemental material, with training details in Appendix D and error bars reported.They mark total compute and resource reporting as not included.Source: Checklist · p. 12

Assets

  • The checklist marks new assets as absent and existing-asset creators as cited, but asset licenses as not reported.Source: Checklist · p. 12

Human-subjects and data questions

  • Consent, personally identifiable information, and participant-risk questions are marked not applicable.The checklist also marks crowd-sourcing and human-subject compensation questions not applicable.Source: Checklist · p. 12

Outline of the Supplementary Material

The supplement documents the FB training algorithm and extends its theory to approximate solutions and rewards defined over goals or features.

Outline of the Supplementary Material

  • The supplement is organized into algorithm pseudocode, extended theoretical results and proofs, and appendices describing environments, architectures, implementation, and additional experiments.Source: Outline of the Supplementary Material · p. 13

B Extended Results: Approximate Solutions and General Goals

  • The extended results cover approximate solutions, general goals, exact representations in finite spaces, and test-time estimation of z_R from a distribution different from training.The supplement also includes proofs and additional experimental details.Source: Outline of the Supplementary Material · p. 13

A Algorithm

  • The unsupervised algorithm collects replay-buffer transitions under sampled z values, then trains F and B with a Bellman loss and an orthonormalization loss for B.It also uses target networks updated by Polyak averaging.Source: A Algorithm · p. 14

B.1 The Forward-Backward Representation With a Goal or Feature Space

  • FB representations extend to rewards that depend on known state-action features, allowing prior feature information to reduce the model and free B to support nonlinear functions of those features.The supplement also connects this extension to successor features and states a full-support condition for the reference measure.Source: §B.1 · pp. 15–16

B.2 Existence of Exact 𝐹𝐵Solutions, Influence of Dimension 𝑑, Uniqueness

The supplement establishes exact FB representations in finite settings and shows that navigation can require far fewer dimensions than representing all state-action pairs; it also treats approximate training.

B.2 Existence of Exact 𝐹𝐵Solutions, Influence of Dimension 𝑑, Uniqueness

  • Exact finite-space FB representations exist when d ≥ #S × #A, but the constructed solutions are nonunique under invertible changes of representation.The representation can be made L2(ρ)-orthonormal when B's components are linearly independent.Source: §B.2 · pp. 17–18
  • An n-dimensional toric grid supports optimal policies for every target with a representation of dimension 2n, rather than the general finite-space threshold d ≥ #S × #A.The general existence result assumes finite state and action spaces and a reference measure positive on every state-action pair.Source: §B.2 · p. 17
  • An injective B suffices to construct policies optimal for reaching every state, although the authors cannot show that their learning algorithm reaches this solution.Source: §B.2 · p. 17

B.3 Approximate Solutions Provide Approximately Optimal Policies

  • The supplement establishes that approximate F and B, and approximate reward embeddings, still yield approximately optimal policies.Source: §B.3 · p. 18

B.3.1 Influence of Approximate 𝐹and 𝐵: Optimality Gaps are Proportional to 𝑚𝜋−𝐹⊤𝐵

The optimality gap is controlled by the FB successor-state approximation error, with guarantees depending on the approximation notion and reward class. Pointwise analysis further identifies errors near optimal trajectories as particularly relevant.

B.3.1 Influence of Approximate 𝐹and 𝐵: Optimality Gaps are Proportional to 𝑚𝜋−𝐹⊤𝐵

  • Optimality gaps are directly controlled by the difference between the successor-state density and its F⊤B approximation, under several approximation notions.These include averaged density error, weak-topology measure error for Lipschitz rewards, and pointwise bounds.Source: §B.3.1 · pp. 18, 20
  • 3ε/(1 − γ) sup |R| bounds the policy's optimality gap when the successor-density approximation error is at most ε in expectation for every state-action.Source: §B.3.1 · pp. 18–19
  • In continuous spaces, weak-topology closeness between F⊤Bρ and the successor measure yields optimal policies for Lipschitz rewards.The paper notes that successor measures are usually distributions in continuous spaces, so density approximation is generally unsuitable.Source: §B.3.1 · p. 18

Then we have the componentwise inequality between state-action vectors

  • Pointwise gap bounds depend on FB errors at state-actions one step away from optimal trajectories, while the optimal trajectories themselves depend on the future reward.The authors therefore motivate minimizing a matrix norm of the error for many z values.Source: Then we have the componentwise… · p. 20

B.3.2 An approximate 𝑧𝑅yields an approximately optimal policy

The section analyzes how estimating a reward representation affects FB policy quality, including when rewards are sampled under a different state distribution. It also characterizes what the learned features capture and qualifications arising from the successor-measure definition.

B.3.2 An approximate 𝑧𝑅yields an approximately optimal policy

  • Estimating z_R from sparse empirical reward samples can overfit those samples; reducing the dimension of B regularizes both the environment model and the reward representation.With few reward samples but ample data for training F and B, the authors favor fitting a reward model over direct empirical estimation of z_R.Source: §B.3.2 · p. 20
  • At unsupervised-loss minima, F and B encode successor and predecessor features of one another, formalizing their roles as representations of future and past states.The stated result assumes an overparameterized smooth model and a data distribution with positive density everywhere.Source: §B.3.2 · pp. 21–22
  • For any bounded reward, the FB policy is optimal for its L2(ρ)-orthogonal projection onto the span of B, and is optimal for the reward itself when it lies in that span.The section also relates the size of the approximation bound to how well B captures relevant reward components and spatially smooth directions.Source: §B.3.2 · pp. 22–23
  • For rewards outside span(B), the bound favors B features that capture spatially smooth directions across policies, though the connection to FB training is not proved rigorously.The authors note that the training objective captures singular directions for the learned policies, while its relation to optimal-policy directions is unclear.Source: §B.3.2 · pp. 21, 23

B.5 Estimating 𝑧𝑅from a Different State Distribution at Test Time

  • At test time, reward observations may come from a distribution different from the training distribution used for F and B.A reward model fitted to test-time observations can be averaged over training state-actions to estimate z_R.Source: §B.5 · p. 23
  • A linear reward model over B can avoid retaining the training set at test time: precompute Cov_ρ(B) during training and combine it with test-distribution statistics.The estimate equals z_R when the test and training distributions coincide or when the reward is linear over B.Source: §B.5 · pp. 23–24

B.6 A Note on the Measure 𝑀𝜋and its Density 𝑚𝜋

  • In continuous spaces, the t=0 term gives the successor measure a Dirac component, so its density is better viewed as a distribution rather than an ordinary function.The authors report that weak approximations are covered by Theorem 8 and have not been problematic in theory or practice.Source: §B.6 · p. 24
  • Starting successor-state counts at t=1 can yield an ordinary function under nonsingularity and positive-density conditions, but works well only for rewards depending on state rather than action.Excluding t=0 also requires an immediate-reward term in the Q-function.Source: §B.6 · p. 24

C Proofs

This section establishes theoretical properties of FB-greedy policies, including finite-state construction, conditions for recovering target expectations from reward-representation estimates, and bounds relating approximation errors to policy performance. It also describes the experimental environments and network architectures.

C Proofs

  • Theorem 4’s FB-greedy policy is optimal for its specified reward, with Theorem 2 obtained as a special case.The proof identifies the policy and its Q-function as an optimal Bellman pair.Source: C Proofs · pp. 25–26
  • In finite spaces, a construction using d = |S×A| makes the FB score greedy policy optimal for every reward.The construction represents successor-measure densities and shows the resulting score equals the policy’s Q-function.Source: C Proofs · pp. 26–27
  • Propositions 17–18 bound the effects of reward changes and Q-function approximation on optimal values and greedy-policy performance.The reward-to-optimal-Q map is Lipschitz in sup norm; the greedy-policy result bounds errors using the deviation between the score and the policy’s Q-function.Source: C Proofs · pp. 27–29
  • Theorem 8 transfers successor-model approximation error into an optimality bound, while Proposition 19 gives a pointwise bound for comparing policies optimal under different rewards.The theorem proof reduces its claims to a general norm bound; the proposition bounds one policy’s gap on the other policy’s reward.Source: C Proofs · pp. 29–30
  • Theorem 9 shows that any score function f is an optimal Q-function for a reward constructed from f and its greedy policy, and the proof also applies to continuous state spaces.The proof uses a policy-comparison bound to relate the constructed reward to the target reward.Source: C Proofs · pp. 30–31
  • Theorem 12 constructs a reward corresponding to the estimated reward representation and uses a high-probability error bound to obtain a policy guarantee.The proof applies Proposition 10 after bounding the reward-estimation error.Source: C Proofs · pp. 31–32
  • Under the overparameterization assumption, Theorem 13’s stationarity condition forces the relevant residual to vanish ρ-almost everywhere.The proof uses the ability to realize any L2 function as a parameter derivative and establishes the corresponding result for F and B.Source: C Proofs · p. 33
  • Theorem 14 proves that the FB policy is optimal for bounded rewards in the span of B and considers arbitrary rewards through their L2(ρ) projection onto that span.The proof identifies the reward representation for both the in-span reward and its projection.Source: C Proofs · pp. 33–34
  • When the test distribution matches the target distribution, covariance terms cancel; for rewards linear in B, both estimates agree across test distributions if Cov_test B is invertible.In the linear case, both estimates equal (Cov_ρ B)w₀.Source: C Proofs · p. 34

D.1 Environments

  • The discrete maze is an 11 × 11 deterministic four-room gridworld, while the continuous maze uses noisy 0.1-unit moves, impassable walls, and 441 radial-basis features.Continuous-maze moves add zero-mean Gaussian noise with standard deviation 0.01; positions are clipped to [0, 1]².Source: §D.1 · p. 35
  • FetchReach uses 10-dimensional joint-position and velocity states with six discretized actions; Ms. Pacman uses four stacked 84 × 84 grayscale frames and 148 reachable screen-coordinate goals.FetchReach adapts the simulated robotic-arm environment [PAR+18]; Ms. Pacman uses an annotated RAM interface [ARO+19] rather than the position heuristic in.Source: §D.1 · p. 35

D.2 Architectures

  • Maze and FetchReach forward/backward networks use three 256-unit ReLU layers; Ms. Pacman uses a convolutional forward network and a three-layer position-input backward network.The Ms. Pacman forward network has three convolutional layers followed by three 256-unit fully connected layers; goal-oriented DQN uses corresponding maze and Ms. Pacman architectures.Source: §D.2 · pp. 35–36

D.3 Implementation Details

The appendix specifies FB training and evaluation procedures, then presents additional performance comparisons and visualizations of policies and embeddings.

D.3 Implementation Details

  • FB training runs for 800 epochs with replay-buffer cycles; mazes use uniform behavior, while FetchReach and Ms. Pacman use ε-greedy behavior, and evaluation uses smaller exploration except in the discrete maze.The backward representation acts on state in mazes, gripper position in FetchReach, and agent screen position in Ms. Pacman.Source: §D.3 · p. 36
  • The training procedure uses Polyak-averaged target networks, an orthonormality regularizer for B, and Adam updates with regularization coefficient 1.The empirical loss is estimated from sampled transition and target state-action minibatches; the specified batch size is 128.Source: §D.3 · pp. 36–37
  • Training samples z from a rescaled Gaussian construction using a Cauchy scalar so its norm spans non-negative values with a heavy tail.The vector is scaled to keep each component on the order of 1; positive rescaling preserves the optimal policy.Source: §D.3 · p. 36

D.4 Experimental results

  • Figures 8–11 compare FB at different dimensions with DQN: policy quality in discrete maze and success rates in continuous maze, FetchReach, and Ms. Pacman, across epochs and post-training evaluations.For Ms. Pacman, the post-training success rate is averaged over 184 handcrafted goals; FB-50 and F-100 were trained for only 200 epochs.Source: Figure 8 · p. 38, Figure 9 · p. 38, Figure 10 · p. 38

D.4.3 Embedding Visualization

  • Policy visualizations show score maps and trajectories for target avoidance, selecting the closest of two equally rewarding goals, and trading off a small nearby reward against a larger distant one.The appendix also shows a full Ms. Pacman trajectory for reaching a target while avoiding forbidden positions.Source: Figure 13 · p. 40, Figure 14 · p. 41, Figure 15 · p. 42
  • Two-dimensional t-SNE visualizations show FB embeddings recovering maze room-and-door structure and Ms. Pacman cycle structure; the Ms. Pacman F embedding uses visual inputs while B uses agent position.Additional figures show continuous-maze states, F and B embeddings, and F embeddings for different z vectors.Source: Figure 17 · p. 44, Figure 19 · p. 44, Figure 18 · p. 44

Related papers

Loading 2103.07945v3…