Learning One Representation to Optimize All Rewards
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 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.
The unsupervised FB 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 Ms. Pacman, and the FetchReach virtual robot arm.
We also illustrate how the agent can immediately adapt to new tasks beyond goal-oriented RL.
Introduction and Related Work
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 (Andrychowicz et al., 2017;
Plappert et al., 2018) 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 (Talvitie, 2017; Ke et al., 2018).
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.
Borsa et al.
(2018) 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 Borsa et al.
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 Blier et al.
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 𝐹 : 𝑆 × 𝐴 × 𝑍 → 𝑍 and 𝐵 : 𝑆 × 𝐴 → 𝑍 from state-actions into a representation space 𝑍 ≃ R 𝑑 , with policies 𝜋 𝑧 ( 𝑠 ) :
= arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 .
Once a reward is specified, a value of 𝑧 is computed from reward values and 𝐵 ; then 𝜋 𝑧 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 𝐹 and 𝐵 : if training is perfect, then the policies are provably optimal for all rewards (Theorem 1).
In finite spaces, perfect training is possible with large enough dimension 𝑑 (Proposition 3).
With imperfect training, sub-optimality is proportional to the training error (Theorem 5).
∙ We provide a TD-like algorithm to train 𝐹 and 𝐵 to minimize this loss, with function approximation, adapted from recent methods for successor states Blier et al.
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 compare this method to 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.
Problem and Notation
We consider a reward-free Markov decision process (MDP) ℳ = ( 𝑆, 𝐴, 𝑃, 𝛾 ) with state space 𝑆 (discrete or continuous), action space 𝐴 (discrete for simplicity, but this is not essential)
, transition probabilities 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 ) from state 𝑠 to 𝑠 ′ given action 𝑎 , and discount factor 0 < 𝛾 < 1 Sutton & Barto (2018).
If 𝑆 is finite, 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 ) can be viewed as a matrix; in general, for each ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 , 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) is a probability measure on 𝑠 ′ ∈ 𝑆 .
The notation 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) covers all cases.
All functions are assumed to be measurable.
Given ( 𝑠 0 , 𝑎 0 ) ∈
𝑆 × 𝐴 and a policy 𝜋 : 𝑆 → Prob( 𝐴 ), we denote Pr( ·| 𝑠 0 , 𝑎 0 , 𝜋 ) and E [ ·| 𝑠 0 , 𝑎 0 , 𝜋 ] the probabilities and expectations under state-action sequences ( 𝑠 𝑡 , 𝑎 𝑡 ) 𝑡 ≥starting with ( 𝑠 0 , 𝑎 0 ) and following policy 𝜋 in the environment, defined by sampling 𝑠 𝑡 ∼ 𝑃 (d 𝑠 𝑡 | 𝑠 𝑡 − 1 , 𝑎 𝑡 − 1 ) and 𝑎 𝑡 ∼ 𝜋 ( 𝑠 𝑡 ).
Given a reward function 𝑟 : 𝑆 × 𝐴 → R , the 𝑄 -function of 𝜋 for 𝑟 is 𝑄 𝜋𝑟 ( 𝑠 0 , 𝑎 0 ) :
= ∑︀ 𝑡 ≥ 0 𝛾 𝑡 E [ 𝑟 ( 𝑠 𝑡 , 𝑎 𝑡 ) | 𝑠 0 , 𝑎 0 , 𝜋 ].
We assume that rewards are bounded, so that all 𝑄 -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 𝜋 ( 𝑠 ) =
arg max 𝑎 𝑄 ( 𝑠, 𝑎 ) instead of 𝜋 ( 𝑠 ) ∈ arg max 𝑎 𝑄 ( 𝑠, 𝑎 ).
Ties may be broken any way.
We consider the following informal problem: Given a reward-free MDP ( 𝑆, 𝐴, 𝑃, 𝛾 ), can we compute a convenient object 𝐸 such that, once a reward function 𝑟 :
𝑆 × 𝐴
→ R is specified, we can easily (with no planning) compute from 𝐸 and 𝑟 a policy 𝜋 such that 𝑄 𝜋𝑟 is close to maximal?
Outline of the Method
The method has three phases: an unsupervised learning phase where we learn a pair of representations 𝐹 and 𝐵 in a reward-free way, by observing state transitions in the environment; a reward estimation phase where we estimate a policy parameter 𝑧 𝑅 from reward observations, or directly set 𝑧 𝑅 if the reward is known (as in goal-oriented RL); and an exploitation phase , where we directly set a policy based on 𝐹 , 𝐵 , and 𝑧 𝑅 .
The last two phases may be merged in some cases.
The unsupervised learning phase.
We set a representation space 𝑍
= R 𝑑 , used to represent both states and reward functions.
We learn a pair of “forward” and “backward” representations
[Equation]
For each 𝑧 ∈ 𝑍 , we define the policy 𝜋 𝑧 ( 𝑠 ) :
= arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 .
No rewards are used in this phase, and no family of tasks has to be specified manually.
𝐹 and 𝐵 are trained such that, for any state-actions ( 𝑠 0 , 𝑎 0 ), ( 𝑠 ′ , 𝑎 ′ ), and any 𝑧 ∈ 𝑍 ,
[Equation]
(We actually learn probability densities, which make more sense in continuous spaces, see (7)–(8).)
Since 𝜋 𝑧 depends on 𝐹 , (2) is a fixed point equation, similarly to the dependency between the optimal 𝑄 -function and optimal policy in ordinary 𝑄 -learning.
Intuitively, 𝐹 is a representation of the future of a state under a certain policy.
𝐵 represents the past of a state, or the ways to reach that state (see also Appendix E.5).
If 𝐹 ⊤ 𝐵 is large, then it is possible to reach the second state from the first.
This is akin to a model of the environment, without synthesizing state trajectories.
This unsupervised learning phase is fully principled: If learning of 𝐹 and 𝐵 is successful, then this is guaranteed to provide all optimal policies (Theorem 1).
Namely, for any reward function 𝑟 ( 𝑠, 𝑎 ), the optimal policy for 𝑟 is the policy 𝜋 𝑧 with 𝑧 :=
E ( 𝑠,𝑎 ) [ 𝑟 ( 𝑠, 𝑎 ) 𝐵
( 𝑠, 𝑎 )].
For instance, 𝜋 𝐵 ( 𝑠,𝑎 ) is the optimal policy to reach ( 𝑠, 𝑎 ).
Approximate solutions still provide approximately optimal policies, with optimality gap directly proportional to the error on 𝐹 ⊤ 𝐵 (Appendix, Theorem 5), so the guarantees are non-empty.
𝐹 and 𝐵 can be learned off-policy from observed transitions in the environment, via the
Bellman equation for (2),
[Equation]
in discrete spaces.
We leverage efficient algorithms for successor states Blier et al.
(2021), also valid in continuous spaces, and which learn without using the sparse reward 1 𝑠 0 = 𝑠 ′ ,𝑎 0 = 𝑎 ′ (see Section 4).
In finite spaces, exact solutions 𝐹 and 𝐵 exist (Appendix, Prop. 3), provided the dimension 𝑑 is large enough.
In infinite spaces, arbitrarily good approximations can be obtained by increasing 𝑑 , corresponding to a rank- 𝑑 approximation of the pairwise cumulated transition probabilities in (2).
The dimension 𝑑 controls how many types of rewards can be optimized well; even a relatively small 𝑑 can provide useful behaviors, see the experiments and Appendix E.2.
The algorithm is linear in 𝑑 , so 𝑑 can be taken as large as the neural network models can handle.
For exploration in this phase, we use the policies being learned: the exploration policy chooses a random value of 𝑧 from some distribution (e.g., Gaussian), and follows 𝜋 𝑧 for some time (Appendix, Algorithm 1).
However, the algorithm can also work from an existing dataset of off-policy transitions.
The optimality guarantee and training algorithm are detailed in Section 4.
The reward estimation phase.
Once rewards are available, we estimate a reward representation (policy parameter) 𝑧 𝑅 by weighing the representation 𝐵 by the reward:
[Equation]
where the expectation must be computed over the same distribution of state-actions ( 𝑠, 𝑎 ) used to learn 𝐹 and 𝐵 .
There are several scenarios.
If the reward is known explicitly, this phase is unnecessary.
For instance, if the reward is to reach a target state-action ( 𝑠 0 , 𝑎 0 ) while avoiding some forbidden state-actions ( 𝑠 1 , 𝑎 1 ) , ..., ( 𝑠 𝑘 , 𝑎 𝑘 ), one may directly set
[Equation]
where the constant 𝜆 adjusts the negative reward for visiting a forbidden state.
This can be used for goal-oriented RL.
If the reward is known algebraically as a function 𝑟 ( 𝑠, 𝑎 ), then 𝑧 𝑅 may be computed by averaging the function 𝑟 ( 𝑠, 𝑎 ) 𝐵
( 𝑠, 𝑎 ) over a replay buffer from the unsupervised training phase.
If the reward is black-box as in standard RL algorithms, then the exploration policy has to be run again for some time, and 𝑧 𝑅 is obtained by averaging 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 ) over the states visited.
If it is not possible to run the exploration policy again, one may learn a model ^ 𝑟 ( 𝑠, 𝑎 ) of 𝑟 ( 𝑠, 𝑎 ) on some reward observations from any source, then estimate 𝑧 𝑅 =
E
[^ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )] by averaging the model over a replay buffer from the unsupervised training phase.
In this case, reward identification may be done in parallel to exploitation.
An approximate value for 𝑧 𝑅 still provides an approximately optimal policy (Appendix,
Prop. 6 and Thm.
Since reward functions are represented by 𝑑 -dimensional vectors 𝑧 𝑅 , loss of precision necessarily occurs.
In practice, for small 𝑑 we notice some blurring of rewards between nearby states (Fig. 3), for reasons discussed in Section 4.
The exploitation phase.
Once the reward representation 𝑧 𝑅 has been estimated, the 𝑄 -function is estimated as
[Equation]
The corresponding policy 𝜋 𝑧 𝑅 ( 𝑠 ) =
arg max 𝑎 𝑄 ( 𝑠, 𝑎 ) is used for exploitation.
Fine-tuning was not needed in our experiments, but it is possible to fine-tune the 𝑄 -function using actual rewards, by setting 𝑄 ( 𝑠, 𝑎 ) =
𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 + 𝑞 𝜃 ( 𝑠, 𝑎 ) where the fine-tuning model 𝑞 𝜃 is initialized to 0 and learned via any standard 𝑄 -learning method.
The Forward-Backward Representation of an MDP Encodes the Optimal Policies
We now present in more detail the algorithm used to learn 𝐹 and 𝐵 , and the associated theoretical guarantee.
Background on successor states.
We start with some definitions and an algorithm from
Blier et al.
(2021) for successor states of a single policy.
The successor state operator 𝑀 𝜋 of a policy 𝜋 is defined as follows.
For each state-action ( 𝑠 0 , 𝑎 0 ), 𝑀 𝜋 ( 𝑠 0 , 𝑎 0 , · ) is a measure over the state-action space 𝑆 × 𝐴 , representing the expected discounted time spent in each set 𝑋 ⊂ 𝑆 × 𝐴 , namely,
[Equation]
for each 𝑋 ⊂ 𝑆 × 𝐴 .
Viewing 𝑀 as a measure deals with both discrete and continuous spaces.
In practice, 𝑀 𝜋 is represented by a function 𝑚 𝜋 ( 𝑠 0 , 𝑎 0 , 𝑠 ′ , 𝑎 ′ ) taking a pair of state-actions and returning a number.
Namely, assume access to a training set of transitions in the MDP (off-policy, reward-free).
Let 𝜌 (d 𝑠, d 𝑎 ) be the probability distribution of state-actions in this training set; we assume 𝜌 > 0 everywhere.
The measure 𝑀 𝜋 ( 𝑠 0 , 𝑎 0 , · ) may be represented as a density with respect to 𝜌 :
[Equation]
where 𝑚 𝜋 is an ordinary function (or a distribution, see Appendix E.4).
Algorithms exist to train a parametric model 𝑚 𝜋𝜃 ( 𝑠 0 , 𝑎 0 , 𝑠 ′ , 𝑎 ′ ) of 𝑚 𝜋 from a training set of off-policy observed transitions ( 𝑠 0 , 𝑎, 𝑠 1 ) in the MDP, based on a Bellman-like equation for 𝑀 Blier et al.
Namely, sample a transition ( 𝑠 0 , 𝑎, 𝑠 1 ) from the training set, generate an action 𝑎 1 ∼ 𝜋 ( 𝑎 1 | 𝑠 1 ), sample a state-action ( 𝑠 ′ , 𝑎 ′ ) from the training set, independently from ( 𝑠 0 , 𝑎, 𝑠 1 ).
Then update the parameter 𝜃 by 𝜃 ← 𝜃 + 𝜂 𝛿𝜃 with learning rate 𝜂 and
[Equation]
[Equation]
The true successor state density 𝑚 𝜋 is a fixed point of this update in expectation Blier et al.
(2021) (it is the only fixed point in the tabular or overparameterized case).
Variants exist, such as using a target network for 𝑚 𝜋𝜃 ( 𝑠 1 , 𝑎 1 , 𝑠 ′ , 𝑎 ′ ) on the right-hand side, as in DQN.
The naive way to estimate successor states via a Bellman equation from (2) would involve a sparse reward term 1 𝑠
= 𝑠 ′ ,𝑎 = 𝑎 ′ , resulting in infinitely sparse rewards in continuous spaces.
However, in expectation, this term’s contribution is known algebraically: (9) avoids sparse rewards in a principled way thanks to the 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠 0 , 𝑎, 𝑠 0 , 𝑎 ) term Blier et al.
It is not necessary to know the training distribution 𝜌 : all quantities are expressed as expectations under 𝜌 , namely, expectations over the training set.
When training under data distribution 𝜌 , one learns the density 𝑚 𝜋 of 𝑀 𝜋 with respect to 𝜌 ; this will be useful below.
The forward-backward representation provides the optimal policies.
The next theorem contains the core idea of our algorithms.
In short, if we can learn two representations 𝐹 and 𝐵 of state-actions such that 𝐹 ⊤ 𝐵 approximates the successor states of certain policies, then we can compute all optimal policies from 𝐹 and 𝐵 .
Theorem 1 (Forward-backward representation of an MDP) .
Consider an MDP with state space 𝑆 and action space 𝐴 .
Let 𝑍
= R 𝑑 be some representation space.
Let
[Equation]
be two functions.
For each 𝑧 ∈ 𝑍 , define the policy
[Equation]
Let 𝜌 be any probability distribution on 𝑆 × 𝐴 (e.g., the distribution of state-actions under some exploration scheme).
Assume that 𝐹 and 𝐵 have been chosen (trained) to satisfy the following: for any 𝑧 ∈ 𝑍 , and
any state-actions ( 𝑠, 𝑎 ) and ( 𝑠 ′ , 𝑎 ′ ) , the quantity 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ ) is equal to the successor state density 𝑚
𝜋 𝑧
( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ ) of policy 𝜋 𝑧 with respect to 𝜌 .
Then, for any bounded reward function 𝑟 : 𝑆 × 𝐴 → R , the following holds.
Set
[Equation]
Then 𝜋 𝑧 𝑅 is an optimal policy for reward 𝑟 in the MDP.
Moreover, the optimal 𝑄 -function 𝑄 * for reward 𝑟 is
[Equation]
The theoretical guarantee extends to approximate training of 𝐹 and 𝐵 , with optimality gap proportional to the 𝐹 ⊤ 𝐵 training error.
This is important, as exact finite- 𝑑 representations are often impossible in continuous spaces.
Namely: if, for some reward 𝑟 , the error ⃒⃒ 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 )
⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ )
− 𝑚 𝜋 𝑧𝑅 ( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ )
⃒⃒ is at most 𝜀 on average over ( 𝑠 ′ , 𝑎 ′ ) ∼ 𝜌 for every ( 𝑠, 𝑎 ), then 𝜋 𝑧 𝑅 is 3 𝜀 ‖ 𝑟 ‖ ∞ / (1 − 𝛾 )-optimal
for 𝑟 (Appendix, Theorem 5, with additional results on weaker norms).
A similar statement holds if we use ¯ 𝑚 ( 𝑠, 𝑧, 𝑠 ′ , 𝑎 ′ )
+ 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ ) instead of just 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ ) to learn successor states (Appendix, Theorem 2).
Here ¯ 𝑚 is any function that does not depend on 𝑎 .
Since ¯ 𝑚 has no rank restriction, the finite rank approximation only applies to the advantage function.
Learning 𝐹 and 𝐵 .
Theorem 1 suggests to choose a parametric model for the representations 𝐹 and 𝐵 , and then train 𝐹 and 𝐵 such that 𝐹 ⊤ 𝐵 approximates the successor state density 𝑚 , via any successor state learning algorithm.
This unfolds as follows (Appendix, Algorithm 1).
Assume access to a training dataset of transitions ( 𝑠 0 , 𝑎, 𝑠 1 ) in the MDP, e.g., from some exploration policies.
The samples ( 𝑠 0 , 𝑎 ) in this dataset follow some unknown distribution 𝜌 ; we make no assumptions on 𝜌 .
Then:
∙ In the unsupervised learning phase, train 𝐹 and 𝐵 so that 𝐹 ⊤ 𝐵 approximates the successor state density 𝑚 𝜋 𝑧 for every 𝑧 .
This can use any successor state learning algorithm, adapted to incorporate 𝑧 .
Here we use (9).
For this, choose a parametric model 𝐹 𝜃 , 𝐵 𝜃 for the representations 𝐹 and 𝐵 .
At each step, pick a 𝑧 at random, pick a batch of transitions ( 𝑠 0 , 𝑎, 𝑠 1 ) from the training set, set the next actions 𝑎 1 :
= 𝜋 𝑧 ( 𝑠 1 ), and pick a batch of state-actions ( 𝑠 ′ , 𝑎 ′ ) from the training set, independently from 𝑧 and ( 𝑠 0 , 𝑎, 𝑠 1 ).
Then apply the successor state update (9) to the parametric model 𝑚 𝜋 𝑧 𝜃 ( 𝑠 0 , 𝑎, 𝑠 ′ , 𝑎 ′ )
=
𝐹 𝜃 ( 𝑠 0 , 𝑎, 𝑧 )
⊤ 𝐵 𝜃 ( 𝑠 ′ , 𝑎 ′ )
(Algorithm 1).
For sampling 𝑧 , we use a fixed distribution (rescaled Gaussians, see Appendix G).
Any number of values of 𝑧 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 𝜋 𝑧 = arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 with a regularized version 𝜋 𝑧 = softmax( 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝑧/𝜏 ) with fixed temperature 𝜏 (Appendix G).
Since there is unidentifiability between 𝐹 and 𝐵 (Appendix, Remark 4), we normalize 𝐵 via an auxiliary loss in Algorithm 1.
∙ In the reward estimation phase, we estimate the reward representation 𝑧 𝑅 via (12), following one of the use cases from Section 3.
The distribution 𝜌 in (12) is the same as the one used to define the density 𝑚 𝜋 ; thus, 𝑧 𝑅 must be estimated over the same distribution of state-actions used to train 𝐹 and 𝐵 .
If the reward is known directly, 𝑧 𝑅 can be set by hand following (12).
For instance, for the optimal policy to reach state-action ( 𝑠, 𝑎 ), we directly set 𝑧 𝑅 =
𝐵 ( 𝑠, 𝑎 ).
(A unit reward at ( 𝑠, 𝑎 ) corresponds to 𝑧 𝑅 =
𝐵 ( 𝑠, 𝑎 ) /𝜌
( 𝑠, 𝑎 ) in (12), but 𝜌 is unknown and scaling the reward yields the same optimal policy.)
This extends to finite combinations of states.
If 𝑟 is provided as a function, then E ( 𝑠,𝑎 ) ∼ 𝜌 [ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )] may be estimated directly on ( 𝑠, 𝑎 ) in a replay buffer.
The case of 𝑟 accessible only by samples has been discussed in Section 3.
Stochastic rewards may be used to estimate 𝑧 𝑅 .
Indeed, if E ˜ 𝑟 ( 𝑠, 𝑎 ) =
𝑟 ( 𝑠, 𝑎 ) then E [˜ 𝑟
( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )]
= 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 ).
∙ In the exploitation phase, the exploitation policy can be directly set to 𝜋 𝑧 𝑅 , which selects the action arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 .
Incorporating prior information.
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 𝐵 can be trained on that part of the state only, with the same theoretical guarantees (Appendix, Theorem 2).
𝐹 still needs to use the full state as input.
This way, the FB model (2) does not have to learn how often every ( 𝑠 ′ , 𝑎 ′ ) is reached, only the part of interest in ( 𝑠 ′ , 𝑎 ′ ).
More generally, if 𝜙 : 𝑆 × 𝐴 → 𝐺 is a feature map to some features 𝑔 = 𝜙 ( 𝑠, 𝑎 ), and if we know that the reward will be a function 𝑅 ( 𝑔 ), then the same theorem holds with 𝐵 ( 𝑔 ) everywhere instead of 𝐵 ( 𝑠, 𝑎 ), and with the successor density 𝑚 𝜋 ( 𝑠, 𝑎, 𝑔 ) instead of 𝑚 𝜋
( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ )
(Appendix, Theorem 2).
Learning the latter is done by replacing 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠 0 , 𝑎, 𝑠 0 , 𝑎 ) with 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠 0 , 𝑎, 𝜙 ( 𝑠 0 , 𝑎 )) in the first term in (9) Blier et al.
Rewards can be arbitrary functions of 𝑔 , so this is more general than Borsa et al.
(2018) which only considers rewards linear in 𝑔 .
For instance, in
MsPacman below, we let 𝑔 be the position of the agent, so we can optimize any reward function that depends on this position.
Experiments
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.
5. 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 ( 𝑥, 𝑦 ) ∈ [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 (Plappert et al., 2018) 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 (Rauber et al., 2018).
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 𝑠 𝑡 is the concatenation of ( 𝑥 𝑡 − 12 , 𝑥 𝑡 − 8 , 𝑥 𝑡 − 4 , 𝑥 𝑡 ) 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 ( 𝑥, 𝑦 ) on the screen; these can be reached only by learning to avoid monsters.
For all environments, we run algorithms for 800 epochs.
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 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 for a sampled 𝑧 .
At evaluation time, 𝜀 -greedy policies are also used, with a smaller 𝜀 .
More details are given in Appendix G.
5. Goal-Oriented Setting: Quantitative Comparisons
[Figure] [Figure]
We investigate the FB representation over goal-reaching tasks and compare it to goal-oriented baselines: DQN 1 , 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 𝑑 , and
DQN and DQN+HER respectively in Discrete Maze and Ms. Pacman (similar results in
FetchReach and Continuous Maze are provided in the Appendix).
The performance of FB consistently increases with the dimension 𝑑 and the best dimension matches the performance of the goal-oriented baseline.
In Ms. Pacman, DQN totally fails to learn and we had to add
HER to make it work.
In Discrete Maze, we observe a drop of performance for 𝑑 = 25: 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 𝑑 = 25 vs 𝑑 = 75 (Fig. 3).
With 𝑑 = 25 the agent often stops right next to its goal.
To evaluate the sample efficiency of FB, after each epoch, we evaluate the agent onrandomly 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.
5. More Complex Rewards: Qualitative Results
[Figure]
[Figure]
[Figure]
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 𝑔while avoiding some forbidden positions 𝑔 1 , . . .
𝑔 𝑘 , we set 𝑧 𝑅 =
𝐵 ( 𝑔 0 ) − 𝜆 ∑︀ 𝑘𝑖 =1 𝐵 ( 𝑔 𝑖 ) and run the corresponding 𝜀 -greedy policy defined by 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 .
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; the full series of frames along the trajectory is in the Appendix.)
Fig.
4 (left) provides a contour plot of max 𝑎 ∈ 𝐴 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 )
⊤ 𝑧 𝑅 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 𝑔 0 and 𝑔 1 , by setting 𝑧 𝑅 =
𝐵 ( 𝑔 0 )
+ 𝐵 ( 𝑔 1 ).
The optimal 𝑄 -function is not a linear combination of the 𝑄 -functions for 𝑔 0 and 𝑔 1 .
Fig.
6 shows successful trajectories generated by the policy 𝜋 𝑧 𝑅 .
On the contour plot of max 𝑎 ∈ 𝐴 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 in Fig.
4 (right), the two rewarding positions appear as basins of attraction.
Similar results for a third task are shown in the Appendix: introducing a “distracting” small reward next to the initial position of the agent, with a larger reward further away.
The Appendix includes embedding visualizations for different 𝑧 and for Discrete Maze and Ms. Pacman.
Conclusion
The FB representation is a learnable mathematical object that “summarizes” a reward-free
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.
A Outline of the Supplementary Material
The Appendix is organized as follows.
∙ Appendix B provides an extended related work.
∙ Appendix C discusses how incoporate prior information as well as what kinod of representation is learned.
∙ Appendix D presents the pseudo-code of the unsupervised phase of FB algorithm.
∙ Appendix E provides extended theoretical results on approximate solutions and general goals: – Section E.1 formalizes the forward-backward representation with a goal or feature space.
– Section E.2 establishes the existence of exact FB representations in finite spaces.
– Section E.3 shows how approximate solutions provide approximately optimal policies.
– Section E.4 presents a note of the measure 𝑀 𝜋 and its density 𝑚 𝜋 .
– Section E.5 shows how 𝐹 and 𝐵 are successor and predecessor features of each other.
∙ Appendix F provides proofs of all theoretical results above.
∙ Appendix G provides additional information about our experiments: – Section G.1 describes the environments.
– Section G.2 describes the different architectures used for FB as well as the goal-oriented DQN.
– Section G.3 provides implementation details and hyperparameters.
– Section G.4 provides additional experimental results.
B Extended Related Work
Borsa et al.
(2018) learn optimal policies for rewards that are linear combinations of a finite number of feature functions provided in advance by the user.
The approach of Borsa et al.
(2018) 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 Blier et al.
where Borsa et al.
(2018) use successor features .
Schematically, and omitting actions, successor features start with user-provided features 𝜙 , then learn 𝜓 such that 𝜓 ( 𝑠 0 ) =
∑︀ 𝑡 ≥
𝛾 𝑡 E [ 𝜙 ( 𝑠 𝑡 ) | 𝑠 0 ].
This limits applicability to rewards that are linear combinations of 𝜙 .
Here we use successor states , namely, we learn two representations 𝐹 and 𝐵 such that 𝐹 ( 𝑠 0 )
⊤ 𝐵 ( 𝑠 ′ )
=
∑︀ 𝑡 ≥
𝛾 𝑡 Pr( 𝑠 𝑡 = 𝑠 ′ | 𝑠 0 ).
This does not require any user-provided input.
Thus we learn two representations instead of one.
The learned backward representation 𝐵 is absent from Borsa et al.
𝐵 plays a different role than the user-provided features 𝜙 of
Borsa et al.
(2018): learning of 𝐹 given 𝐵 is not learning of 𝜓 given 𝜙 .
The features 𝜙 of
Borsa et al.
(2018) are split between our learned 𝐵 and the functions 𝜙 we can use if the reward is known to depend on only part of the state.
We use a similar parameterization of policies by 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 as in Borsa et al.
(2018), for similar reasons, although 𝑧 encodes a different object.
Successor representations where first defined in Dayan (1993) for finite spaces, corresponding to an older object from Markov chains, the fundamental matrix Kemeny & Snell (1960);
Brémaud (1999); Grinstead & Snell (1997).
Stachenfeld et al.
(2017) 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., (Barreto et al., 2017; Zhang et al., 2017; Grimm et al., 2019;
Hansen et al., 2019).
Blier et al.
(2021) moves from successor features to successor states by providing pointwise occupancy map estimates even in continuous spaces, without using the sparse reward 1 𝑠 𝑡 = 𝑠 ′ .
We adapt successor state learning algorithms from Blier et al.
(2021), which also introduced simpler versions of 𝐹 and 𝐵 for a single, fixed policy.
There is a long literature on goal-oriented RL.
For instance, Schaul et al.
(2015) learn goaldependent value functions, regularized via an explicit matrix factorization.
Goal-dependent value functions have been investigated in earlier works such as (Foster & Dayan, 2002) and (Sutton et al., 2011).
Hindsight experience replay (HER) (Andrychowicz et al., 2017) improve 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.
For finite state spaces, Jin et al.
(2020) provide an algorithm to constitute a training set by reward-free interactions with a finite environment, such 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.
Planning still has to be done for each reward.
C Discussion
Incorporating prior information.
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 𝑑 to accomodate all rewards.
Prior information can be incorporated in several ways into the 𝐹𝐵 method, by focussing on certain types of rewards or parts of the space.
First, 𝐹 and 𝐵 may be trained only on values of 𝑧 𝑅 corresponding to a category of rewards specified in advance.
For instance, training only on 𝑧 𝑅 =
𝐵 ( 𝑠 ) with 𝑠 in some set 𝑆 ′ ⊂
𝑆 learns the optimal policies for reaching all states 𝑠 ∈ 𝑆 ′ .
Any prior (e.g., Bayesian) on rewards may be incorporated this way.
In our experiments, 𝑧 was just sampled from a rescaled Gaussian.
Second, 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 𝐵 can be trained on that part of the state only, with the same theoretical guarantees (Appendix, Theorem 2).
𝐹 still needs to use the full state as input.
This way, the FB model (2) does not have to learn how often every ( 𝑠 ′ , 𝑎 ′ ) is reached, only the part of interest in ( 𝑠 ′ , 𝑎 ′ ).
More generally, if 𝜙 : 𝑆 × 𝐴 → 𝐺 is a feature map to some features 𝑔 = 𝜙 ( 𝑠, 𝑎 ), and if we know that the reward will be a function 𝑅 ( 𝑔 ), then the same theorem holds with 𝐵 ( 𝑔 ) everywhere instead of 𝐵 ( 𝑠, 𝑎 ), and with the successor density 𝑚 𝜋 ( 𝑠, 𝑎, 𝑔 ) instead of 𝑚 𝜋
( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ )
(Appendix, Theorem 2).
Learning the latter is done by replacing 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠 0 , 𝑎, 𝑠 0 , 𝑎 ) with 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠 0 , 𝑎, 𝜙 ( 𝑠 0 , 𝑎 )) in the first term in (9) Blier et al.
Rewards can be arbitrary functions of 𝑔 , so this is more general than Borsa et al.
(2018) which only considers rewards linear in 𝑔 .
For instance, in MsPacman below, we let 𝑔 be the position of the agent, so we can optimize any reward function that depends on this position.
Third, if we are interested only in part of the space, focusing the exploration policy on that part will provide representations 𝐹 and 𝐵 that minimize the loss over that part of the space.
This could lead to optimal policies for rewards “around a typical set of behaviors”, presumably a more realistic goal than fully agnostic unsupervised training in complex environments.
What kind of representation is learned?
Which information on the reward gets lost?
As reward functions are represented by a finite-dimensional object 𝑧 𝑅 , some approximation is incurred on rewards (unless fine-tuning is used at exploitation time).
By (12), a reward is represented by 0 if and only if it is 𝐿 2 ( 𝜌 )-orthogonal to all components of
𝐵 .
So the representation 𝐵 is not only part of a transition probability model, but also a reward regularizer.
The loss associated to (2), and minimized by a Bellman-like training, is
[Equation]
where 𝑃 𝑡 is the distribution of states reached at step 𝑡 by policy 𝜋 𝑧 starting at ( 𝑠, 𝑎 ), and 𝑃 𝑧 is some user-chosen distribution, such as a Gaussian or a prior over tasks (see above).
Thus, 𝐹 and 𝐵 provide a rank- 𝑑 approximation to the transition probabilities in the environment.
For a single 𝑧 , this loss is minimized when 𝐹 and 𝐵 are the SVD of the cumulated transition matrix
∑︀ 𝑡 𝛾 𝑡 𝑃 𝑡 , truncated to the largest 𝑑 singular values.
The SVD is computed in the 𝐿 2 ( 𝜌 ) norm: the data distribution influences the representations learned by focusing on the part of the space covered by 𝜌 .
For a fixed policy, on a finite space, ∑︀ 𝑡 𝛾 𝑡 𝑃 𝑡 =
(Id − 𝛾𝑃 ) − 1 as matrices.
Thus the largest singular values of ∑︀ 𝑡 𝛾 𝑡 𝑃 𝑡 correspond to the smallest singular values of the Markov chain
Laplacian Id − 𝛾𝑃 .
The associated singular vectors loosely correspond to long-range (lowfrequency) behavior Mahadevan & Maggioni (2007): presumably, 𝐹 and 𝐵 will represent spatially smooth rewards first, while local detail may be lost (Fig. 3 and Section 5).
Stachenfeld et al.
(2017) argue for the cognitive relevance of low-dimensional approximations of successor representations.
The features learned in 𝐹 and 𝐵 may have broader interest.
D Algorithm
The unsupervised phase of the FB algorithm is described in Algorithm 1.
do 3: /* Collect episodes 4: for episode 𝑒 = 1 , . . .
𝐸 do 5:
Sample 𝑧 ∼ 𝜈 6:
Observe an initial state 𝑠7: for 𝑡 = 1 , . .
do 8:
Select an action 𝑎 𝑡 according to some behaviour policy (e.g the 𝜀 -greedy with respect to 𝐹 𝜃 𝑠 𝑡 , 𝑎, 𝑧 ) ⊤ 𝑧 ) 9:
Observe next state 𝑠
Store transition ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 1 ) in the replay buffer 𝒟 11: end for 12: end for 13: /* Perform 𝑁 stochastic gradient descent updates 14: for 𝑛 = 1 .
𝑁 do 15:
Sample a mini-batch of transitions , 𝑠 𝑖 +1 ) } 𝑖 ∈ 𝐼 ⊂ 𝐼 |
Sample a mini-batch of target state-action pairs
( 𝑠 𝑖 , 𝑎 ′ 𝑖 ) } 𝑖 ∈ 𝐼 ⊂
| 𝐼 | = 𝑏 .
Sample a mini-batch of { 𝑧 𝑖 } 𝑖 ∈
𝐼 ∼ 𝜈 of size | 𝐼 | = 𝑏 .
Set 𝜋
[Equation]
[Equation]
𝑧 𝑖 )
⊤ 𝐵 𝜔 − ( 𝑠 ′ 𝑗 , 𝑎 ′ 𝑗 ) )
︀ 2 − 1 𝑏
∑︀ 𝑖 ∈
𝐼 𝐹 𝜃 ( 𝑠 𝑖 , 𝑎 𝑖 , 𝑧 𝑖 )
⊤ 𝐵 𝜔 ( 𝑠 𝑖 , 𝑎 𝑖 ) 20: /* Compute orthonormality regularization loss 21:
L reg ( 𝜔 ) = 1 𝑏 2 ∑︀ 𝑖,𝑗 ∈
𝐼 2 𝐵 𝜔 ( 𝑠 𝑖 , 𝑎 𝑖 )
⊤ stop-gradient ( 𝐵 𝜔 ( 𝑠 ′ 𝑗 , 𝑎 ′ 𝑗 )) · stop-gradient ( 𝐵 𝜔 ( 𝑠 𝑖 , 𝑎 𝑖 )
⊤ 𝐵 𝜔 ( 𝑠 ′ 𝑗 , 𝑎 ′ 𝑗 )) − 1 𝑏 ∑︀ 𝑖 ∈
𝐼 𝐵 𝜔 ( 𝑠 𝑖 , 𝑎 𝑖 )
⊤ stop-gradient ( 𝐵 𝜔 ( 𝑠 𝑖 , 𝑎 𝑖 )) 22:
Update 𝜃 ← 𝜃 − 𝜂 ∇ 𝜃 L ( 𝜃, 𝜔 ) and 𝜔 ← 𝜔 − 𝜂 ∇ 𝜔 ( L ( 𝜃, 𝜔 ) + 𝜆 · L reg ( 𝜔 )) 23: end for 24: /
* Update target network parameters 25: 𝜃 − ← 𝛼𝜃 − + (1 − 𝛼 ) 𝜃 26: 𝜔 − ← 𝛼𝜔 − + (1 − 𝛼 ) 𝜔 27: end for
E Extended Results: Approximate Solutions and General Goals E. The Forward-Backward Representation With a Goal or Feature Space
Here we state a generalization of Theorem 1 covering the extensions mentioned in Section 4.
First, this covers rewards depending on the state-action ( 𝑠, 𝑎 ) only via certain features 𝑔 = 𝜙 ( 𝑠, 𝑎 ) where 𝜙 is a known function from state-actions to some goal state 𝐺 (for instance, rewards depending only on some components of the state).
Then it is enough to compute 𝐵 as a function of the goal 𝑔 .
Theorem 1 corresponds to 𝜙 =
Id.
This also allows us to cover successor features as in Borsa et al.
(2018), defined by userprovided features 𝜙 .
Indeed, fixing 𝐵 to Id and setting our 𝜙 to the 𝜙 of Borsa et al.
(2018) (or fixing 𝐵 to their 𝜙 and our 𝜙 to Id) will represent the same set of rewards and policies as in Borsa et al.
, although with a slightly different learning algorithm and up to a linear change of variables for 𝐹 and 𝑧 (given by the covariance of 𝜙 , see Appendix E.5), namely, optimal policies for rewards linear in 𝜙 .
However, keeping the same 𝜙 but letting 𝐵 free (with larger 𝑑 ) can provide optimal policies for rewards that are arbitrary functions of 𝜙 , linear or not.
For this, we extend successor state measures to values in goal spaces, representing the discounted time spent at each goal by the policy.
Namely, given a policy 𝜋 , let 𝑀 𝜋 be the the successor state measure of 𝜋 over goals 𝑔 :
[Equation]
for each state-action ( 𝑠, 𝑎 ) and each measurable set d 𝑔 ⊂ 𝐺 .
This will be the object approximated by 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝐵 ( 𝑔 ).
Second, we use a more general model of successor states, 𝑚 ≈ 𝐹 ⊤ 𝐵 + ¯ 𝑚 where ¯ 𝑚 does not depend on the action, so that the 𝐹 ⊤ 𝐵 part is enough to compute advantages; this lifts the constraint that the model of 𝑚 has rank at most 𝑑 , because there is no restriction on the rank on ¯ 𝑚 .
For simplicity we state the result with deterministic rewards, but this extends to stochastic rewards, because the expectation 𝑧 𝑅 will be the same.
Theorem 2 (Forward-backward representation of an MDP, with features as goals) .
Consider an MDP with state space 𝑆 and action space 𝐴 .
Let 𝜙 : 𝑆 × 𝐴 → 𝐺 be a function from state-actions to some goal space 𝐺
= R 𝑘 .
Let 𝑍
=
R 𝑑 be some representation space.
Let
[Equation]
be three functions.
For each 𝑧 ∈ 𝑍 , define the policy
[Equation] [Equation]
Let 𝜌 be any probability distribution over the goal space 𝐺 .
For each 𝑧 ∈ 𝑍 , let 𝑚 𝜋 𝑧 be the density of the successor state measure of 𝜋 𝑧 with respect to 𝜌 :
𝑀 𝜋 𝑧
( 𝑠, 𝑎, d 𝑔 ) =
𝑚
𝜋 𝑧
( 𝑠, 𝑎, 𝑔 ) 𝜌 (d 𝑔 ) .
Assume that 𝐹 and 𝐵 have been chosen (trained) to satisfy the following: for any 𝑧 ∈ 𝑍 , any state-actions ( 𝑠, 𝑎 ) , and any goal 𝑔 ∈ 𝐺 , one has
[Equation]
Let 𝑅 : 𝑆 × 𝐴 → R be any bounded reward function, and assume that this reward function depends only on 𝑔 = 𝜙 ( 𝑠, 𝑎 ) , namely, that there exists a function 𝑟 : 𝐺
→ R such that 𝑅 ( 𝑠, 𝑎 )
=
𝑟 ( 𝜙 ( 𝑠, 𝑎 )) .
Set
[Equation]
Then:
𝜋 𝑧 𝑅 is an optimal policy for reward 𝑅 in the MDP. 2.
For any 𝑧 ∈ 𝑍 , the 𝑄 -function of policy 𝜋 𝑧 for the reward 𝑅 is equal to
[Equation]
and the optimal 𝑄 -function is obtained when 𝑧 = 𝑧 𝑅 .
The advantages 𝑄 𝜋 𝑧 ( 𝑠, 𝑎 ) − 𝑄 𝜋
𝑧
( 𝑠, 𝑎 ′ ) do not depend on ¯ 𝑉 .
Here ¯
[Equation]
and in particular ¯ 𝑉 =
if ¯ 𝑚 = 0 .
If ¯ 𝑚 = 0 , then for any state-action ( 𝑠, 𝑎 ) one has
[Equation]
(We do not claim that ¯ 𝑉 is the value function and 𝐹 ⊤ 𝑧 𝑅 the advantage function, only that the sum is the 𝑄 -function.
When ¯ 𝑚 = 0, the term 𝐹 ⊤ 𝑧 𝑅 is the whole 𝑄 -function.)
The last point of the theorem is a form of policy improvement.
Indeed, by the second point, 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 𝑅 is the estimated 𝑄 -function of policy 𝜋 𝑧 for rewards 𝑟 .
If 𝑧 𝑅 falls outside of the training distribution for 𝐹 , the values of 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) may not be safe to use; in that case, it may be useful to use a finite set 𝑍 ′ ⊂ 𝑍 of values of 𝑧 closer to the training distribution, and use the estimate sup 𝑧 ∈ 𝑍 ′ 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 𝑅 instead of 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 .
This has been used, e.g., in
Borsa et al.
(2018), but in the end it was not necessary in our experiments.
E. Existence of Exact 𝐹𝐵 Representations in Finite Spaces
We now prove existence of an exact solution for finite spaces if the representation dimension 𝑑 is at least # 𝑆 × # 𝐴 .
Solutions are never unique: one may always multiply 𝐹 by an invertible matrix 𝐶 and multiply 𝐵 by ( 𝐶 ⊤ ) − 1 , see Remark 4 below (this allows us to impose orthonormality of 𝐵 in the experiments).
The constraint 𝑑 ≥ # 𝑆 × # 𝐴 can be largely overestimated depending on the tasks of interest, though.
For instance, we prove below that in an 𝑛 -dimensional toric grid 𝑆 = { 1 , . . .
, 𝑘 } 𝑛 , 𝑑 = 2 𝑛 is enough to obtain optimal policies for reaching every target state (a set of tasks smaller than optimizing all possible rewards).
Proposition 3 (Existence of an exact 𝐹𝐵 representation for finite state spaces) .
Assume that the state and action spaces 𝑆 and 𝐴 of an MDP are finite.
Let 𝑍 =
R 𝑑 with 𝑑 ≥ # 𝑆 × # 𝐴 .
Let 𝜌 be any probability distribution on 𝑆 × 𝐴 , with 𝜌 ( 𝑠, 𝑎 ) > 0 for any ( 𝑠, 𝑎 ) .
Then there exists 𝐹 : 𝑆 × 𝐴 × 𝑍 → 𝑍 and 𝐵 : 𝑆 × 𝐴 → 𝑍 , such that 𝐹 ⊤ 𝐵 is equal to the successor state density of 𝜋 𝑧 with respect to 𝜌 :
[Equation]
for any 𝑧 ∈ 𝑍 and any state-actions ( 𝑠, 𝑎 ) and ( 𝑠 ′ , 𝑎 ′ ) , where 𝜋 𝑧 is defined as in Theoremby 𝜋 𝑧 ( 𝑠 ) =
arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝑧 .
In practice, even a small 𝑑 can be enough to get optimal policies for reaching arbitrary many states (as opposed to optimizing all possible rewards).
Let us give an example with 𝑆 a toric 𝑛 -dimensional grid of size 𝑘 .
Let us start with 𝑛 = 1.
Take 𝑆 = { 0 , . .
, 𝑘 − 1 } to be a length- 𝑘 cycle with three actions 𝑎 ∈ {− 1 , 0 , 1 } (go left, stay in place, go right).
Take 𝑑 = 2, so that 𝑍 = R 2 ≃ C .
We consider the tasks of reaching an arbitrary target state 𝑠 ′ , for every 𝑠 ′ ∈ 𝑆 .
Thus the goal state is 𝐺 = 𝑆 in the notation of Theorem 2, and 𝐵 only depends on 𝑠 ′ .
The policy for such a reward is 𝜋 𝑧 𝑅 =
𝜋 𝐵 ( 𝑠 ′ ) .
For a state 𝑠 ∈ { 0 , . . .
, 𝑘 − 1 } and action 𝑎 ∈ {− 1 , 0 , 1 } , define
[Equation]
Then one checks that 𝜋 𝐵 ( 𝑠 ′ ) is the optimal policy for reaching 𝑠 ′ , for every 𝑠 ′ ∈ 𝑆 .
Indeed, 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 = cos(
𝜋
( 𝑠 + 𝑎 − 𝑠 ′ )
/𝑘 ).
This is maximized for the action 𝑎 that brings 𝑠 closer to 𝑠 ′ .
So the policies will be optimal for reaching every target 𝑠 ′ ∈ 𝑆 , despite the dimension being only 2.
By taking the product of 𝑛 copies of this example, this also works on the 𝑛 -dimensional toric grid 𝑆 = { 0 , . . .
, 𝑘 − 1 } 𝑛 with 2 𝑛 + 1 actions (add ± 1 in each direction or stay in place), with a representation of dimension 𝑑 = 2 𝑛 in C 𝑛 , namely, by taking 𝐵 ( 𝑠 ) 𝑗 := 𝑒 2 𝑖𝜋𝑠
𝑗
/𝑘 for ecah direction 𝑗 and likewise for 𝐹 .
Then 𝜋 𝐵 ( 𝑠 ′ ) is the optimal policy for reaching 𝑠 ′ for every 𝑠 ′ ∈ 𝑆 .
More generally, if one is only interested in the optimal policies for reaching states, then it is easy to show that there exist functions 𝐹 :
𝑆 × 𝐴 → 𝑍 and 𝐵 : 𝑆 → 𝑍 such that the policies 𝜋 𝑧 describe the optimal policies to reach each state: it is enough that 𝐵 be injective (typically requiring 𝑑 = dim( 𝑆 )).
Indeed, for any state 𝑠 ∈ 𝑆 , let 𝜋 * 𝑠 be the optimal policy to reach 𝑠 .
We want 𝜋 𝑧 to be equal to 𝜋 * 𝑠 for 𝑧 =
𝐵 ( 𝑠 ) (the value of 𝑧 𝑅 for a reward located at 𝑠 ).
This translates as arg max 𝑎 𝐹 ( 𝑠 ′ , 𝑎, 𝐵 ( 𝑠 ))
⊤ 𝐵 ( 𝑠 ) =
𝜋 * 𝑠 ( 𝑠 ′ ) for every other state 𝑠 ′ .
This is realized just by letting 𝐹 be any function such that 𝐹 ( 𝑠 ′ , 𝜋 * 𝑠 ( 𝑠 ′ ) , 𝐵 ( 𝑠 )) :
= 𝐵 ( 𝑠 ) and 𝐹 ( 𝑠 ′ , 𝑎, 𝐵 ( 𝑠 )) :
= − 𝐵 ( 𝑠 ) for every other action 𝑎 .
As soon as 𝐵 is injective, there exists such a function 𝐹 .
(Unfortunately, we are not able to show that the learning algorithm reaches such a solution.)
Let us turn to uniqueness of 𝐹 and 𝐵 .
Remark 4.
Let 𝐶 be an invertible 𝑑 × 𝑑 matrix.
Given 𝐹 and 𝐵 as in Theorem 1, define 𝐵 ′
[Equation]
𝑠
[Equation]
𝑎
[Equation]
together with the policies 𝜋 ′ 𝑧 ( 𝑠 ) :
= arg max 𝑎 𝐹 ′
( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 .
For each reward 𝑟 , define 𝑧 ′ 𝑅 :
= E ( 𝑠,𝑎 ) ∼ 𝜌 [ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ′
( 𝑠, 𝑎 )] .
Then this operation does not change the policies or estimated 𝑄 -values: for any reward, we have 𝜋 ′ 𝑧 ′ 𝑅
=
𝜋 𝑧 𝑅 , and 𝐹 ′ ( 𝑠, 𝑎, 𝑧 ′ 𝑅 ) ⊤ 𝑧 ′ 𝑅
=
𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 .
In particular, assume that the components of 𝐵 are linearly independent.
Then, taking 𝐶 =
(︀ E ( 𝑠,𝑎 ) ∼ 𝜌 𝐵 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 ) ⊤ )︀ − 1 / 2 , 𝐵 ′ is 𝐿 2 ( 𝜌 ) -orthonormal.
So up to reducing the dimension 𝑑 to rank( 𝐵 ) , we can always assume that 𝐵 is 𝐿 2 ( 𝜌 ) -orthonormal.
Reduction to orthonormal 𝐵 will be useful in some proofs below.
Even after imposing that 𝐵 be orthonormal, solutions are not unique, as one can still apply a rotation matrix on the variable 𝑧 .
E. Approximate Solutions Provide Approximately Optimal Policies
Here we prove that the optimality in Theorems 1 and 2 is robust to approximation errors during training.
We deal in turn with the approximation errors on ( 𝐹, 𝐵 ) during unsupervised training, and on 𝑧 𝑅 during the reward estimation phase.
E.3. Influence of Approximate 𝐹 and 𝐵
In continuous spaces, Theorems 1 and 2 are somewhat spurious: the equality 𝐹 ⊤ 𝐵 = 𝑚 will never hold exactly with finite representation dimension 𝑑 .
Instead, 𝐹 ⊤ 𝐵 will only be a rank- 𝑑 approximation of 𝑚 .
Even in finite spaces, since 𝐹 and 𝐵 are learned by a neural network, we can only expect that 𝐹 ⊤ 𝐵 ≈ 𝑚 in general.
We now prove that approximate solutions still provide approximately optimal policies.
We start with the error from learning 𝐹 and 𝐵 in the unsupervised learning phase.
Then we turn to the error from estimating 𝑧 𝑅 in the reward estimation phase (in case the reward is not known explicitly).
We provide this result for different notions of approximate solutions for 𝐹 and 𝐵 : first, in sup norm over ( 𝑠, 𝑎 ) but in expectation over ( 𝑠 ′ , 𝑎 ′ )
(so that a perfect model of the world is not necessary); second, for the weak topology on measures (this is the most relevant in continuous spaces: for instance, a Dirac measure can be approached by a continuous model in the weak topology).
𝐹 and 𝐵 are trained such that 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ ) approximates the successor state density 𝑚
𝜋 𝑧
( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ ).
In the simplest case, we prove that if for some reward 𝑅 ,
[Equation]
for every ( 𝑠, 𝑎 ), then the optimality gap of policy 𝜋 𝑧 𝑅 is at most (3 𝜀/ (1 − 𝛾 )) sup | 𝑅 | for that reward (Theorem 5, first case).
In continuous spaces, 𝑚 𝜋 is usually a distribution (Appendix E.4), so such an approximation will not hold, and it is better to work on the measures themselves rather than their densities, namely, to compare 𝐹 ⊤ 𝐵𝜌 to 𝑀 𝜋 .
We prove that if 𝐹 ⊤ 𝐵𝜌 is close to 𝑀 𝜋 in the weak topology, then the resulting policies are optimal for any Lipschitz reward.
Remember that a sequence of nonnegative measures 𝜇 𝑛 converges weakly to 𝜇 if for any bounded, continuous function 𝑓 , ∫︀ 𝑓 ( 𝑥 ) 𝜇 𝑛 ( 𝑑𝑥 ) converges to ∫︀ 𝑓 ( 𝑥 ) 𝜇 (d 𝑥 ) (Bogachev (2007), §8.1).
The associated topology can be defined via the following Kantorovich–Rubinstein norm on nonnegative measures (Bogachev (2007), §8.3)
[Equation]
where we have equipped the state-action space with any metric compatible with its topology.
The following theorem states that if 𝐹 ⊤ 𝐵 approximates the successor state density of the policy 𝜋 𝑧 (for various sorts of approximations), then 𝜋 𝑧 is approximately optimal.
Given a reward function 𝑟 on state-actions, we denote
[Equation]
and
[Equation]
where we have chosen any metric on state-actions.
The first statement is for any bounded reward.
The second statement only assumes an 𝐹 ⊤ 𝐵 approximation in the weak topology but only applies to Lipschitz rewards.
The third statement is more general and is how we prove the first two: weaker assumptions on 𝐹 ⊤ 𝐵 work on a stricter class of rewards.
Theorem 5 (If 𝐹 and 𝐵 approximate successor states, then the policies 𝜋 𝑧 yield approximately optimal returns) .
Let 𝐹 : 𝑆 × 𝐴 × 𝑍 → 𝑍 and 𝐵 :
𝑆 × 𝐴
→ 𝑍 be any functions, and define the policy 𝜋 𝑧 ( 𝑠 )
=
arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝑧 for each 𝑧 ∈ 𝑍 .
Let 𝜌 be any positive probability distribution on 𝑆 × 𝐴 , and for each policy 𝜋 , let 𝑚 𝜋 be the density of the successor state measure 𝑀 𝜋 of 𝜋 with respect to 𝜌 .
Let ^ 𝑚
[Equation]
𝑠
[Equation]
𝑎 [Equation] 𝑠 ′
[Equation] [Equation]
be the estimates of 𝑚 and 𝑀 obtained via the model 𝐹 and 𝐵 .
Let 𝑟 : 𝑆 × 𝐴 → R be any bounded reward function.
Let 𝑉 * be the optimal value function for this reward 𝑟 .
Let ^
𝑉 𝜋 𝑧 be the value function of policy 𝜋 𝑧 for this reward.
Let 𝑧 𝑅 = E ( 𝑠,𝑎 ) ∼ 𝜌 [ 𝑟 ( 𝑠, 𝑎 ) 𝐵
( 𝑠, 𝑎 )] .
Then:
If E ( 𝑠 ′ ,𝑎 ′ ) ∼ 𝜌 | ^ 𝑚 𝑧 𝑅
( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ ) − 𝑚 𝜋 𝑧𝑅 ( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ )
| ≤ 𝜀 for any ( 𝑠, 𝑎 ) in 𝑆 × 𝐴 , then ⃦⃦⃦ ^
𝑉 𝜋 𝑧𝑅 − 𝑉 * ⃦⃦⃦ ∞ ≤ 3 𝜀 ‖ 𝑟 ‖ ∞ / (1 − 𝛾 ) .
If 𝑟 is Lipschitz and ⃦⃦⃦ ^ 𝑀 𝑧 𝑅 ( 𝑠, 𝑎, · ) − 𝑀 𝜋 𝑧𝑅 ( 𝑠, 𝑎, · ) ⃦⃦⃦ KR ≤ 𝜀 for any ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 , then ⃦⃦⃦ ^
𝑉 𝜋 𝑧𝑅 − 𝑉 * ⃦⃦⃦ ∞ ≤ 3 𝜀 max( ‖ 𝑟 ‖ ∞ , ‖ 𝑟 ‖ Lip ) / (1 − 𝛾 ) .
More generally, let ‖·‖ 𝐴 be a norm on functions and ‖·‖ 𝐵 a norm on measures, such that ∫︀ 𝑓 d 𝜇 ≤ ‖ 𝑓 ‖ 𝐴
‖ 𝜇 ‖ 𝐵 for any function 𝑓 and measure 𝜇 .
Then for any reward function 𝑟 such that ‖ 𝑟 ‖ 𝐴 < ∞ ,
[Equation]
Moreover, the optimal 𝑄 -function is close to 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 )
⊤ 𝑧 𝑅 :
[Equation]
E.3.
An approximate 𝑧 𝑅 yields an approximately optimal policy
We now turn to the second source of approximation: computing 𝑧 𝑅 in the reward estimation phase.
This is a problem only if the reward is not specified explicitly.
We deal in turn with the effect of using a model of the reward function, and the effect of estimating 𝑧 𝑅 = E ( 𝑠,𝑎 ) ∼ 𝜌 𝑟 ( 𝑠, 𝑎 ) 𝐵
( 𝑠, 𝑎 ) via sampling.
Which of these options is better depends a lot on the situation.
If training of 𝐹 and 𝐵 is perfect, by construction the policies are optimal for each 𝑧 𝑅 : thus, estimating 𝑧 𝑅 from rewards sampled at 𝑁 states ( 𝑠 𝑖 , 𝑎 𝑖 ) ∼ 𝜌 will produce the optimal policy for exactly that empirical reward, namely, a nonzero reward at each ( 𝑠 𝑖 , 𝑎 𝑖 ) but zero everywhere else, thus overfitting the reward function.
Reducing the dimension 𝑑 reduces this effect, since rewards are projected on the span of the features in 𝐵 : 𝐵 plays both the roles of a transition model and a reward regularizer.
This appears as a √︀ 𝑑/𝑁 factor in Theorem 8 below.
Thus, if both the number of samples to train 𝐹 and 𝐵 and the number of reward samples are small, using a smaller 𝑑 will regularize both the model of the environment and the model of the reward.
However, if the number of samples to train 𝐹 and 𝐵 is large, yielding an excellent model of the environment, but the number of reward samples is small, then learning a model of the reward function will be a better option than direct empirical estimation of 𝑧 𝑅 .
The first result below states that reward misidentification comes on top of the approximation error of the 𝐹 ⊤ 𝐵 model.
This is relevant, for instance, if a reward model ^ 𝑟 is estimated by an external model using some reward values.
Proposition 6 (Influence of estimating 𝑧 𝑅 by an approximate reward) .
Let ^ 𝑟 : 𝑆 × 𝐴 → R be any reward function.
Let ^ 𝑧 𝑅 =
E ( 𝑠,𝑎 ) ∼ 𝜌
[^ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )] .
Let 𝜀 𝐹𝐵 be the error attained by the 𝐹 ⊤ 𝐵 model in Theorem 5 for reward ^ 𝑟 ; namely, assume that ⃦⃦⃦ ^
𝑉 𝜋 ^ 𝑧𝑅 − ^ 𝑉 * ⃦⃦⃦ ∞ ≤ 𝜀 𝐹𝐵 with ^ 𝑉 * the optimal value function for ^ 𝑟 .
Then the policy 𝜋 ^ 𝑧 𝑅 ( 𝑠 ) =
arg max 𝑎 𝐹 ( 𝑠, 𝑎, ^ 𝑧 𝑅 ) ⊤ ^ 𝑧 𝑅 defined by the model ^ 𝑟 is (︂ ‖ 𝑟 − ^ 𝑟 ‖ ∞ (1 − 𝛾 ) + 𝜀 𝐹𝐵 )︂ -optimal for reward 𝑟 .
This still assumes that the expectation ^ 𝑧 𝑅 = E ( 𝑠,𝑎 ) ∼ 𝜌
[^ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )] is computed exactly for the model ^ 𝑟 .
If ^ 𝑟 is given by an explicit model, this expectation can in principle be computed on the whole replay buffer used to train 𝐹 and 𝐵 , so variance would be low.
Nevertheless, we provide an additional statement which covers the influence of the variance of the estimator of 𝑧 𝑅 , whether this estimator uses an external model ^ 𝑟 or a direct empirical average reward observations 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 ).
Definition 7.
The skewness 𝜁 ( 𝐵 ) of 𝐵 is defined as follows.
Assume 𝐵 is bounded.
Let 𝐵 1 , . . .
, 𝐵 𝑑 : 𝑆 × 𝐴 → R be the functions of ( 𝑠, 𝑎 ) defined by each component of 𝐵 .
Let ⟨ 𝐵 ⟩ be the linear span of the ( 𝐵 𝑖 ) 1 ≤ 𝑖 ≤ 𝑑 as functions on 𝑆 × 𝐴 .
Set
[Equation]
Theorem 8 (Influence of estimating 𝑧 𝑅 by empirical averages) .
Assume that 𝑧 𝑅 = E ( 𝑠,𝑎 ) ∼ 𝜌 [ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )] is estimated via
[Equation]
34) using 𝑁 independent samples ( 𝑠 𝑖 , 𝑎 𝑖 ) ∼ 𝜌 , where the 𝑟 𝑖 are random variables such that E
[^ 𝑟 𝑖 | 𝑠 𝑖 , 𝑎 𝑖 ]
= 𝑟 ( 𝑠 𝑖 , 𝑎 𝑖 ) , Var[^ 𝑟 𝑖 | 𝑠 𝑖 , 𝑎 𝑖 ] ≤ 𝑣 for some 𝑣 ∈ R , and the ^ 𝑟 𝑖 are mutually independent given ( 𝑠 𝑖 , 𝑎 𝑖 ) 𝑖
=1 ,...,𝑁 .
Let 𝑉 * be the optimal value function for reward 𝑟 , and let ^ 𝑉 be the value function of the estimated policy 𝜋 ^ 𝑧 𝑅 for reward 𝑟 .
Then, for any 𝛿 > 0 , with probability at least 1 − 𝛿 ,
[Equation]
which is therefore the bound on the optimality gap of 𝜋 ^ 𝑧 𝑅 for 𝑟 .
Here 𝜀 𝐹𝐵 is the error due to the 𝐹 ⊤ 𝐵 model approximation, defined as in Proposition 6.
The proofs are not direct, because 𝐹 is not continuous with respect to 𝑧 .
Contrary to 𝑄 -values, successor states are not continuous in the reward: if an action has reward 1 and the reward for another action changes from 1 − 𝜀 to 1 + 𝜀 , the return values change by at most 2 𝜀 , but the actions and states visited by the optimal policy change a lot.
So it is not possible to reason by continuity on each of the terms involved.
E. A Note on the Measure 𝑀 𝜋 and its Density 𝑚 𝜋
In finite spaces, the definition of the successor state density 𝑚 𝜋 via
[Equation]
with respect to the data distribution 𝜌 poses no problem, as long as the data distribution is positive everywhere.
In continuous spaces, this can be understood as the (Radon–Nikodym) density of 𝑀 𝜋 with respect to 𝜌 , assuming 𝑀 𝜋 has no singular part with respect to 𝜌 .
However, this is never the case: in the definition (7) of the successor state measure 𝑀 𝜋 , the term 𝑡 = 0 produces a
Dirac measure 𝛿 𝑠,𝑎 .
So 𝑀 𝜋 has a singular component due to 𝑡 = 0, and 𝑚 𝜋 is better thought of as a distribution.
When 𝑚 𝜋 is a distribution, a continuous parametric model 𝑚 𝜋𝜃 learned by (9) can approximate 𝑚 𝜋 in the weak topology only: 𝑚 𝜋𝜃 𝜌 approximates 𝑀 𝜋 for the weak convergence of measures.
Thus, for the forward-backward representation, 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ )
𝜌 (d 𝑠 ′ , d 𝑎 ′ ) weakly approximates 𝑀 𝜋
𝑧
( 𝑠, 𝑎, d 𝑠 ′ , d 𝑎 ′ ).
We have not found this to be a problem either in theory or practice, and Theorem 5 covers weak approximations.
Alternatively, one may just define successor states starting at 𝑡 = 1 in (7).
This only works if rewards 𝑟 ( 𝑠, 𝑎 ) depend on the state 𝑠
but not the action 𝑎 (e.g., in goal-oriented settings).
If starting the definition at 𝑡 = 1, 𝑚 𝜋 is an ordinary function provided the transition kernels 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) of the environment are non-singular, 𝜌 has positive density, and 𝜋 (d 𝑎 | 𝑠 ) is non-singular as well.
But starting at 𝑡 = 1 induces some changes in the theorems:
∙ In the learning algorithm (9) for successor states, the term 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠, 𝑎, 𝑠, 𝑎 ) becomes 𝛾 𝜕 𝜃 𝑚 𝜋𝜃 ( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ ).
∙ The expression for the 𝑄 -function in Theorem 1 becomes 𝑄 * ( 𝑠, 𝑎 ) =
𝑟 ( 𝑠, 𝑎 ) + 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 , and likewise in Theorem 2.
The 𝑟 ( 𝑠, 𝑎 ) term covers the immediate reward at a state, since we have excluded 𝑡 = 0 from the definition of successor states.
∙ In general the expression for optimal policies becomes
[Equation]
) which cannot be computed from 𝑧 and 𝐹 alone in the unsupervised training phase.
The algorithm only makes sense for rewards that depend on 𝑠
but not on 𝑎 (e.g., in goal-oriented settings): then the policy 𝜋 𝑧 is equal to 𝜋 𝑧 ( 𝑠 ) :
= arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 again.
E. 𝐹 and 𝐵 as Successor and Predecessor Features of Each Other
We now give a statement that specifies how 𝐹 encodes the future of a state while 𝐵 encodes the past of a state.
Namely: if 𝐹 and 𝐵 minimize their unsupervised loss, then 𝐹 is equal to the successor features from the dual features of 𝐵 , and 𝐵 is equal to the predecessor features from the dual features of 𝐹 (Theorem 9).
This statement holds for a fixed 𝑧 and the corresponding policy 𝜋 𝑧 .
So, for the rest of this section, 𝑧 is fixed.
By “dual” features we mean the following.
Define the 𝑑 × 𝑑 covariance matrices
[Equation]
Then (Cov 𝐹 ) − 1 / 2 𝐹 ( 𝑠, 𝑎, 𝑧 ) is 𝐿 2 ( 𝜌 )-orthonormal and likewise for 𝐵 .
The “dual” features are (Cov 𝐹 ) − 1 𝐹 ( 𝑠, 𝑎, 𝑧 ) and (Cov 𝐵 ) − 1 𝐵 ( 𝑠, 𝑎 ), without the square root: these are the least square solvers for 𝐹 and 𝐵 respectively, and these are the ones that appear below.
The unsupervised FB loss (14) for a fixed 𝑧 is
[Equation]
Thus, minimizers in dimension 𝑑 correspond to an SVD of the successor state density in 𝐿 2 ( 𝜌 ), truncated to the largest 𝑑 singular values.
Theorem 9.
Consider a smooth parametric model for 𝐹 and 𝐵 , and assume this model is overparameterized.
4 Also assume that the data distribution 𝜌 has positive density everywhere.
Let 𝑧 ∈ 𝑍 .
Assume that for this 𝑧 , 𝐹 and 𝐵 lie in 𝐿 2 ( 𝜌 ) and achieve a local extremum of ℓ ( 𝐹, 𝐵 ) within this parametric model.
Namely, the derivative 𝜕 𝜃 ℓ ( 𝐹, 𝐵 ) of the loss with respect to the parameters 𝜃 of 𝐹 is 0 , and likewise for 𝐵 .
Then 𝐹 is equal to (Cov 𝐵 ) − 1 times the successor features of 𝐵 : for any ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 ,
[Equation]
and 𝐵 is equal to (Cov 𝐹 ) − 1 times the predecessor features of 𝐹 :
[Equation]
𝜌 -almost everywhere.
Here the covariances have been defined in (38) , and 𝑃 𝑡 ( ·| 𝑠, 𝑎, 𝜋 ) denotes the law of ( 𝑠 𝑡 , 𝑎 𝑡 ) under trajectories starting at ( 𝑠, 𝑎 ) and following policy 𝜋 .
The same result holds when working with features 𝜙 ( 𝑠 ′ , 𝑎 ′ ) , just by applying it to 𝐵 ∘ 𝜙 .
Note that in the FB framework, we may normalize either 𝐹 or 𝐵 (Remark 4), but not both.
F Proofs
The first proposition is a direct consequence of the definition (15) of successor states with a goal space 𝐺 .
Proposition 10.
Let 𝜙 : 𝑆 × 𝐴 → 𝐺 be a map to some goal space 𝐺 .
Let 𝜋 be some policy, and let 𝑀 𝜋 be the successor state measure (15) of 𝜋 in goal space 𝐺 .
Let 𝑚 𝜋 be the density of 𝑀 𝜋 with respect to some positive probability measure 𝜌 on 𝐺 .
Let 𝑟 : 𝐺 → R be some function on 𝐺 , and define the reward function 𝑅 ( 𝑠, 𝑎 ) :
= 𝑟 ( 𝜙 ( 𝑠, 𝑎 )) on 𝑆 × 𝐴 .
Then the 𝑄 -function 𝑄 𝜋 of policy 𝜋 for reward 𝑅 is
[Equation]
Proof of Proposition 10.
For each time 𝑡 ≥ 0, let 𝑃 𝜋𝑡 ( 𝑠 0 , 𝑎 0 , d 𝑔 ) be the probability distribution of 𝑔 = 𝜙 ( 𝑠 𝑡 , 𝑎 𝑡 ) over trajectories of the policy 𝜋 starting at ( 𝑠 0 , 𝑎 0 ) in the MDP.
Thus, by the definition (15),
[Equation]
[Equation]
The 𝑄 -function of 𝜋 for the reward 𝑅 is by definition (the sums and integrals are finite since 𝑅 is bounded)
[Equation]
[Equation]
[Equation]
by definition of the density 𝑚 𝜋 .
Proof of Theorems 1 and 2.
Theorem 1 is a particular case of Theorem 2 ( 𝜙 = Id and ¯ 𝑚 = 0), so we only prove the latter.
Let 𝑅 ( 𝑠, 𝑎 ) =
𝑟 ( 𝜙 ( 𝑠, 𝑎 )) be a reward function as in the theorem.
The 𝑄 -function of 𝜋 for the reward 𝑅 is, by Proposition 10,
[Equation]
The assumptions state that for any 𝑧 ∈ 𝑍 , 𝑚 𝜋 𝑧
( 𝑠, 𝑎, 𝑔 ) is equal to 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝐵 ( 𝑔 )+ ¯ 𝑚
( 𝑠, 𝑧, 𝑔 ).
Therefore, for any 𝑧 ∈ 𝑍 we have
[Equation]
by definition of 𝑧 𝑅 and ¯ 𝑉 .
This proves the claim (20) about 𝑄 -functions.
By definition, the policy 𝜋 𝑧 selects the action 𝑎 that maximizes 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 .
Take 𝑧 = 𝑧 𝑅 .
Then
[Equation]
since the last term does not depend on 𝑎 .
This quantity is equal to 𝑄 𝜋 𝑧𝑅 ( 𝑠, 𝑎 ).
Therefore,
[Equation]
and by the above, 𝑄 𝜋 𝑧𝑅 ( 𝑠, 𝑎 ) is indeed equal to the 𝑄 -function of policy 𝜋 𝑧 𝑅 for the reward 𝑅 .
Therefore, 𝜋 𝑧 𝑅 and 𝑄 𝜋 𝑧𝑅 constitute an optimal Bellman pair for reward 𝑅 .
Since 𝑄 𝜋 𝑧𝑅 ( 𝑠, 𝑎 ) is the 𝑄 -function of 𝜋 𝑍 𝑅 , it satisfies the Bellman equation
[Equation]
by (57).
This is the optimal Bellman equation for 𝑅 , and 𝜋 𝑧 𝑅 is the optimal policy for 𝑅 .
We still have to prove the last statement of Theorem 2.
Since 𝜋 𝑧 𝑅 is an optimal policy for 𝑅 , for any other policy 𝜋 𝑧 and state-action ( 𝑠, 𝑎 ) we have
[Equation]
Using the formulas above for 𝑄 𝜋 , with ¯ 𝑚 = 0, this rewrites as
[Equation] [Equation]
as needed.
Thus 𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 ≥ sup
𝑧 ∈ 𝑍 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 𝑅 , and equality occurs by taking 𝑧 = 𝑧 𝑅 .
This ends the proof of Theorem 2.
Proof of Proposition 3.
Assume 𝑑
=
# 𝑆 × # 𝐴 ; extra dimensions can just be ignored by setting the extra components of 𝐹 and 𝐵 to 0.
With 𝑑 = # 𝑆 × # 𝐴 , we can index the components of 𝑍 by pairs ( 𝑠, 𝑎 ).
First, let us set 𝐵 ( 𝑠, 𝑎 ) :
= 1 𝑠,𝑎 .
Let 𝑟 : 𝑆 × 𝐴 → R be any reward function.
Let 𝑧 𝑅 ∈ R
# 𝑆 × # 𝐴 be defined as in Theorem 1, namely, 𝑧 𝑅 = E ( 𝑠,𝑎 ) ∼ 𝜌 [ 𝑟 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )].
With our choice of 𝐵 , the components of 𝑧 𝑅 are ( 𝑧 𝑅 ) 𝑠,𝑎
= 𝑟 ( 𝑠, 𝑎 ) 𝜌 ( 𝑠, 𝑎 ).
Since 𝜌 > 0, the correspondence 𝑟 ↔ 𝑧 𝑅 is bijective.
Let us now define 𝐹 .
Take 𝑧 ∈ 𝑍 .
Since 𝑟 ↔ 𝑧 𝑅 is bijective, this 𝑧 is equal to 𝑧 𝑅 for some reward function 𝑟 .
Let 𝜋 𝑧 be an optimal policy for this reward 𝑟 in the MDP.
Let 𝑀 𝜋 𝑟 be the successor state measure of policy
𝜋 𝑧 , namely:
[Equation]
Now define 𝐹 ( 𝑠, 𝑎, 𝑧 ) by setting its ( 𝑠 ′ , 𝑎 ′ ) component to 𝑀 𝜋 𝑧 ( 𝑠, 𝑎, 𝑠 ′ , 𝑎 ′ ) /𝜌
( 𝑠 ′ , 𝑎 ′ ) for each ( 𝑠 ′ , 𝑎 ′ ):
[Equation]
Then we have
[Equation]
because by our choice of 𝐵 , 𝐵 ( 𝑠 ′ , 𝑎 ′ )
𝑠 ′′ 𝑎 ′′ = 1 𝑠 ′ =
𝑠
[Equation]
,𝑎 ′ =
𝑎 [Equation]
Thus, 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ ) is the density of the successor state measure 𝑀 𝜋 𝑧 of policy 𝜋 𝑧 with respect to 𝜌 , as needed.
We still have to check that 𝜋 𝑧 satisfies 𝜋 𝑧 ( 𝑠 ) =
arg max 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 (since this is not how it was defined).
Since 𝜋 𝑧 was defined as an optimal policy for the reward 𝑟 associated with 𝑧 , it satisfies
[Equation]
with 𝑄 𝜋 𝑧
( 𝑠, 𝑎 ) the 𝑄 -function of policy 𝜋 𝑧 for the reward 𝑟 .
This 𝑄 -function is equal to the cumulated expected reward
[Equation]
since 𝑧 is equal to E ( 𝑠 ′ ,𝑎 ′ )
∼ 𝜌 [ 𝑟 ( 𝑠 ′ , 𝑎 ′ )
𝐵 ( 𝑠 ′ , 𝑎 ′ )].
This proves that 𝜋 𝑧 ( 𝑠 ) =
arg max 𝑎 𝑄 𝜋 𝑧 ( 𝑠,𝑎 ) =
arg max 𝑎 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 .
So this choice of 𝐹 and 𝐵 satisfies all the properties claimed.
We will rely on the following two basic results in 𝑄 -learning.
Proposition 11 ( 𝑟 ↦→ 𝑄 * is Lipschitz in sup-norm) .
Let 𝑟 1 , 𝑟 2 :
𝑆 × 𝐴
→ R be two bounded reward functions.
Let 𝑄 * 1 and 𝑄 * 2 be the corresponding optimal 𝑄 -functions, and likewise for the 𝑉 -functions.
Then
[Equation] [Equation]
and
[Equation]
Proof.
Assume sup 𝑆 × 𝐴 | 𝑟 1 − 𝑟 2 | ≤ 𝜀 for some 𝜀 ≥ 0.
For any policy 𝜋 , let 𝑄 𝜋 1 be its 𝑄 -function for reward 𝑟 1 , and likewise for 𝑟 2 .
Let 𝜋 1 and 𝜋be optimal policies for 𝑟 1 and 𝑟 2 , respectively.
Then for any ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 ,
[Equation] [Equation]
(76) =
[Equation]
𝑡
[Equation] [Equation]
( 𝑠
[Equation]
𝑎 )] (77) ≥ ∑︁ 𝑡 [Equation] [Equation] 𝑡
[Equation]
𝑡
[Equation]
|
[Equation]
𝑎
[Equation]
( 𝑠
[Equation]
𝑎 )
[Equation]
(78) =
∑︁ 𝑡
[Equation]
𝛾 𝑡 E [ 𝑟 [Equation] 𝑠 𝑡
[Equation]
𝑎 𝑡
[Equation]
| 𝜋
[Equation]
( 𝑠
[Equation]
𝑎 0 ) = ( 𝑠
[Equation]
𝑎 )] − 𝜀 1 − 𝛾 (79) = 𝑄 * 2 ( 𝑠, 𝑎 ) − 𝜀 1 − 𝛾 (80) and likewise in the other direction, which ends the proof for 𝑄 -functions.
The case of 𝑉 -functions follows by restricting to the optimal actions at each state 𝑠 .
Proposition 12.
Let 𝑓 : 𝑆 × 𝐴 → R be any function, and define a policy 𝜋 𝑓 by 𝜋 𝑓 ( 𝑠 ) :
= arg max 𝑎 𝑓 ( 𝑠, 𝑎 ) .
Let 𝑟 : 𝑆 × 𝐴 → R be some bounded reward function.
Let 𝑄 * be its optimal 𝑄 -function, and let 𝑄 𝜋
𝑓 be the 𝑄 -function of 𝜋 𝑓 for reward 𝑟 .
Then
[Equation]
and
[Equation]
Proof.
Define 𝜀
( 𝑠, 𝑎 ) :
= 𝑄 𝜋 𝑓
( 𝑠, 𝑎 ) − 𝑓
( 𝑠, 𝑎 ).
The 𝑄 -function 𝑄 𝜋 𝑓 satisfies the Bellman equation
[Equation]
for any ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 .
Substituting 𝑄 𝜋 𝑓 = 𝑓 + 𝜀 , this rewrites as
[Equation]
by definition of 𝜋 𝑓 , where we have set
[Equation]
(86) is the optimal Bellman equation for the reward 𝑟 − 𝜀 ′ .
Therefore, 𝑓 is the optimal 𝑄 -function for the reward 𝑟 − 𝜀 ′ .
Since 𝑄 * is the optimal 𝑄 -function for reward 𝑟 , by
Proposition 11, we have
[Equation]
By construction of 𝜀 ′ , sup 𝑆 × 𝐴 | 𝜀 ′ | ≤ 2 sup 𝑆 × 𝐴 | 𝜀 | = 2 sup 𝑆 × 𝐴 | 𝑓 − 𝑄 𝜋 𝑓 | .
This proves the first claim.
The second claim follows by the triangle inequality | 𝑄 𝜋 𝑓 − 𝑄 * | ≤ | 𝑄 𝜋 𝑓 − 𝑓 | + | 𝑓 − 𝑄 * | and 21 − 𝛾 + 1 ≤ 31 − 𝛾 .
Proof of Theorem 5.
By construction of the Kantorovich–Rubinstein norm, the second claim of Theorem 5 is a particular case of the third claim, with ‖ 𝑓 ‖ 𝐴 :=
max( ‖ 𝑓 ‖ ∞ , ‖ 𝑓 ‖ Lip ) and ‖ 𝜇 ‖ 𝐵 :=
‖ 𝜇 ‖ KR .
Likewise, since 𝑚 is the density of 𝑀 with respect to 𝜌 , the first claim is an instance of the third, by taking ‖ 𝑓 ‖ 𝐴 :=
‖ 𝑓 ‖ ∞ and ‖ 𝜇 ‖ 𝐵 :=
⃦⃦⃦ d 𝜇 d 𝜌 ⃦⃦⃦ 𝐿 1 ( 𝜌 ) .
Therefore, we only prove the third claim.
Let 𝑧 ∈ 𝑍 and let 𝑟 : 𝑆 × 𝐴 → R be any reward function.
By Proposition 10 with 𝐺 = 𝑆 × 𝐴 and 𝜙 = Id, the 𝑄 -function of policy 𝑄 𝜋 𝑧 for this reward is
[Equation]
Let 𝜀 𝑧 ( 𝑠, 𝑎, d 𝑠 ′ , d 𝑎 ′ ) be the difference of measures between the model 𝐹 ⊤ 𝐵𝜌 and 𝑀 𝜋 𝑧 : 𝜀 𝑧 ( 𝑠, 𝑎, d 𝑠 ′ , d 𝑎 ′ ) :
=
𝑀 𝜋 𝑧 ( 𝑠, 𝑎, d 𝑠 ′ , d 𝑎 ′ ) − ^ 𝑀 𝑧 ( 𝑠, 𝑎, d 𝑠 ′ , d 𝑎 ′ ) (90) =
𝑀 𝜋 𝑧
[Equation]
We want to control the optimality gap in terms of sup 𝑠,𝑎 ‖ 𝜀 𝑧 ( 𝑠, 𝑎, · ) ‖ 𝐵 .
By definition of 𝜀 𝑧 ,
[Equation]
since 𝑧 𝑅 = ∫︀ 𝑟 ( 𝑠 ′ , 𝑎 ′ )
𝐵 ( 𝑠 ′ , 𝑎 ′ )
𝜌 (d 𝑠 ′ , d 𝑎 ′ ).
Therefore,
[Equation]
for any reward 𝑟 and any 𝑧 ∈ 𝑍 (not necessarily 𝑧 = 𝑧 𝑅 ).
Let 𝑄 * be the optimal 𝑄 -function for reward 𝑟 .
Define 𝑓
( 𝑠, 𝑎 ) :
=
𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 .
By definition, the policy 𝜋 𝑧 𝑅 is equal to arg max 𝑎 𝑓 ( 𝑠, 𝑎 ).
Therefore, by Proposition 12,
[Equation]
[Equation]
and
[Equation]
But by the above,
[Equation]
Therefore, for any reward function 𝑟 ,
[Equation]
This inequality transfers to the value functions, hence the result.
In addition, using again 𝑓 ( 𝑠, 𝑎 ) =
𝐹 ( 𝑠, 𝑎, 𝑧 𝑅 ) ⊤ 𝑧 𝑅 , we obtain
[Equation]
Proof of Proposition 6.
This is just a triangle inequality.
By assumption, the difference between the value function of 𝜋 ^ 𝑧 𝑅 and the optimal value function 𝑉 * ^ 𝑟 is at most 𝜀 𝐹𝐵 .
Then by Proposition 11, the difference between 𝑉 * ^ 𝑟 and 𝑉 * 𝑟 is bounded by 11 − 𝛾 sup 𝑆 × 𝐴 | ^ 𝑟 − 𝑟 | .
Proof of Theorem 8.
We proceed by building a reward function ^ 𝑟 corresponding to ^ 𝑧 𝑅 .
Then we will bound ^ 𝑟 − 𝑟 and apply Proposition 6.
First, by Remark 4, up to reducing 𝑑 , we can assume that 𝐵 is 𝐿 2 ( 𝜌 )-orthonormal.
For any function 𝜙 : ( 𝑠, 𝑎 ) → R , define 𝑧 𝜙 := E ( 𝑠,𝑎 ) ∼ 𝜌 [ 𝜙 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 )].
For each 𝑧 ∈ 𝑍 , define 𝜙 𝑧 via 𝜙 𝑧 ( 𝑠, 𝑎 ) :
= 𝐵 ( 𝑠, 𝑎 ) ⊤ 𝑧 .
Then, if 𝐵 is 𝐿 2 ( 𝜌 )-orthonormal, we have 𝑧 𝜙 𝑧 = 𝑧 .
(Indeed, 𝑧 𝜙 𝑧 = E [( 𝐵 ( 𝑠, 𝑎 ) ⊤ 𝑧 ) 𝐵
( 𝑠, 𝑎 )]
=
E [ 𝐵 ( 𝑠, 𝑎 )( 𝐵 ( 𝑠, 𝑎 ) ⊤ 𝑧 )]
= (︀ E
[ 𝐵 ( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 ) ⊤ ] )︀ 𝑧 .)
Define the function
[Equation]
[Equation]
using the functions 𝜙 𝑧 defined above.
By construction, 𝑧 ^ 𝑟 = 𝑧 𝑅 + 𝑧 𝜙 ^ 𝑧𝑅 − 𝑧𝑅 =
^ 𝑧 𝑅 .
Therefore, the policy 𝜋 ^ 𝑧 𝑅 associated to ^ 𝑧 𝑅 is the policy associated to the reward ^ 𝑟 .
We will now apply Proposition 6 to 𝑟 and ^ 𝑟 .
For this, we need to bound ‖ 𝜙 ^ 𝑧 𝑅 − 𝑧 𝑅 ‖ ∞ .
Let 𝐵 1 , 𝐵 2 , . . .
, 𝐵 𝑑 be the components of 𝐵 as functions on 𝑆 × 𝐴 .
For any 𝑧 ∈ 𝑍 , we have
[Equation]
since the 𝐵 𝑖 are 𝐿 2 ( 𝜌 )-orthonormal.
Moreover, by construction, 𝜙 𝑧 lies in the linear span ⟨ 𝐵 ⟩ of the functions ( 𝐵 𝑖 ).
Therefore
[Equation]
by the definition of 𝜁 ( 𝐵 ) (Definition 7).
Therefore,
[Equation]
Let us now bound ^ 𝑧 𝑅 − 𝑧 𝑅 :
[Equation]
The first term satisfies
[Equation]
because the ^ 𝑟 𝑖 are independent conditionally to ( 𝑠 𝑖 , 𝑎 𝑖 ), and because 𝐵 is deterministic.
The expectation of this over ( 𝑠 𝑖 , 𝑎 𝑖 ) is
[Equation]
which is thus a bound on the first term.
The second term satisfies
[Equation]
since the ( 𝑠 𝑖 , 𝑎 𝑖 ) are independent with distribution 𝜌 .
By the Cauchy–Schwarz inequality (applied to each component of 𝐵 ), this is at most
[Equation]
Therefore,
[Equation]
Since 𝐵 is orthonormal in 𝐿 2 ( 𝜌 ), we have ‖ 𝐵 ‖ 2 𝐿 2 ( 𝜌 ) =
𝑑 .
Putting everything together, we find
[Equation]
Therefore, by the Markov inequality, for any 𝛿 > 0, with probability 1 − 𝛿 ,
[Equation]
hence the conclusion by Proposition
[Equation]
Proof of Theorem 9.
Let
[Equation]
so that
[Equation]
Let us first take the derivative with respect to the parameters of 𝐹 .
This is 0 by assumption, so we find
[Equation]
where
[Equation]
Since the model is overparameterized [Equation] we can realize any 𝐿 2 function 𝑓 ( 𝑠, 𝑎 ) as the derivative 𝜕 𝜃 𝐹 ( 𝑠, 𝑎, 𝑧 ) for some direction 𝜃 .
Therefore, the equation 0 = ∫︀ 𝜕 𝜃 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝐺 ( 𝑠, 𝑎 ) 𝜌 (d 𝑠, d 𝑎 ) implies that 𝐺 ( 𝑠, 𝑎 ) is 𝐿 2 ( 𝜌 )-orthogonal to any function 𝑓
( 𝑠, 𝑎 ) in 𝐿 2 ( 𝜌 ).
Therefore, 𝐺 ( 𝑠, 𝑎 ) vanishes 𝜌 -almost everywhere, namely
[Equation]
Now, since 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ ) is a real number, 𝐹 ( 𝑠, 𝑎, 𝑧 )
⊤ 𝐵 ( 𝑠 ′ , 𝑎 ′ )
=
𝐵 ( 𝑠 ′ , 𝑎 ′ )
⊤ 𝐹 ( 𝑠, 𝑎, 𝑧 ).
Therefore, the right-hand-side above rewrites as
[Equation]
so that
[Equation]
Unfolding the definition of 𝑚 yields the statement for 𝐹 .
The proof for 𝐵 is similar.
G Experimental Setup
In this section we provide additional information about our experiments.
G. Environments
∙ Discrete maze: is the 11 × 11 classical tabular gridworld with foor rooms.
States are represented by one-hot unit vectors, 𝑆 = { 0 , 1 } 121 .
There are five available actions , 𝐴 = { left , right , up , down , do nothing } .
The dynamics are deterministic and the walls are impassable.
∙ Continuous maze: is a two dimensional environment with impassable walls.
States are represented by their Cartesian coordinates ( 𝑥, 𝑦 ) ∈ 𝑆 = [0 , 1] 2 .
There are five available actions, 𝐴 = { left , right , up , down , do nothing } .
The execution of one of the actions moves the agent 0 .
1 units in the desired direction, and normal random noise with zero mean and standard deviation 0 .
01 is added to the position of the agent (that is, a move along the x axis would be 𝑥 ′ =
𝑥 ± 0 .
1 + 𝒩 (0 , 0 . 01), where 𝒩 (0 , 0 . 01) is a normal variable with mean 0 and standard deviation 0 . 01).
If after a move the agent ends up outside of [0 , 1] 2 , the agent’s position is clipped (e.g if 𝑥 < 0 then we set 𝑥 = 0).
If a move make the agent cross an interior wall, this move is undone.
For all algorithms, we convert a state 𝑠
= ( 𝑥, 𝑦 ) into feature vector 𝜙 ( 𝑠 ) ∈
R 441 by computing the activations of a regular 21 × 21 grid of radial basis functions at the point ( 𝑠, 𝑦 ).
Especially, we use Gaussian functions: 𝜙 ( 𝑠 ) =
(︁ exp( − ( 𝑥 − 𝑥 𝑖 ) 2 +
[Equation]
𝑦
[Equation]
𝑦 𝑖 ) 2 𝜎 ) , . . . ,
exp( − ( 𝑥
[Equation]
𝑥
[Equation]
) 2 +( 𝑦 − 𝑦
[Equation]
) 2 2 𝜎) )
︁ where ( 𝑥 𝑖 , 𝑦 𝑖 ) is the center of the 𝑖 𝑡ℎ Gaussian and 𝜎 = 0 .
∙ FeatchReach: is a variant of the simulated robotic arm environment from (Plappert et al., 2018) using discrete actions instead of continuous actions
States are 10dimensional vectors consisting of positions and velocities of robot joints. We
discretise the original 3-dimensional action space into 6 possible actions using action stepsize of 1 (The same way as in https ://github.com/ paulorauber
hpg
, the implementation of hindsight policy gradient (Rauber et al., 2018)).
The goal space is -dimensional
space representing of the position of the object to reach
∙ Ms. Pacman is a variant of the Atari 2600 game Ms. Pacman, where an episode ends when the agent is captured by a monster (Rauber et al., 2018).
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 𝑠 𝑡 is the concatenation of ( 𝑥 𝑡 − 12 , 𝑥 𝑡 − , 𝑥 𝑡 − 𝑥 𝑡 frames,
i
e. an 84 × 84 ×tensor.
An action repeat of 12 is used.
As Ms. Pacman is not originally a multi-goal domain, we define the set of goals as the set of the reachable coordinate pairs 𝑥, 𝑦 ) on the screen; these can be reached only by learning to avoid monsters.
In contrast with (Rauber et al. , who use a heuristic to find the agent’s position from the screen’s pixels, we use the Atari annotated RAM interface wrapper (Anand et al.,
G. Architectures
We use the same architecture for discrete maze, continuous maze and FeatchReach.
Both forward and backward networks are represented by a feedforward neural network with three hidden layers, each with 256 ReLU units.
The forward network receives a concatenation of a state and a 𝑧 vector as input and has | 𝐴 | × 𝑑 as output dimension.
The backward network receives a state as input (or gripper’s position for FeatchReach) and has 𝑑 as output dimension.
For goal-oriented DQN, the 𝑄 -value network is also a feedforward neural network with three hidden layers, each with 256 ReLU units.
It receives a concatenation of a state and a goal as input and has | 𝐴 | as output dimension.
For Ms. Pacman, the forward network is represented by a convolutional neural network given by a convolutional layer with 32 filters (8 × 8, stride 4); convolutional layer with 64 filters (4 × 4, stride 2); convolutional layer with 64 filters (3 × 3, stride 1); and three fully-connected layers, each with 256 units.
We use ReLU as activation function.
The 𝑧 vector is concatenated with the output of the third convolutional layer.
The output dimension of the final linear layer is | 𝐴 | × 𝑑 .
The backward network acts only on agent’s position, a 2-dimensional input.
It is represented by a feedforward neural network with three hidden layers, each with 256 ReLU units.
The output dimension is 𝑑 .
For goal-oriented DQN, the 𝑄 -value network is represented by a convolutional neural network with the same architecture as the one of the forward network.
The goal’s position is concatenated with the output of the third convolutional layer.
The output dimension of the final linear layer is | 𝐴 | .
G. Implementation Details
For all environments, we run the algorithms for 800 epochs.
Each epoch consists of 25 cycles where we interleave between gathering some amount of transitions, to add to the replay buffer 𝒟 (old transitions are thrown when we reach the maximum of its size), 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 ( 𝜀 = 0 . 2) with respect to the current approximation 𝐹 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 for a sampled 𝑧 .
At evaluation time, 𝜀 -greedy policies are also used, with a smaller 𝜀 = 0 .
02 for all envrionments except from discrete maze where we use Boltzmann policy with temperature 𝜏 = 1.
We train each models for three different random seeds.
For generality, we will keep using the notation 𝐵 ( 𝑠, 𝑎 ) while in our experiments 𝐵 acts only on 𝜙 ( 𝑠, 𝑎 ), a part of the state-action space.
For discrete and continuous mazes, 𝜙 ( 𝑠, 𝑎 ) =
𝑠 , for FetchReach, 𝜙 ( 𝑠, 𝑎 ) the position of arm’s gripper and for Ms. Pacman, 𝜙 ( 𝑠, 𝑎 ) is the 2-dimensional position ( 𝑥, 𝑦 ) of the agent on the screen.
We denote by 𝜃 and 𝜔 the parameters of forward and backward networks respectively and 𝜃 − and 𝜔 − the parameters of their corresponding target networks.
Both 𝜃 − and 𝜔 − are updated after each cycle using Polyak averaging; i.e 𝜃 − ← 𝛼𝜃 − + (1 − 𝛼 ) 𝜃 and 𝜔 − ← 𝛼𝜔 − + (1 − 𝛼 ) 𝜔 where 𝛼 = 0 .
95 is the Polyak coefficient.
During training, we sample 𝑧 from a rescaled Gaussian that we denote 𝜈 .
Especially, we sample a 𝑑 -dimensional standard Gaussian variable 𝑥 ∼ 𝒩 (0 , Id) ∈ R 𝑑 and a scalar centered
Cauchy variable 𝑢 ∈ R of scale 0 .
5, then we set 𝑧 = √ 𝑑𝑢 𝑥
‖ 𝑥
We use a Cauchy distribution to ensure that the norm of 𝑧 spans the non-negative real numbers space while having a heavy tail: with a pure Gaussian, the norm of 𝑧 would be very concentrated around a single value.
We also scale by √ 𝑑 to ensure that each component of 𝑧 has an order of magnitude of 1.
Before being fed to 𝐹 , 𝑧 is preprocessed by 𝑧 ← 𝑧 √ 1+ ‖ 𝑧 ‖ 22 /𝑑 ; this way, 𝑧 ranges over a bounded set in R 𝑑 , and this takes advantage of optimal policies being equal for a reward 𝑅 and for 𝜆𝑅 with 𝜆 > 0.
To update network parameters, we compute an empirical loss by sampling 3 mini-batches, each of size 𝑏 = 128, of transitions { ( 𝑠 𝑖 , 𝑎 𝑖 , 𝑠 𝑖 +1 ) }
𝑖 ∈ 𝐼 ⊂ 𝒟 , of target state-action pairs { ( 𝑠 ′ 𝑖 , 𝑎 ′ 𝑖 ) }
𝑖 ∈ 𝐼 ⊂ 𝒟 and of { 𝑧 𝑖 } 𝑖 ∈
𝐼 ∼ 𝜈 :
[Equation]
∑︁
[Equation] ) (125) where we use the Boltzmann policy 𝜋 𝑧 𝑖 ( · | 𝑠 𝑖 +1 )
=
softmax ( 𝐹 𝜃 − ( 𝑠 𝑖 +1 , · , 𝑧 𝑖 ) ⊤ 𝑧 𝑖 /𝜏 ) with fixed temperature 𝜏 = 200 to avoid the instability and discontinuity caused by the argmax operator.
Since there is unidentifiability between 𝐹 and 𝐵 (Appendix, Remark 4), we include a gradient to make 𝐵 closer to orthonormal, E ( 𝑠,𝑎 ) ∼ 𝜌
𝐵
( 𝑠, 𝑎 ) 𝐵 ( 𝑠, 𝑎 ) ⊤
≈ Id:
[Equation]
To compute an unbiased estimate of the latter gradient, we use the following auxiliary empirical loss:
[Equation]
Finally, we use the Adam optimizer and we update 𝜃 and 𝜔 by taking a gradient step on L ( 𝜃, 𝜔 ) and L ( 𝜃, 𝜔 ) + 𝜆 · L reg ( 𝜔 ) respectively, where 𝜆 is a regularization coefficient that we set to 1 for all experiments.
We summarize the hyperparameters used for FB algorithm and goal-oriented DQN in tableand 2 respectively. Hyperparameters
[Figure]
[Table]
FetchReach
Ms. Pacman
[Figure]
evaluation 𝜀
Boltzman with 𝜏 =
[Figure]
[Table] [Table]
G. Experimental results
In this section [Table] we provide additional experimental results [Table]
G.4. Goal-Oriented Setup: Quantitative Comparisons
[Figure] [Figure]
[Figure]
G.4. More Complex Rewards: Qualitative Results G.4. Embedding Visualization
[Figure]
[Figure] [Figure]
[Figure]
[Equation]