Learning One Representation to Optimize All Rewards
Year: 2021

Abstract:
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 rewardfree 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 reward 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. 2

heading: 1 Introduction
text: We consider one kind of unsupervised reinforcement learning problem: Given a Markov decision process (MDP) but no reward information, is it possible to learn and store a compact object that, for any reward function specified later, provides the optimal policy for that reward, with a minimal amount of additional computation? In a sense, such an object would encode in a compact form the solutions of all possible planning problems in the environment. This is a step towards building agents that are fully controllable after first exploring their environment in an unsupervised way.
Goal-oriented RL methods [ACR+17, PAR+18] compute policies for a series of rewards specified in advance (such as reaching a set of target states), but cannot adapt in real time to new rewards, such as weighted combinations of target states or dense rewards.
Learning a model of the world is another possibility, but it still requires explicit planning for each new reward; moreover, synthesizing accurate trajectories of states over long time ranges has proven difficult [Tal17, KST+18].
⇤Work done during an internship at Facebook Artificial Intelligence Research Paris. 2Code: https://github.com/ahmed-touati/controllable_agent
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Instead, we exhibit an object that is both simpler to learn than a model of the world, and contains the information to recover near-optimal policies for any reward provided a posteriori, without a planning phase.
[BBQ+18] learn optimal policies for all rewards that are linear combinations of a finite number of feature functions provided in advance by the user. This limits applications: e.g., goal-oriented tasks would require one feature per goal state, thus using infinitely many features in continuous spaces. We reuse a policy parameterization from [BBQ+18], but introduce a novel representation with better properties, based on state occupancy prediction instead of expected featurizations. We use theoretical advances on successor state learning from [BTO21]. We obtain the following.
• We prove the existence of a learnable “summary” of a reward-free discrete or continuous MDP, that provides an explicit formula for optimal policies for any reward specified later. This takes the form of a pair of representations F : S ⇥A⇥ Z ! Z and B : S ⇥A! Z from state-actions into a representation space Z ' Rd, with policies ⇡z(s) := argmaxa F (s, a, z)>z. Once a reward is specified, a value of z is computed from reward values and B; then ⇡z is used. Rewards may be specified either explicitly as a function, or as target states, or by samples as in usual RL setups. • We provide a well-principled unsupervised loss for F and B. If FB training is perfect, then the policies are provably optimal for all rewards (Theorem 2). With imperfect training, sub-optimality is proportional to the FB training error (Theorems 8–9). In finite spaces, perfect training is possible with large enough dimension d (Proposition 6). Explicitly, F and B are trained so that F (s, a, z)>B(s0, a0) approximates the long-term probability to reach s0 from s if following ⇡z . This is akin to a model of the environment, without synthesizing state trajectories. • We provide a TD-like algorithm to train F and B for this unsupervised loss, with function approximation, adapted from recent methods for successor states [BTO21]. No sparse rewards are used: every transition reaches some state s0, so every step is exploited. As usual with TD, learning seeks a fixed point but the loss itself is not observable. • We prove viability of the method on several environments from mazes to pixel-based MsPacman and a virtual robotic arm. For single-state rewards (learning to reach arbitrary states), we provide quantitative comparisons with goal-oriented methods such as HER. (Our method is not a substitute for HER: in principle they could be combined, with HER improving replay buffer management for our method.) For more general rewards, which cannot be tackled a posteriori by trained goal-oriented models, we provide qualitative examples. • We also illustrate qualitatively the sub-optimalities (long-range behavior is preserved but local blurring of rewards occurs) and the representations learned.

heading: 2 Problem and Notation
text: Let M = (S,A, P, ) be a reward-free Markov decision process with state space S (discrete or continuous), action space A (discrete for simplicity, but this is not essential), transition probabilities P (s0|s, a) from state s to s0 given action a, and discount factor 0 < < 1 [SB18]. If S is finite, P (s0|s, a) can be viewed as a matrix; in general, for each (s, a) 2 S⇥A, P (ds0|s, a) is a probability measure on s0 2 S. The notation P (ds0|s, a) covers all cases. Given (s0, a0) 2 S ⇥A and a policy ⇡ : S ! Prob(A), we denote Pr(·|s0, a0,⇡) and E[·|s0, a0,⇡] the probabilities and expectations under state-action sequences (st, at)t 0 starting with (s0, a0) and following policy ⇡ in the environment, defined by sampling st ⇠ P (dst|st 1, at 1) and at ⇠ ⇡(st). For any policy ⇡ and state-action (s0, a0), define the successor measure M⇡(s0, a0, ·) as the measure over S ⇥A representing the expected discounted time spent in each set X ⇢ S ⇥A:
M⇡(s0, a0, X) := X
t 0 t Pr ((st, at) 2 X | s0, a0, ⇡) (1)
for each X ⇢ S ⇥A. Viewing M as a measure deals with both discrete and continuous spaces. Given a reward function r : S ⇥ A ! R, the Q-function of ⇡ for r is Q⇡r (s0, a0) :=P
t 0 t E[r(st, at)|s0, a0,⇡]. We assume that rewards are bounded, so that all Q-functions are well-defined. We state the results for deterministic reward functions, but this is not essential. We abuse notation and write greedy policies as ⇡(s) = argmaxa Q(s, a) instead of ⇡(s) 2 argmaxa Q(s, a). Ties may be broken any way.
We consider the following informal problem: Given a reward-free MDP (S,A, P, ), can we compute a convenient learnable object E such that, once a reward function r : S ⇥A! R is specified, we can easily (with no planning) compute, from E and r, a policy ⇡ whose performance is close to maximal?

heading: 3 Encoding All Optimal Policies via the Forward-Backward Representation
text: We first present forward-backward (FB) representations of a reward-free MDP as a way to summarize all optimal policies via explicit formulas. The resulting learning procedure is described in Section 4.
Core idea. The main algebraic idea is as follows. Assume, at first, that S is finite. For a fixed policy, the Q-function depends lineary on the reward: namely, Q⇡r (s, a) = P s0,a0 M ⇡(s, a, s0, a0)r(s0, a0)
where M⇡(s, a, s0, a0) = P
t 0 t Pr ((st, at) = (s0, a0)|s, a,⇡). This rewrites as Q⇡r = M⇡r
viewing everything as vectors and matrices indexed by state-actions.
Now let (⇡z)z2Rd be any family of policies parameterized by z. Assume that for each z, we can find d⇥ (S ⇥ A)-matrices Fz and B such that M⇡z = F>zB. Then Q⇡zr = F>zBr. Specializing to zR := Br, the Q-function of policy ⇡zR on reward r is Q ⇡zR r = F>zRzR. So far ⇡z was unspecified; but if we define ⇡z(s) := argmaxa(F>z z)sa at each state s, then by definition, ⇡zR is the greedy policy with respect to F>zRzR. At the same time, F > zRzR is the Q-function of ⇡zR for reward r: thus, ⇡zR is the greedy policy of its own Q-function, and is therefore optimal for reward r.
Thus, if we manage to find F , B, and ⇡z such that ⇡z = argmaxF>z z and F>zB = M⇡z for all z 2 Rd, then we obtain the optimal policy for any reward r, just by computing Br and applying policy ⇡Br.
This criterion on (F,B,⇡z) is entirely unsupervised. Since F and B depend on ⇡z but ⇡z is defined via F , this is a fixed point equation. An exact solution exists for d large enough (Appendix, Prop. 6), while a smaller d provides lower-rank approximations M⇡z ⇡ F>zB. In Section 4 we present a well-grounded algorithm to learn such F , B, and ⇡z .
In short, we learn two representations F and B such that F (s0, a0, z)>B(s0, a0) is approximately the long-term probability M⇡z (s0, a0, s0, a0) to reach (s0, a0) if starting at (s0, a0) and following policy ⇡z . Then all optimal policies can be computed from F and B. We think of F as a representation of the future of a state, and B as the ways to reach a state (Appendix B.4): if F>B is large, then the second state is reachable from the first. This is akin to a model of the environment, without synthesizing state trajectories.
General statement. In continuous spaces with function approximation, Fz and B become functions S ⇥A! Rd instead of matrices; since Fz depends on z, F itself is a function S ⇥A⇥ Rd ! Rd. The sums over states will be replaced with expectations under the data distribution ⇢. Definition 1 (Forward-backward representation). Let Z = Rd be a representation space, and let ⇢ be a measure on S ⇥A. A pair of functions F : S ⇥A⇥Z ! Z and B : S ⇥A! Z, together with a parametric family of policies (⇡z)z2Z , is called a forward-backward representation of the MDP with respect to ⇢, if the following conditions hold for any z 2 Z and (s, a), (s0, a0) 2 S ⇥A: ⇡z(s) = argmax
a F (s, a, z)>z, M⇡z (s0, a0, ds, da) = F (s0, a0, z) >B(s, a)⇢(ds, da) (2)
where M⇡ is the successor measure defined in (1), and the last equality is between measures. Theorem 2 (FB representations encode all optimal policies). Let (F,B, (⇡z)) be a forward-backward representation of a reward-free MDP with respect to some measure ⇢.
Then, for any bounded reward function r : S ⇥A! R, the following holds. Set
zR :=
Z
s,a r(s, a)B(s, a) ⇢(ds, da). (3)
assuming the integral exists. Then ⇡zR is an optimal policy for reward r in the MDP. Moreover, the optimal Q-function Q? for reward r is Q?(s, a) = F (s, a, zR)>zR.
For instance, for a single reward located at state-action (s, a), the optimal policy is ⇡zR with zR = B(s, a). (In that case the factor ⇢(ds, da) does not matter because scaling the reward does not change the optimal policy.)
We present in Section 4 an algorithm to learn FB representations. The measure ⇢ will be the distribution of state-actions visited in a training set or under an exploration policy: then zR = E(s,a)⇠⇢[r(s, a)B(s, a)] can be obtained by sampling from visited states. In finite spaces, exact FB representations exist, provided the dimension d is larger than #S ⇥#A (Appendix, Prop. 6). In infinite spaces, arbitrarily good approximations can be obtained by increasing d, corresponding to a rank-d approximation of the cumulated transition probabilities M⇡ . Importantly, the optimality guarantee extends to approximate F and B, with optimality gap proportional to F>B M⇡z/⇢ (Appendix, Theorems 8–9 with various norms on F>B M⇡/⇢). For instance, if, for some reward r, the error F (s0, a0, zR)>B(s, a) M⇡zR (s0, a0, ds, da)/⇢(ds, da) is at most " on average over (s, a) ⇠ ⇢ for every (s0, a0), then ⇡zR is 3" krk1 /(1 )-optimal for r.
These results justify using some norm over F>B M⇡z/⇢ , averaged over z 2 Rd, as a training loss for unsupervised reinforcement learning. (Below, we average over z 2 Rd from a fixed rescaled Gaussian. If prior information is available on the rewards r, the corresponding distribution of zR may be used instead.)
If B is fixed in advance and only F is learned, the method has similar properties to successor features based on B (Appendix B.4). But one may set a large d and let B be learned: arguably, by Theorem 2, the resulting features “linearize” optimal policies as much as possible. The features learned in F and B may have broader interest.

heading: 4 Learning and Using Forward-Backward Representations
text: Our algorithm starts with an unsupervised learning phase, where we learn the representations F and B in a reward-free way, by observing state transitions in the environment, generated from any exploration scheme. Then, in a reward estimation phase, we estimate a policy parameter zR = E[r(s, a)B(s, a)] from some reward observations, or directly set zR if the reward is known (e.g., set zR = B(s, a) to reach a known target (s, a)). In the exploitation phase, we directly use the policy ⇡zR(s) = argmaxa F (s, a, zR)>zR.
The unsupervised learning phase. No rewards are used in this phase, and no family of tasks has to be specified manually. F and B are trained off-policy from observed transitions in the environment. The first condition of FB representations, ⇡z(s) = argmaxa F (s, a, z)>z, is just taken as the definition of ⇡z given F . In turn, F and B are trained so that the second condition (2), F (·, z)>B = M⇡z/⇢, holds for every z. Here ⇢ is the (unknown) distribution of state-actions in the training data. Training is based on the Bellman equation for the successor measure M⇡ ,
M⇡(s0, a0, {(s0, a0)}) = s0=s0, a0=a0 + Es1⇠P (ds1|s0,a0) M⇡(s1,⇡(s1), {(s0, a0)}). (4)
We leverage a well-principled algorithm from [BTO21] in the single-policy setting: it learns the successor measure of a policy ⇡ without using the sparse reward s0=s0, a0=a0 (which would vanish in continuous spaces). Other successor measure algorithms could be used, such as C-learning [ESL21].
The algorithm from [BTO21] uses a parametric model m⇡✓ (s0, a0, s 0, a0) to represent M⇡(s0, a0, ds0, da0) ⇡ m⇡✓ (s0, a0, s0, a0)⇢(ds0, da0). It is not necessary to know ⇢, only to sample states from it. Given an observed transition (s0, a0, s1) from the training set, generate an action a1 ⇠ ⇡(a1|s1), and sample another state-action (s0, a0) from the training set, independently from (s0, a0, s1). Then update the parameter ✓ by ✓ ✓ + ⌘ ✓ with learning rate ⌘ and ✓ := @✓m ⇡ ✓ (s0, a0, s0, a0)+@✓m ⇡ ✓ (s0, a0, s
0, a0) ⇥( m⇡✓ (s1, a1, s0, a0) m⇡✓ (s0, a0, s0, a0)) (5) This computes the density m⇡ of M⇡ with respect to the distribution ⇢ of state-actions in the training set. Namely, the true successor state density m⇡ = M⇡/⇢ is a fixed point of (5) in expectation [BTO21, Theorem 6] (and is the only fixed point in the tabular or overparameterized case). Variants exist, such as using a target network for m⇡✓ (s1, a1, s 0, a0) on the right-hand side, as in DQN.
Thus, we first choose a parametric model F✓, B✓ for the representations F and B, and set m⇡z✓ (s0, a0, s
0, a0) := F✓(s0, a0, z)>B✓(s0, a0). Then we iterate the update (5) over many stateactions and values of z. This results in Algorithm 1. At each step, a value of z is picked at random, together with a batch of transitions (s0, a0, s1) and a batch of state-actions (s0, a0) from the training set, with (s0, a0) independent from z and (s0, a0, s1).
For sampling z, we use a fixed distribution (rescaled Gaussians, see Appendix D). Any number of values of z may be sampled: this does not use up training samples. We use a target network with soft updates (Polyak averaging) as in DDPG. For training we also replace the greedy policy ⇡z = argmaxa F (s, a, z)
>z with a regularized version ⇡z = softmax(F (s, a, z)>z/⌧) with fixed temperature ⌧ (Appendix D). Since there is unidentifiability between F and B (Appendix, Remark 7), we normalize B via an auxiliary loss in Algorithm 1.
For exploration in this phase, we use the policies being learned: the exploration policy chooses a random value of z from some distribution (e.g., Gaussian), and follows ⇡z for some time (Appendix, Algorithm 1). However, the algorithm can also work from an existing dataset of off-policy transitions.
The reward estimation phase. Once rewards are available, we estimate a reward representation (policy parameter) zR by weighing the representation B by the reward:
zR := E[r(s, a)B(s, a)] (6)
where the expectation must be computed over the same distribution ⇢ of state-actions (s, a) used to learn F and B (see Appendix B.5 for using a different distribution). Thus, if the reward is black-box as in standard RL algorithms, then the exploration policy has to be run again for some time, and zR is obtained by averaging r(s, a)B(s, a) over the states visited. An approximate value for zR still provides an approximately optimal policy (Appendix, Prop. 10 and Thm. 12).
If the reward is known explicitly, this phase is unnecessary. For instance, if the reward is to reach a target state-action (s0, a0) while avoiding some forbidden state-actions (s1, a1), ..., (sk, ak), one may directly set
zR = B(s0, a0) X B(si, ai) (7)
where the constant sets the negative reward for forbidden states and adjusts for the unknown ⇢(dsi, dai) factors in (3). This can be used for goal-oriented RL.
If the reward is known algebraically as a function r(s, a), then zR may be computed by averaging the function r(s, a)B(s, a) over a replay buffer from the unsupervised training phase. We may also use a reward model r̂(s, a) of r(s, a) trained on some reward observations from any source.
The exploitation phase. Once the reward representation zR has been estimated, the Q-function is estimated as
Q(s, a) = F (s, a, zR) >zR. (8)
The corresponding policy ⇡zR(s) = argmaxa Q(s, a) is used for exploitation.
Fine-tuning was not needed in our experiments, but it is possible to fine-tune the Q-function using actual rewards, by setting Q(s, a) = F (s, a, zR)>zR + q✓(s, a) where the fine-tuning model q✓ is initialized to 0 and learned via any standard Q-learning method.
Incorporating prior information on rewards in B. Trying to plan in advance for all possible rewards in an arbitrary environment may be too generic and problem-agnostic, and become difficult in large environments, requiring long exploration and a large d to accommodate all rewards. In practice, we are often interested in rewards depending, not on the full state, but only on a part or some features of the state (e.g., a few components of the state, such as the position of an agent, or its neighbordhood, rather than the full environment).
If this is known in advance, the representation B can be trained on that part of the state only, with the same theoretical guarantees (Appendix, Theorem 4). F still needs to use the full state as input. This way, the FB model of the transition probabilities (1) only has to learn the future probabilities of the part of interest in the future states, based on the full initial state (s0, a0). Explicitly, if ' : S⇥A! G is a feature map to some features g = '(s, a), and if we know that the reward will be a function R(g), then Theorem 2 still holds with B(g) everywhere instead of B(s, a), and with the successor measure M⇡(s0, a0, dg) instead of M⇡(s0, a0, ds0, da0) (Appendix, Theorem 4). Learning is done by replacing @✓m⇡✓ (s0, a0, s0, a0) with @✓m ⇡ ✓ (s0, a0,'(s0, a0)) in the first term in (5) [BTO21]. Rewards can be arbitrary functions of g, so this is more general than [BBQ+18] which only considers rewards linear in g. For instance, in MsPacman below, we let g be the 2D position (x, y) of the agent, so we can optimize any reward function that depends on this position.
Limitations. First, this method does not solve exploration: it assumes access to a good exploration strategy. (Here we used the policies ⇡z with random values of z, corresponding to random rewards.)
Next, this task-agnostic approach is relevant if the reward is not known in advance, but may not bring the best performance on a particular reward. Mitigation strategies include: increasing d; using prior information on rewards by including relevant variables into B, as discussed above; and fine-tuning the Q-function at test time based on the initial F>B estimate.
As reward functions are represented by a d-dimensional vector zR = E[r.B], some information about the reward is necessarily lost. Any reward uncorrelated to B is treated as 0. The dimension d controls how many types of rewards can be optimized well. A priori, a large d may be required. Still, in the experiments, d ⇡ 100 manages navigation in a pixel-based environment with a huge state space. Appendix B.2 argues theoretically that d = 2n is enough for navigation on an n-dimensional grid. The algorithm is linear in d, so d can be taken as large as the neural network models can handle.
We expect this method to have an implicit bias for long-range behavior (spatially smooth rewards), while local details of the reward function may be blurred. Indeed, F>B is optimized to approximate the successor measure M⇡ = P t
tP t⇡ with P t⇡ the t-step transition kernel for each policy ⇡. The rank-d approximation will favor large eigenvectors of P⇡ , i.e., small eigenvectors of the Markov chain Laplacian Id P⇡. These loosely correspond to long-range (low-frequency) behavior [MM07]: presumably, F and B will learn spatially smooth rewards first. Indeed, experimentally, a small d leads to spatial blurring of rewards and Q-functions (Fig. 3). Arguably, without any prior information this is a reasonable prior. [SBG17] have argued for the cognitive relevance of low-dimensional approximations of successor representations.
Variance is a potential issue in larger environments, although this did not arise in our experiments. Learning M⇡ requires sampling a state-action (s0, a0) and an independent state-action (s0, a0). In large spaces, most state-action pairs will be unrelated. A possible mitigation is to combine FB with strategies such as Hindsight Experience Replay [ACR+17] to select goals related to the current state-action. The following may help a lot: the update of F and B decouples as an expectation over (s0, a0), times an expectation over (s0, a0). Thus, by estimating these expectations by a moving average over a dataset, it is easy to have many pairs (s0, a0) interact with many (s0, a0). The cost is handling full d⇥ d matrices. This will be explored in future work.

heading: 5 Experiments
text: We first consider the task of reaching arbitrary goal states. For this, we can make quantitative comparisons to existing goal-oriented baselines. Next, we illustrate qualitatively some tasks that cannot be tackled a posteriori by goal-oriented methods, such as introducing forbidden states. Finally, we illustrate some of the representations learned.

heading: 5.1 Environments and Experimental Setup
text: We run our experiments on a selection of environments that are diverse in term of state space dimensionality, stochasticity and dynamics.
• Discrete Maze is the classical gridworld with four rooms. States are represented by one-hot unit vectors. • Continuous Maze is a two dimensional environment with impassable walls. States are represented by their Cartesian coordinates (x, y) 2 [0, 1]2. The execution of one of the actions moves the agent in the desired direction, but with normal random noise added to the position of the agent. • FetchReach is a variant of the simulated robotic arm environment from [PAR+18] using discrete actions instead of continuous actions. States are 10-dimensional vectors consisting of positions and velocities of robot joints. • Ms. Pacman is a variant of the Atari 2600 game Ms. Pacman, where an episode ends when the agent is captured by a monster [RUMS18]. States are obtained by processing the raw visual input directly from the screen. Frames are preprocessed by cropping, conversion to grayscale and downsampling to 84⇥ 84 pixels. A state st is the concatenation of (xt 12, xt 8, xt 4, xt) frames, i.e. an 84 ⇥ 84 ⇥ 4 tensor. An action repeat of 12 is used. As Ms. Pacman is not originally a multi-goal domain, we define the goals as the 148 reachable coordinates (x, y) on the screen; these can be reached only by learning to avoid monsters.
For all environments, we run algorithms for 800 epochs, with three different random seeds. Each epoch consists of 25 cycles where we interleave between gathering some amount of transitions, to add to the replay buffer, and performing 40 steps of stochastic gradient descent on the model parameters. To collect transitions, we generate episodes using some behavior policy. For both mazes, we use a uniform policy while for FetchReach and Ms. Pacman, we use an "-greedy policy with respect to the current approximation F (s, a, z)>z for a sampled z. At evaluation time, "-greedy policies are also used, with a smaller ". More details are given in Appendix D.

heading: 5.2 Goal-Oriented Setting: Quantitative Comparisons
text: We investigate the FB representation over goal-reaching tasks and compare it to goal-oriented baselines: DQN3, and DQN with HER when needed. We define sparse reward functions. For Discrete Maze, the reward function is equal to one when the agent’s state is equal exactly to the goal state. For Discrete Maze, we measured the quality of the obtained policy to be the ratio between the true expected discounted reward of the policy for its goal and the true optimal value function, on average over all states. For the other environments, the reward function is equal to one when the distance of the agent’s position and the goal position is below some threshold, and zero otherwise. We assess policies by computing the average success rate, i.e the average number of times the agent successfully reaches its goal.
Figs. 1 and 2 show the comparative performance of FB for different dimensions d, and DQN respectively in FetchReach and Ms. Pacman (similar results in Discrete and Continuous Mazes are provided in Appendix D). In Ms. Pacman, DQN totally fails to learn and we had to add HER to make it work. The performance of FB consistently increases with the dimension d and the best dimension matches the performance of the goal-oriented baseline.
In Discrete Maze, we observe a drop of performance for d = 25 (Appendix D, Fig. 8): this is due to the spatial smoothing induced by the small rank approximation and the reward being nonzero only if the agent is exactly at the goal. This spatial blurring is clear on heatmaps for d = 25 vs d = 75 (Fig. 3). With d = 25 the agent often stops right next to its goal.
To evaluate the sample efficiency of FB, after each epoch, we evaluate the agent on 20 randomly selected goals. Learning curves are reported in Figs. 1 and 2 (left). In all environments, we observe no loss in sample efficiency compared to the goal-oriented baseline. In Ms. Pacman, FB even learns faster than DQN+HER.

heading: 5.3 More Complex Rewards: Qualitative Results
text: We now investigate FB’s ability to generalize to new tasks that cannot be solved by an already trained goal-oriented model: reaching a goal with forbidden states imposed a posteriori, reaching the nearest of two goals, and choosing between a small, close reward and a large, distant one.
First, for the task of reaching a target position g0 while avoiding some forbidden positions g1, . . . gk , we set zR = B(g1) Pk i=1 B(gi) and run the corresponding "-greedy policy defined
3Here DQN is short for goal-oriented DQN, Q(s, a, g).
by F (s, a, zR)>zR. Fig. 5 shows the resulting trajectories, which succeed at solving the task for the different domains. In Ms. Pacman, the path is suboptimal (though successful) due to the sudden appearance of a monster along the optimal path. (We only plot the initial frame; see the full series of frames along the trajectory in Appendix D, Fig. 16.) Fig. 4 (left) provides a contour plot of maxa2A F (s, a, zR)>zR for the continuous maze and shows the landscape shape around the forbidden regions.
Next, we consider the task of reaching the closest target among two equally rewarding positions g0 and g1, by setting zR = B(g0) +B(g1). The optimal Q-function is not a linear combination of the Q-functions for g0 and g1. Fig. 6 shows successful trajectories generated by the policy ⇡zR . On the contour plot of maxa2A F (s, a, zR)>zR in Fig. 4 (right), the two rewarding positions appear as basins of attraction. Similar results for a third task are shown in Appendix D: introducing a “distracting” small reward next to the initial position of the agent, with a larger reward further away.
(star shape while avoiding forbidden positions (red shape )

heading: 5.4 Embedding Visualizations
text: We visualize the learned FB state embeddings for Continuous Maze by projecting them into 2- dimensional space using t-SNE [VdMH08] in Fig. 7. For the forward embeddings, we set z = 0 corresponding to the uniform policy. We can see that FB partitions states according to the topology induced by the dynamics: states on opposite sides of walls are separated in the representation space
and states on the same side lie together. Appendix D includes embedding visualizations for different z and for Discrete Maze and Ms. Pacman.

heading: 6 Related work
text: [BBQ+18] learn optimal policies for rewards that are linear combinations of a finite number of feature functions provided in advance by the user. This approach cannot tackle generic rewards or goal-oriented RL: this would require introducing one feature per possible goal state, requiring infinitely many features in continuous spaces.
Our approach does not require user-provided features describing the future tasks, thanks to using successor states [BTO21] where [BBQ+18] use successor features. Schematically, and omitting actions, successor features start with user-provided features ', then learn such that (s0) =P
t 0 t E['(st) | s0]. This limits applicability to rewards that are linear combinations of '. Here we use successor state probabilities, namely, we learn two representations F and B such that F (s0)>B(s0) = P t 0 t Pr(st = s0 | s0). This does not require any user-provided input.
Thus we learn two representations instead of one. The learned backward representation B is absent from [BBQ+18]. B plays a different role than the user-provided features ' of [BBQ+18]: if the reward is known a priori to depend only on some features ', we learn B on top of ', which represents all rewards that depend linearly or nonlineary on '. Up to a change of variables, [BBQ+18] is recovered by setting B = Id on top of ', or B = ' and ' = Id, and then only training F .
We use a similar parameterization of policies by F (s, a, z)>z as in [BBQ+18], for similar reasons, although z encodes a different object.
Successor representations where first defined in [Day93] for finite spaces, corresponding to an older object from Markov chains, the fundamental matrix [KS60, Bré99, GS97]. [SBG17] argue for their relevance for cognitive science. For successor representations in continuous spaces, a finite number of features ' are specified first; this can be used for generalization within a family of tasks, e.g., [BDM+17, ZSBB17, GHB+19, HDB+19]. [BTO21] moves from successor features to successor states by providing pointwise occupancy map estimates even in continuous spaces, without using the sparse reward st=s0 . We borrow a successor state learning algorithm from [BTO21]. [BTO21] also introduced simpler versions of F and B for a single, fixed policy; [BTO21] does not consider the every-optimal-policy setting.
There is a long literature on goal-oriented RL. For instance, [SHGS15] learn goal-dependent value functions, regularized via an explicit matrix factorization. Goal-dependent value functions have been investigated in earlier works such as [FD02] and [SMD+11]. Hindsight experience replay (HER) [ACR+17] improves the sample efficiency of multiple goal learning with sparse rewards. A family of rewards has to be specified beforehand, such as reaching arbitrary target states. Specifying rewards a posteriori is not possible: for instance, learning to reach target states does not extend to reaching the nearest among several goals, reaching a goal while avoiding forbidden states, or maximizing any dense reward.
Hierarchical methods such as options [SPS99] can be used for multi-task RL problems. However, policy learning on top of the options is still needed after the task is known.
For finite state spaces, [JKSY20] use reward-free interactions to build a training set that summarizes a finite environment, in the sense that any optimal policies later computed on this training set instead of the true environment are provably "-optimal, for any reward. They prove tight bounds on the necessary set size. Policy learning still has to be done afterwards for each reward.
Acknowledgments. The authors would like to thank Léonard Blier, Diana Borsa, Alessandro Lazaric, Rémi Munos, Tom Schaul, Corentin Tallec, Nicolas Usunier, and the anonymous reviewers for numerous comments, technical questions, references, and invaluable suggestions for presentation that led to an improved text.

heading: 7 Conclusion
text: The forward-backward representation is a learnable mathematical object that “summarizes” a rewardfree MDP. It provides near-optimal policies for any reward specified a posteriori, without planning. It is learned from black-box reward-free interactions with the environment. In practice, this unsupervised method performs comparably to goal-oriented methods for reaching arbitrary goals, but is also able to tackle more complex rewards in real time. The representations learned encode the MDP dynamics and may have broader interest.

Review summaries:
This paper proposes an algorithm learning a kind of universal value function (UVFA), so that given a buffer with a good coverage of transitions it returns the optimal behavior for any reward function or any goal state by computing a simple average.  The idea comes from the fact that the Q-function for a given policy $\pi$ -- that induces a reaching probability from any ($s,a$) $M^\pi(s, a, s', a')$ -- can be expressed as the scalar produce $Q^\pi(s,a) = \sum_{s',a'} M^\pi(s,a, s', a') r(s',a') = \langle M(s,a); r\rangle$. Then, if we learn $F(s,a,z)$ and $B(z)$ such that $M^{\pi_z}(s,a) = F(s,a, z)^T B(z)$ we obtain $Q^\pi_z = \langle F(s,a, z)^T B(z); r\rangle = F(s,a, z_r)^T z_r$ where $z_r = \langle B(z); r\rangle$.  The greedy policy given such a Q-function should maximize (over action) $F(s,a, z_r)^T z_r$. Now, there is a fix point {$F^*, B^*, \pi^*$} such that for all $z$, $Q^{\pi^*}(s, a) = F^*(s,a, z)^T z$, which coincides with the optimal Q-function for the given $z$.   The paper provides an algorithm that learns these optimal $F^*$ and $B^*$. Then, given any reward function, they compute $z_r = \langle B^*(z); r\rangle$ and return the optimal Q-function $Q^*(s,a) = F^*(s,a, z_r)^T z_r$. 

This paper considered the representation learning in reward-free MDPs to encode the optimal policy given any reward function without planning. The proposed unsupervised learning algorithm FB offered two features $F$ and $B$ that encoded the information of the transition to extract the optimal policy. They showed that as long as $F^\top B$ approximates the probability of visiting one state-action to another, then we can extract the optimal policy given any reward function directly. Motived by the nice theoretical properties of $F$ and $B$, they showed that the unsupervised learning algorithm can achieve the same level of learning compared with DQN, DQN with HER, etc. on several benchmark environments. They also incorporated some engineering tricks to make the theory-boosted algorithm more effective.