Learning One Representation to Optimize All Rewards
Authors: Ahmed Touati
Yann Ollivier
Venue: NeurIPS 2021 Poster (2021)
Keywords: reward-free MDP
successor states
self-supervision
provable representation learning

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 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 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.

Learning One Representation to Optimize All Rewards
Ahmed Touati" Mila, University of Montreal ahmed. Poouati@umontres ca
Yann Ollivier Facebook Artificial Intelligence Research Paris yol@fb.com
Abstract
We introduce the reward-free Markov decision process. It provides explicit near-optimal policies for any reward specified 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 reward observations or an explicit reward description (e.g.. target state). The optimal policy for that reward is directly obtained from these representations, with no planning. We assume access to exploration scheme replay buffer for the first phase.
(FB) representation of the dynamics of a
The corresponding unsupervised loss is well-prineipled: if training is perfect, the policies obtained are provably optimal for any reward function. With imperfect training. the sub-optimality is proportional error. The FB representation learns long-range relationships between states and actions, via predictive occupancy map, without having synthesize states as model-based approaches.
the unsupervised approximation
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
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 compact object that, for any reward function specified later, provides the optimal policy for that reward, with a minimal amount of additional computation? In sense, such object would encode in compact form the solutions of all possible planning problems in the environment. This a step towards building agents that are fully controllable after first exploring their environment in unsupervised way.
Goal-oriented RL methods 17 PAR+IS
compute policies for series of rewards specified in advance (such reaching set of target states), but cannot adapt in real time to new rewards, such as weighted combinations 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 [all7 KST+ 18).
Work done during an internship at Facebook Artificial Intelligence Research Paris. Code: https::/githen coom/aamed-eaee htte::////////ttto////////////// htt.aadd--coueee///ooooolllleeeo
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Instead, we exhibit an object that is both simpler to learn than 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 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 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 [BT021) We obtain the following.
We prove the existence of a learnable "summary" " of a reward freee discrete or continuous MDP, that provides an explicit formula for optimal policies for any reward specified later. This takes the form of pair of representations F: S x A Z Z and B: S A -> Z from state-actions into a representation space 7 = R2, with policies т. (s) == arg max, a F(s, a, 2) z. Once a reward is specified, value of is computed from reward values and B: then πy 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 well-principled unsupervised loss for F' and B. If FB training is perfect, then the policies are provably optimal for all rewards (Theorem is proportional to the FB training error (Theorems with large enough dimension d (Proposition 6
With imperfect training, sub-optimality In finite spaces, perfect training is possible
Explicitly, F and B are trained so that Fss, a, 2) B(s', a²) approximates the long-term probability to reach s' from if following π.. 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 [BTOII. No sparse rewards are used: every transition reaches some state s', 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 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 posteriori by trained goal-oriented models, we provide qualitative examples.
We also illustrate qualitatively the sub-optimalities (Iong-range behavior is preserved but local blurring of rewards occurs) and the representations learned.
2 Problem and Notation
Let M P,7) be 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(s'ls,a) a) from state to s' given action a, and discount factor 0 < 7 SBB. If S is finite, P(s'|s, a) can be viewed as a matrix; in general, for each (s,a) E S x A, P(ds' s, a) is a probability measure s' S. The notation P(ds') [s, a) covers all cases.
Given (Sg. ao) E S x A and a policy T: S + Prob(A), we denote (r(- ag, T) and E|-|so. ag. π. the probabilities and expectations under state-action sequences (st, at )=20 starting with (so. ag) and following policy π in the environment, defined by sampling s, ~ P(ds 9t- a(-1) and aj ~ π (st).
For any policy π and state-action (su, do). define the successor measure M (so, 4o, .) as the measure over 5 x A representing the expected discounted time spent in each set X C A:
M M Co; X) == Σ 7 Pr ((se, at) E X 80; ao; T)
for each X C S. X A. Viewing M as measure deals with both discrete and continuous spaces.
Given a reward function r: S x A + R, the Q-function of π for r is Q7 ag) == Ei>0 of E[r(s, (s)sss 0g. TT. 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 T (s)) = arg max, a ((s, a) instead of (s) E arg max, Q(s,a). Ties may be broken any way.

---

We consider the following informal problem: Given a reward-free MDP AS, P, can we compute a convenient learnable object E such that, once reward function r: 8 A - IR is specified. we can easily (with no planning) compute, from E. and r, policy whose performance is close to maximal?
3 Encoding All Optimal Policies via the Forward-Backward Representation
We first present forward-backward (FB) representations of reward-free MDP as way summarize all optimal policies via explicit formulas. The resulting learning procedure is described in Section
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 (s, a) = Esr M*(s, a, s', aa)r(s'',)))))) a') (s,a,
where M (8, s' a') = Σi20 7 Pr ((tt, ox) (s', a 11) a, T). This rewrites as QF = M viewing everything as vectors and matrices indexed by state-actions.
Now let (T=)>ER= be any family of policies parameterized by z. Assume that for each 2, we can find d x (S x A)-matrices Fz and B such that M = F B. Then Q = F.. Br. Specializing to 2R. So far n; was unspecified; but if we define (s) ;= arg max, (F' 2) at each state s, then by definition, RaR is the greedy policy with respect to FT. zR- At the same time, FT Ex 2R is the Q-function of Ten for reward r: thus, Ten is the greedy policy of its own Q-function, and therefore optimal for reward r.
2R := Br, the Q-function of policy Tzn on reward is Qr =n = =
Thus, if we manage to find F, B, and T. such that n. = arg max F and F = for all E Rd, then we obtain the optimal policy for any reward r, just by computing Br and applying policy TBr.
This criterion on (F, B, π2) is entirely unsupervised. Since F and B depend on тं but n. is defined via F, this is fixed point equation. An exact solution exists for d large enough (Appendix, Prop. 6 while a smaller d provides lower-rank approximations n x F B. In Section we present a well-grounded algorithm learn such F. B, and uz.
In short, we learn two representations F. and B such that F(s0, ag, 2) B(s', a') is approximately the long-term probability = do, a', a²) to reach (s', a') if starting at (50, an) and following policy T.. Then all optimal policies can be computed from and B. We think of F as a representation of the future of state, and B the ways reach state (Appendix B.4): if F 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, F; and B become functions S x A - Rdd instead of matrices; since F= depends on F itself is a function S x A x Rd - Rd The sums over states will be replaced with expectations under the data distribution p. Definition
representation). Let 7 = Rd be a representation space, and let be a measure on S x A. A pair of functions F: S A + Z and B: S A + Z, together with a parametric family of policies (Tz)s is called a barward-bakwaad with respect to p. if the following conditions hold for any z E Z and (s, a), (so, ao) S A:
fforward-hackwand
representation of the MDP
Te (s) arg max Fss, 0, 2)2, 2,
M "s (so, ao ds, da) F ao, T)TBB B(s, p)p(ds da) (2)
where R 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, (T))) be a -aforward-backward representation of a reward-free MDP with respect to some measure p.
Then, for any bounded reward function r: S x A + R. the following holds. Set
r(s, B(s, (ss, a) da). s,a
ZR
assuming the integral exists. Then WaR is optimal policy for reward in the MDP. Moreover, the optimal Q-function Q" for reward r is Q²(ssa) = F(s, a, T 2n-
For instance, for a single reward located at state-action (s, a), the optimal policy is TzR with in = Bs, a). (In that case the factor p(ds, da) does not matter because scaling the reward does not change the optimal policy.)
3

---

We present in Section an algorithm to learn FB representations. The measure p will be the distribution of state-actions visited in training set or under an exploration policy: then zn B(e.a)-p(((s a.) B(s,a) 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. G. In infinite spaces, arbitrarily good approximations can be obtained by increasing d, corresponding a rank-d approximation of the cumulated transition probabilities A/". Importantl, the optimality guarantee extends to approximate F' and B, with optimality gap proportional to F B F (Appendix, Theorems with various norms on F B M /p). For instance, if, for some reward r, the error [F(o, do, zn) B(s, a) = (8o; ag ds, da) p(ds, da) is at most on average over (s, a) ~ p for every (so, ao). then u is 3= ||rll /11 - for
"/=*/pp /p/- averaged over 2 E R", as training These results justify using some norm over F B - M
loss for unsupervised reinforcement learning. (Below, average over E R from 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 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.
4 Learning and Using Forward-Backwand Representations
Our algorithm starts with an unsupervised learning phase, where we learn the representations F and B in 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 ER B(r(s, a) B(s, a)] from some reward observations, or directly set 2R if the reward is known (c.g., set 2R = B(s, a) to reach a known target (s, a)). In the exploitation phase, we directly use the policy zn (s) arg max, F(s, a, 2R) 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, (s) = arg max,, F(s, a, 2) z, is just taken as the definition of π. given F. In turn, I' and / are trained so that the second condition (
2(.,2)TB B = M /p. holds for every Here p is the (unknown) distribution of state-actions in the training data. Training is based on the Bellman equation for the successor measure M
M (so, 4o, ((ssss''))) = a st, agma! + 7 Egg (idssPP(ddllesddd (e0.0)) (s1, ((),a (4)
We leverage well-principle algorithm from [T021 in the single-policy setting: it learns the successor measure of policy without using the sparse reward 1s0=s', ao=a' (which would vanish in continuous spaces). Other successor measure algorithms could be used, such as C-learning [ES.2
The algorithm from [TО21 uses parametric model m5 og ,aa)
represent (so,,, co ds', dd') se mž (s0, an, s', aa))))))) p(dd''ss da'). is not necessary to know p. only to sample states from it. Given an observed transition (so, ag, s1) from the training set, generate an action a1 ~ ग [ag]s8) and sample another state-action (s', aa) from the training set, independently from (Sp, ag: $1). Then update the parameter by 0 t + 17 80 with learning rate 17 and
50 = a (sp, ag. 80, on) + +gmg (So, 4o, s..')) x (7 mž (#1, .1,8', 8,, ,a'd)'==== mă (So, ag a')))) (5)
This computes the density m² of M with respect to the distribution p of state-actions in the training set. Namely, the true successor state density m = M²/p is a fixed point of (5) in expectation [BO21 Theorem 6] (and is the only fixed point in the tabular or overparameterized case). Variants exist, such as using a target network for m ($1, a1, s'a'') on the right-hand side, as in DQN.
Thus, we first choose a parametric model Fo. Bg for the representations F and B. and set m (so, d', aa) := Fo do, =) Bg a²). Then we iterate the update (5) over many state- actions and values of z. This results in Algorithm At each step, value of = is picked at random, together with a batch of transitions (sg. ao, 81) and a batch of state-actions (s',a') from the training set, with (s', aa) independent from and (so, ag 81).

---

For sampling 2, we use fixed distribution (rescaled Gaussians, see Appendix D) Any number of values of may be sampled: this does not use up training samples. We use target network with soft updates (Polyak averaging) as in DDPG. For training we also replace the greedy policy т= = arg max, F F(s, o, 2) with a regularized version =N = softmax F(s, a, =) 2)/ with fixed temperature r (Appendix D). Since there is unidentifabilltyy between F and B (Appendix, Remark we normalize B via an auxiliary loss in Algorithm D
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 n; for some time (Appendix. Algorithm D. 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) z by weighing the representation B by the reward:
2 == E[r(s, B), a)]
(6)
where the expectation must be computed over the same distribution p state-actions (s, a) used to learn F and B (see Appendix [B.5 for using 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 žR is obtained by averaging r(s, a) (ss, a) over the states visited. An approximate value for zg still provides an approximately optimal policy (Appendix, Prop. and Thm. [2].
If the reward is known explicitly, this phase is unnecessary. For instance, if the reward is to reach a target state-action (sp, an) while avoiding some forbidden acttonstiions (s1, a1).... (8k, az), one may directly set
2R = B(so. ao) - > B(s,, a;)
where the constant À sets the negative reward for forbidden states and adjusts for the unknown p(ds,, factors in 3 This can be used for goal-oriented RL.
If the reward is known algebraically as a function r(s, a), then zд may be computed by averaging the function r(s, a (ss, a) over a replay buffer from the unsupervised training phase. We may also use a reward model f(s, a) of r(s, a) trained on some reward observations from any source.
The exploitation phase. Once the reward representation žR has been estimated, the Q-function is estimated as
Q a) = a(s, 2R 2R.
The corresponding policy Tz (s) = arg max, 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, =R) žR + qo a) where the fine-tuning model qe 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 ( only has to learn the future probabilities of the part of interest in the future states, based on the full initial state (sg, ag). Explicitly, if pp: S A G is feature map to some features g = sp(s, a). and if we know that the reward will be a function R(g), then Theorem Sisemmsstill still holds with Bgg) everywhere instead of B(s, a), and with the successor measure " ag dg) instead of M (So. au ds', da') (Appendix, Theorem 4. Learning is done by replacing Öpmz (80: up, 80, (o) with dgmg (8o, on *(00, an)) in the first term in (5) [BTO2]] Rewards can be arbitrary functions of g. so this is more general than [BBQ which only considers rewards linear in g. For instance, in MsPacman below, we let y be the 2D position (z, y) of the agent, so we can optimize any reward function that depends on this position.
5

---

Limitations. First, this method does not solve exploration: it assumee access to good exploration strategy. (Here we used the policies я, with random values of 2, 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 test time based on the initial 1 B estimate.
As reward functions are represented by a d-dimensional vector ž = E-r.B), some information about the reward is necessarily lost. Any reward uncorrelated to B is treated as O. 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 z 100 manages navigation in pixel-based environment with 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, d can be taken 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, 1 B is optimized to approximate the successor measure M = >, 7 P' with Ps the t-step transition kernel for each policy т. The rank-d approximation will favor large eigenvectors of Pr. i.e., small eigenvectors of the Markov chain Laplacian Id P> These loosely correspond to long-range (low-frequeen behavior [MM07] presumably, F and B will learn spatially smooth rewards first. Indeed, experimentally, small d leads to spatial blurring of rewards and Q-functions (Fig. B Arguably, without any prior information this is reasonable prior. [SBGIZ have argued for the cognitive relevance of low-dimensional approximations of successor representations.
Variance is potential issue in larger environments, although this did not arise in our experiments. Learning MM requires sampling a state-action (so, (o) and an independent state-action (s'.a'). d'). In large spaces, most state-action pairs will be unrelated. A possible mitigation is to combine FB with strategies such as Hindsight Experience Replay [AC to select goals related to the current state-action. The following may help a lot: the update of F and B decouples over (sy, ao), times an expectation over (s',a') :a). Thus, by estimating these expectations by a moving average over dataset, it is easy to have many pairs (So, ao) interact with many (s', a'). The cost is handling full d x d matrices. This will be explored in future work.
expectation
5 Experiments
We first consider the task of reaching arbitrary goal states. For this, we can make quantitative comparisons existing goal-oriented baselines. Next, 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.
5.1 Environments and Experimental Setup
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 (xc, y) [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 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 variant of the Atari 2600 game Ms. Pacman, where an episode ends when the agent is captured by a monster [ 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 x 84 pixels. state 8; the concatenation of (I; Tt 41-4, at)) frames, i.e. an 84 x 84 x 4 tensor. An action repeat of 12 is used. As Ms. Pacman is not originally multi-goal domain, we define the goals as the 148 reachable coordinates (æt, 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 uniform policy while for FetchReach and Ms. Pacman, we use an E-greedy policy with respect to the current approximation F(s, a, z) T z for a sampled z. At evaluation time, s-greedy policies are also used, with smaller E. More details are given in Appendix D
5.2 Goal-Oriented Setting: Quantitative Comparisons
We investigate the FB representation over goal-reaching tasks and compare it goal-oriented basclines: DQN and DQN with HER when needed. We define sparse reward functions. For Discrete Maze, the reward function is equal one when the agent's state is equal exactly 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 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 goal.
Figure 1: Comparative performance FB for dif- ferent dimensions DQN FetchReach. Left: success rate averaged over 20 randomly selected goals as function of the first 100 training epochs. Right: success rate averaged over 20 random goals after 800 training epochs.
Comparative performance of for dif- ferent dimensions and DQN in Ms. Pacman. Left: success rate averaged over 20 randomly selected goals as function of the first 200 training epochs. Right: success rate averaged over the goal space after 800 training epochs.
Figure 2:
Figs. I and 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 learn and had to add HER make 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. this due to the spatial smoothing induced by the small rank approximation and the reward being nonzero only the agent is exactly at the goal. This spatial blurring is clear on heatmaps for d 25 vs d = 75 (Fig. With d 25 the agent often stops right next to its goal.
To evaluate the sample efficiency FB, after each epoch, we evaluate the agent on 20 randomly selected goals. Learning curves are reported in Figs. and 2 (left). In all environments, we observe no loss sample efficiency compared goal-oriented baseline. Ms. Pacman, FB even learns faster than HQN+HER.
5.3 More Complex Rewards: Qualitative Results
We now investigate FB's ability generalize new tasks that cannot be solved by already trained goal-oriented model: reaching goal with forbidden states imposed a posteriori, reaching the nearest of two goals, and choosing between a small, close reward and large, distant one.
First, for the task of reaching a target position go d while avoiding some forbidden positions B(g)) Ek= B(g;) and run the corresponding e-greedy policy defined
I1.... we set RR
Here DQN is short for goal-oriented DQN, Q(s, a)
7

---

Figure 7: Visualization of FB embedding vectors on Continuous Maze after projecting them in two-dimensional space with t-SNE. Left: the states be mapped. Middle: the F embedding. Right: the B embedding. The walls appear as large dents; the smaller dents correspond to the number of steps needed to get past wall.
by F(s, 2R) zn. Fig. shows the resulting trajectories, which succeed at solving the task for the different domains. In Ms. Pacman, the path suboptimal (though successful) due to the sudden appearance of monster along the optimal path. (We only plot the initial frame: see the full series of frames along the trajectory in Appendix D Fig. [6] Fig. (left) provides a contour plot of max (8, a, =R) 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 go and Ø1, by setting R (goo) + B(g1). The optimal Q-function not linear combination of the Q-functions for and Fig. 6 shows successful trajectories generated by the policy Zza- On the contour plot of m F(s, a, =R) zR in Fig. (right), the two rewarding positions appear basins of attraction. Similar results for third task are shown in Appendix [D] introducing "distracting" small reward next to the initial position of the agent, with larger reward further away.
Figure
Contour plot of 2R in Continuous Maze. Left: for the task of reaching a target while avoiding a forbidden region, Right: for two equally rewarding targets.
4:
maXaga F(s, a, 2R) T
Figure 3: Heatmap of max. F(s, a, an)) an for
B(xx) Left: d 25. Right: d = 75.
Figure Trajectories generated by the F B policies for the task of reaching the closest among equally rewarding positions (star shapes 1t). (Optimal Q Q-values are not linear over such mix- tures.)
Trajectories generated by the F` B policies for the task of reaching target position (star shape while avoiding forbidden positions (red shape
Figure 5:
5.4 Embedding Visualizations
visualize the learned state embeddings for Continuous Maze by projecting them into dimensional space using t-SNE [VdMH08] in Fig. 7 For the forward embeddings. we set 0 corresponding to the uniform policy. We can see that partitions states according to the topology induced by the dynamics: states on opposite sides of walls separated in the representation space

---

and states on the same side lie together. Appendix D includes embedding visualizations for different and for Discrete Maze and Ms. Pacman.
6
Related work
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 [ЕТОЛ]] where [BBQ 18 successor features. Schematically, and omitting actions, successor features start with user pser-provided features yp, then learn e such that (b())) = Eiz0 7 [8(((())) so]. This limits applicability rewards that are linear combinations of y. Here use successor state probabilities, namely, learn two representations F' and B such that F(5g) B(s') = Σ1z0 7 Pr(s, s' | 5p). This does not require any user-provided input.
Thus learn representations instead of one. The learned backward representation B is absent from [BQQ+18 18). B plays different role than the user- user-provided features of (BBQ+ 18): if the reward is known priori to depend only on some features y. we learn on top of (p, 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 yp B +p and &p Id, and then only training P.,
We use a similar parameterization of policies by F(s, a, 2 as in BBQ 18 for similar reasons, although encodes different object.
Successor representations where first defined in Day93 for finite spaces, corresponding to an older object from Markov chains, the fundamental matrix KS60 Bre99) GS97.. [BG17] argue for their relevance for cognitive science. For successor representations in continuous spaces, finite number of features
specified first; this can be used for generalization within a family of tasks, c.g.. BDM ZSBB17 GHB 19 HDB+ 19). [BTO2] moves from successor features to successor states by providing pointwise occupancy map estimates even in continuous spaces, without using the sparse reward 1,,=s'. We borrow a successor state learning algorithm from [BT21] [BT02]] also introduced simpler versions of F and B for single, fixed policy; [BT021 does not consider the overy-optimal-pone ppolicy setting.
There is long literature on goal-oriented RL. For instance, SSGGS15 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) I 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 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, [KSY20 use reward-free interactions to build a training set that summarizes finite environment, in the sense that any optimal policies later computed on this training set instead of the true environment are provably e-optimal, for any reward. They prove tight bounds on the necessary set size. Policy learning still has be done afterwards for each reward.
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 improved text.
Acknowledgments.

---

7 Conclusion
The forward-backwart representation is learnable mathematical object that "summarizes" reward- free MDP. It provides near-optimal policies for any reward specified a posteriori, without planning, It 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.
References
[ACR+17] Marcin Andrychowicz, Dwight Crow, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In NIPS, 2017.
[ARO*19] 19] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, R Devon Hjelm. Unsupervised state representation learning in atari. arXiv preprint arXiv: (00.00260.0 2019.
Côté, and
[BBQ* + 18] Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado van Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. arXiv preprint arXiv:1812.0720 2018.
17] André Barreto, Will Dabney, Rémi Munos, Jonathan Hunt, Tom Schaul, David Silver, and Hado P van Hasselt. Successor features for transfer in reinforcement learning. In NIPS, 2017.
[BNVB13] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents, Journal of Artificial Intelligence Research, 47:253-279 2013.
[Bog07] Vladimir Bogachev. Measure theory. Springer, 2007.
[Bré99] Pierre Brémaud. Markov chains: Gibbs fields, Monte Carlo simulation, and queues, volume 31. 1999.
[BTO21] Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal- dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123 2021.
[Day93] Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation, 5(4):613-624, 1993.
[ESL21] Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning to achieve goals via recursive classification. In International Conference on Learning Representations, 2021.
[FD02] David Foster and Peter Dayan. Structure the space of value functions. Machine Learning, 49(2):325-346, 2002.
[GHB+19] 19] Christopher Grimm, Irina Higgins, Andre Barreto, Denis Teplyashin, Markus Wulfmeier, Tim Hertweck, Raia Hadsell, and Satinder Singh. Disentangled cumulants help successor representations transfer to new tasks. arXiv preprint 11Xiv:1911.10866 2019.
[GS97] Charles Miller Grinstead and James Laurie Snell. Introduction probability. American Mathematical Soc., 1997.
HDB + 19] Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. arXiv preprint arXiv: 19066050300 2019.
[JSY200] Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free ex- ploration for reinforcement learning. In International Conference on Machine Learning, pages 4870-4879, PMLR, 2020.
10

---

[KS60] J. G. Kemeny and J. L. Snell. Finite Markov Chains. Van Nostrand, New York, 1960.
[KST 18] Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Ben- gio, Devi Parikh, and Dhruv Batra. Modeling the long term future in model-based reinforcement learning. In International Conference Learning Representations, 2018.
[MM07] Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A Laplacian framework for learning representation and control in Markov decision processes. Journal of Machine Learning Research, 8(10), 2007.
[PAR + 18] Matthias Plappert, Marcin Andrychowicz. Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforcement learning: Challenging roboties environments and request for research. arXiv preprint arXiv: wwXXiv:802.09466 2018.
[RUMS18] Paulo Rauber, Avinash Ummadisingu, Filipe Mutz, and Jürgen Schmidhuber. Hindsight policy gradients. In International Conference on Learning Representations, 2018.
[SB18] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018. 2nd edition.
[SBG17] Kimberly Stachenfeld, Matthew M Botvinick, and Samuel Gershman. The hip- pocampus as a predictive map. Nature neuroscience, 20(1 1):1643, 2017.
[SHGS15] Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning. volume 37 of Proceedings of Machine Learning Research, pages 1312 1320, Lille, France, 07-09 Jul 2015. PMLR.
[SMD+11] Richard S Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, pages 761-768. 2011.
[SPS99] Richard Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):111-211, 1999.
[Tl17]] Erik Talvitie. Self-correctigg models for model-based reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
[VdMH08] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
[ZSBB17] Jingwei Zhang, Jost Tobias Springenberg, Joschka Boedecker, and Wolfram Burgard. Deep reinforcement learning with successor features for navigation across similar en- vironments. In 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 2371-2377 IEEE, 2017.


Reviews:
[{"review_id": "xr0BGoDazhm", "reviewer": "Reviewer_ojoH", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Novelty: The paper proposed a novel representation learning idea to learn the forward-backward features $F$ and $B$.  The authors constructed $F$ and $B$ according to a novel observation that $F^\\top B$ should approximate the probability of visiting one state-action to another. They also proposed new unsupervised learning method to train $F$ and $B$ based on a new density measure. \n\nQuality & Clarity: The claims are clear and well-supported to me. The proofs are clean and correct. The organization is okay, but some statements in several sections are obscure and took quite a while to understand. For example, the authors introduced their ideas in the related work section, in a fashion of story-telling. This type of writing may fail to highlight the critical points of the paper, and make the readers confusing because they are not comfortable with the notations. \n\nSignificance: The paper studied a very important problem for both theoretical RL and empirical RL: how to “summarize” an environment, so as to remove the useless information and keep the important part. The FB representation provided valuable insights on this problem. The experiments on Maze and Ms. Pacman also verified that FB is comparable to the classic value-based methods, but FB is more flexible to random goals.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations have been addressed in the paper, including the lack of exploration, the large variance, etc. The lack of exploration can be a serious problem, because rewards in real-world are mostly not random. Perhaps this problem is related to some techniques in batch RL. Moreover, it would be better to show more experimental results, such as in more realistic environments, or compared with more goal-oriented baseline algorithms. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "XoBRlbP_S1h", "reviewer": "Reviewer_MgL9", "summary": "This paper proposes an algorithm that is able to learn the inherent structure of a reward-free MDP in a compact form, which can quickly give a nearly optimal policy after the reward function is specified. The algorithm is developed through a rigorous theoretical analysis and the experiment also shows comparable or even better performance than the algorithm trained with specified rewards.", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I appreciated the novel and rigorous theoretical framework developed for learning MDP representation in this paper. Meanwhile, the experimental result looks promising and the extended analysis in appendix looks thorough. However, the writing of this paper may need some improvement since currently the ordering of some content makes it harder to understand, details in \"Suggestion on writing\" part.\n\nMeanwhile, I also have the following concerns and questions.\n\nQuestions:\n- On page 6, why it is okay to write $z_R$ as stated in Eq. (9) when the reward is known explicitly? Is Eq. (9) consistent with the definition of $z_R$ in Eq. (8)?\n- In the experiment for FetchReach and Ms. Pacman, the behavior policy becomes adaptive, which implies a changing distribution $\\rho$; meanwhile, the theory uses a fixed distribution $\\rho$. Is there any intuition on why a changing distribution $\\rho$ is better under this scenario?\n- Does the distribution of $z$ play any role in the theory development? It seems $z$ is always fixed in theory while the distribution $\\nu$ of $z$ is very carefully designed in experiment.\n\nSuggestions on writing:\n- At page 5, it may be better to explicitly point out that Theorem 6 in [BTO21] gives the update rule in Eq. (7).\n- It may be better to move the explanations about $\\mathscr{L}$ and $\\mathscr{L}_{\\mathsf{reg}}$ at Section D.3 to Section A right below the full Algorithm 1. Or more aggressively, it may be better to move all of these into the main text of Section 4, which can help a lot more in understanding the algorithm than reading the paragraph description in current Section 4. This extra space can be saved from moving some figures and experimental results into the Appendix.\n- In Remark 4, page 17, it may be better to say explicitly that $B'$ being $L^2(\\rho)$-orthonormal means $\\mathbb{E}_{(s, a)\\sim\\rho}\\left[B'(s, a)B'(s, a)^\\top\\right]=\\mathbf{I}$ if I understand correctly.\n\n```\n[BTO21] Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123, 2021.\n```\n\n---\n\nGood paper. My score remains the same after rebuttal.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations of their work in Section 4 by pointing out the potential problems of their algorithm. The potential negative social impact is not addressed, which will generally be all potential malicious usage of any reinforcement learning algorithms.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "SZEzJvdlXLa", "reviewer": "Reviewer_NkUB", "summary": "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.\n\nThe 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$.\nThen, 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$.\n\nThe 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$. \n\nThe 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$.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea is brilliant, and I think will have a strong impact on the topic of UVFA. \n\nHowever, reading this paper was quite painful, in the sense that it is really badly constructed. A lot of quantities are introduced before the notation is explained, or ideas spoiled before the content makes it understandable. I agree that in the abstract and introduction the authors don't have the choice and must explain the idea before having the mathematical notation and tools. But --in my case-- I finally understood what the authors were talking about  at line 183 with the \"sketch of proof\" (page 5). \nI would suggest a lot of improvements regarding the construction of the paper, but to rather summarize:\n\n1) separate introduction and related work, and move the related work later in the paper (for ex before or after experiments)\n2) improve the introduction by a better explanation \"with hands\" of what $F$ and $B$ represent (future and past), how to use $B$ to compute $z_r$ and then $F$ to get the optimal policy. \n3) in problem and notation, give the problem after the notation.\n4) right after the notation, briefly explain the idea like I did in the summary of this review.\n5) put elements of your proof in the paper, so that you prepare the reader to understand the TH1\n6) ends up with TH1\n\nIn the current state, we don't understand the quantities before TH1, then we don't understand TH1, then there is a long explanation that we finally understand by reading the short \"sketch of proof\". \n\nAnother problem of this paper regards the experiment design: did the author used more than one seed? if yes how many? (especially for the pacman expe). The pacman results are promising (FB works better than HER), it would have been great to have more examples of such complex environments.\n\nI think this paper is already nice and important to be communicated, but I would be a bit disappointed because of its actual form (difficult to read), so I am on the fence, marginally above the acceptance.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This paper has no societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "Iq5u6t_vMBh", "reviewer": "Reviewer_Bh3C", "summary": "This paper proposes a new representation in MDPs such that it can be learned using transition samples and used to compute the optimal policy. I am not convinced by the motivation of such a new representation since no computational or statistical efficiency results are established and thus suggest not accepting the paper.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposes a new representation in MDPs such that it can be learned using transition samples and used to compute the optimal policy. \n\nMy main concern is that such a representation might be too restrictive for learning unknown MDPs. As one can always learn the transition kernel well without knowing any reward signals and later compute an optimal policy by a dynamic program procedure using the learned model and reward function provided a posterior, the proposed new algorithm just provides an alternative which might save some computation in the second step while not sure adds how much restriction on the first step. For the tabular MDPs, the dimension d of is already |SA| as shown in the appendix and I am not sure how large the computation cost can be to learn the map F: S $\\times$ A $\\times$ Z $\\rightarrow$ Z and B. To me, it feels equal to or more expensive than learning the transition kernel matrix. Also, it is not clear whether this new representation can help improve the exploration efficiency in the reward-free exploration phase, which is usually the result of many other reward-free works, e.g., see [1]. \n\nIn all, I am not convinced by the motivation of such a new representation since no computational or statistical efficiency results are established and thus suggest not accepting the paper.\n\n[1] Jin, Chi, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. \"Reward-free exploration for reinforcement learning.\" In International Conference on Machine Learning, pp. 4870-4879. PMLR, 2020.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The majority of the reviews agrees that the paper presents an interesting new approach and votes for a clear accept. Accordingly, I recommend acceptance. As most reviewers also pointed out a few issues with the presentation and organisation of the paper, please take into account these comments when preparing the final version. Finally, one reviewer was not that completely convinced by the suggested representation. In the discussion (s)he said (I think expressing the point better than in the original review):

"My concern is that the representation carries too much information. To be specific, (please correct me if my statements are wrong), in their tabular MDP example, the function $F$ has a large continuous input space (dimension>SA) and thus requires function approximations and large computation power. In contrast, learning the MDP only involves $S^2A$ exact parameters, and the planning under a given MDP and reward function is usually not the bottleneck (using dynamic programming and given a discount factor in (0,1)). This issue can be even worse for general MDPs, and the paper seems to not discuss too much on the approximation power of such a representative for other MDPs."

Although I understand that the paper suggests not to learn a particular MDP but a range of tasks, the question whether the arising complexity of the final problem is reasonable when compared to the complexity of a set of single tasks is I think a valid one and should be discussed.
Decision: Accept (Poster)