N EURAL A LL -P AIRS S HORTEST P ATH FOR R EINFORCEMENT L EARNING
A BSTRACT
Having an informative and dense reward function is an important requirement to efﬁciently solve goal-reaching tasks.
While the natural reward for such tasks is a binary signal indicating success or failure, providing only a binary reward makes learning very challenging given the sparsity of the feedback.
Hence, introducing dense rewards helps to provide smooth gradients.
However, these functions are not readily available, and constructing them is difﬁcult, as it often requires a lot of time and domain-speciﬁc knowledge, and can unintentionally create spurious local minima.
We propose a method that learns neural all-pairs shortest paths , used as a distance function to learn a policy for goal-reaching tasks, requiring zero domainspeciﬁc knowledge.
In particular, our approach includes both a self-supervised signal from the temporal distance between state pairs of an episode, and a metricbased regularizer that leverages the triangle inequality for an additional connectivity information between state triples.
This dynamical distance can be either used as a cost function, or reshaped as a reward, and, differently from previous work, is fully self-supervised, compatible with off-policy learning and robust to local minima.
I NTRODUCTION
Teaching an AI agent to autonomously and efﬁciently reach goals is still an open problem in the reinforcement learning (RL) community.
One of the bottlenecks is correctly specifying a learning signal: an intuitive choice would be a binary reward indicating whether the goal is reached or not, but this kind of signal does not provide enough feedback to the learner, as it would require extensive exploration before receiving any information.
As an alternative, hand-crafted reward shaping is often employed on top of the “true" sparse reward - such as the Euclidean distance to the goal (Ng et al., 1999a; Mahmood et al., 2018) - but that might lead to reward hacking, a widely observed phenomenon in which the agent optimizes for the local optima introduced by the shaped reward instead of the real sparse objective, generating unintended behaviors (Pan et al., 2022).
Instead of imposing brittle heuristic shaping, recent research shows that learning dynamical distances is a valid alternative (Hartikainen et al., 2019).
The purpose of these functions is to learn all-pairs shortest path distances (APSP).
One of the ﬁrst works in this direction was presented in the 90s by
Kaelbling (1993) and utilizes a Q-learning agent to construct a goal-conditioned action-value function d ( s, a, g ) , drawing de facto the ﬁrst connection between RL and the APSP problem, based on the
Floyd-Warshall algorithm (Floyd, 1962a).
However, Kaelbling (1993) demonstrate the method in a low dimensional tabular setting, where all the goals were known in advance.
Recent works (Hartikainen et al., 2019; Jurgenson et al., 2020) propose extensions including function approximation, but are constrained to on-policy learning.
We propose a method that learns a dynamical distance function with function approximation, off-policy , without any reinforcement signal, employing only a self-supervised signal from the actual number of time steps separating state pairs, and a second loss inspired by the triangle inequality, as in the all-pairs shortest path algorithm.
1. S HORTEST P ATHS AND RL
The problem of learning dynamical distances is closely related to the one of learning value functions in RL.
The Q-learning algorithm itself can be seen as a generalization of a single-source shortest path problem (SSSP) for directed graphs, in the context of decision making.
The negative Q-function  Q ⇡ ( s, ⇡ ( s )) represents the shortest path distance under the policy ⇡ .
The dynamic programming (DP) equation underlying Q-learning is the Bellman update (Bellman, 1966), while the one for single-source shortest path is provided by the Dijkstra’s algorithm (Dijkstra, 1959) or by other variations like the Bellman-Ford algorithm (Bellman, 1958), depending on the input graph class.
However, the nature of these recursive equations can still be traced to the Bellman’s optimality principle (Sniedovich, 2006).
Goal-conditioned Q-functions (Schaul et al., 2015) generalize distance learning for a multi-goal setting, as they are deﬁned over pairs of states and desired goals.
This is similar to the all-pairs shortest path problem (APSP), which computes the shortest path between any two vertices in a graph.
However, the DP equations behind these problems are different.
Goal-conditioned Q-learning is typically based on the temporal difference update, where the only bootstrapped value is the estimate of the Q-function.
Dedicated algorithms for the solution of the APSP problem are built on different
DP principles.
The Floyd-Warshall algorithm (Floyd, 1962b) for APSP, for example, ﬁnds shortest paths by iteratively enforcing the triangle inequality between state triples: the shortest path between state s i and s k is cast in terms of relaxations , such that the distances are initialized by overestimation, and are updated as d ( s i , s k ) :
=
min
j
[ d ( s i , s j ) + d ( s j , s k )] .
O ( N 3 ) relaxations are needed for the algorithm to converge, where N is the number of vertices in the graph.
Kaelbling (1993) propose a RL connection to the Floyd-Warshall algorithm that also includes actions, formalized as d ( s i , a i , s k ) :
=
min
j
[ d ( s i , a i , s j ) + min a 0 d ( s j , a 0 , s k )] .
Differently from a temporal difference update, this equation makes use of available estimates of d ( s i , a i , s j ) as well.
Several other works (Kaelbling, 1993; Dhiman et al., 2018; Jurgenson et al., 2020) adapt the Floyd-Warshall algorithm to
RL, with results mainly in the tabular setting or for on-policy learning.
These methods rely on having accessibility to the ground truth cost between state pairs.
Our approach lifts this impractical assumption and learns distances in a self-supervised fashion.
Moreover, since we are in a reinforcement learning setting, we can let the agent decide which paths to increase depending on their reachability, in a bottom-up approach, rather than initializing all the distances to inﬁnity as in the Floyd-Warshall algorithm.
To summarize, this paper proposes a way to learn dynamical distances and goal-conditioned policies off-policy, and in a fully self-supervised way, employing:
• a novel off-policy correction term for distances learned via temporal regression, and a corrective bootstrapping loss inspired by the Floyd-Warshall algorithm on graphs, adapted to RL.
• a general method to learn dynamical distance functions from scratch, that can be used as part of a goal-conditioned RL algorithm, or as a cost for model-based RL, and that, differently from previous work, successfully deals with local minima (created e
g. by obstacles) without explicitly encoding any type of information in the distance function.
R ELATED W ORK
Dealing with environments with sparse reinforcement signals complicates long-term credit assignment, as it introduces higher variance in Monte Carlo (MC) learning and higher bias in Temporal
Difference (TD) updates (Arjona-Medina et al., 2019).
Many techniques are focused on accelerating learning by “densifying" the reward signal with auxiliary functions while not affecting the optimal policy of the underlying sparse problem (Grzes & Kudenko, 2008; Randløv & Alstrøm, 1998; Ng et al., 1999b).
Alternatively, other kinds of auxiliary reward functions can be used to facilitate exploration according to different strategies as curiosity (Schmidhuber, 1991), prediction error (Pathak et al., 2017), information gain (Houthooft et al., 2016), or predicting state reachability (Savinov et al., 2018).
On the other side there are methods like Hindsight Experience Replay (HER) (Andrychowicz et al., 2017) that do not try to solve the exploration problem, but introduce auxiliary rewards based on the distance to the achieved goal rather than to the desired one.
As done in HER, we also employ goal-conditioned value functions (Schaul et al., 2015).
Our approach, however, is based on dynamical distances, which are functions expressing the distance between state pairs based on some notion of functional similarity .
In the literature, dynamical distances have been learned via direct regression using temporal regression Hartikainen et al.
(2019), in the form of goal-conditioned policies Ghosh et al.
(2019), via Q-learning by relabeling goals
Eysenbach et al.
; Florensa et al.
(2019), or goal-directed Q-learning with negative transition mining Tian et al.
We extend the on-policy temporal regression presented in Hartikainen et al.
(2019) to be suited for off-policy learning, while simultaneously learning a goal-conditioned policy that uses the learned dynamical distance as a negative reward.
Differently from other methods (Eysenbach et al., 2019; Savinov et al., 2018), we do not explicitly perform graph search.
We instead make use of a term that expresses whether two states were ever part of the same trajectory.
At training time, we only check if the ‘edge’ ( s, a, s 0 ) has ever been visited or not.
In practice, we use locality sensitive hashing (Tang et al., 2017), which allows for generalization based on angular distance similarity between states.
B ACKGROUND
We consider the problem of an agent attempting to solve a goal-reaching task.
The goal reaching
Markov Decision Process (MDP) (Schaul et al., 2015; Andrychowicz et al., 2017) is deﬁned by the tuple ( S , A , G , r,  , ⇢ 0 , ⇢ g ) , where the state, action and goal spaces are indicated by S , A , G 2 R .
In our case, the goal space
G can be an arbitrary subset of the state space.
The initial state s 0 is sampled from the distribution ⇢ 0 , while the goal is sampled from ⇢ g .
Our aim is to learn a stochastic policy ⇡ :
S ⇥ G !
⇤ ( A ) that maximizes the expected discounted return E ⌧ ⇠ ⇢ ⇡ ,s g ⇠ ⇢ g ⇥ P
T  1 t =
t r ( s t , a t , s g ) ⇤ .
3. H INDSIGHT E XPERIENCE R EPLAY
Let us now consider the particular case in which the reward function is sparse, i.e. the indicator function r ( s t , a t , g )
=
I { s t 6 = g } .
This assumption complicates the credit assignment problem and becomes a bottleneck for sample efﬁciency.
To overcome this issue, Andrychowicz et al.
(2017) introduce Hindsight Experience Replay (HER).
The main idea of HER is to relabel past failed trajectories with the indicator function mentioned above, where the goal g is the one that was actually achieved, rather than the desired one.
Although success signals are generally more frequent while using HER, they remain binary, and thus hold little information.
In this work we build upon the relabeling framework introduced in HER, but we propose to compute rewards from a more informative learned dynamical distance.
Our distance, however, is not tied to the goal sampling procedure presented in HER, and can be used with any relabeling scheme or goal-conditioned RL algorithm as a shaped reward.
3. D YNAMICAL D ISTANCES
We learn a distance function that predicts the expected number of timesteps along the shortest path between two states.
This is learned via a self-supervision signal given by the empirical distance between states of a rollout, as done in Hartikainen et al.
(2019), which we extend to an off-policy setting.
In the original paper, the distance d ⇡✓ is associated with a policy ⇡ , which gives rise to the training loss:
[Equation]
where the empirical distance between the states ( j  i ) is relative to the policy ⇡ .
The complete dynamical distance learning (DDL) algorithm uses d ⇡✓ ( s, s g ) as a negative reward to optimize the policy ⇡ and it is guaranteed to converge to the optimal policy ⇡ ⇤ .
However, it requires on-policy data collection, which bottlenecks the efﬁciency of the pipeline.
Moreover, it also requires a small amount of goal proposals to be used during on-policy training, selected by a human operator from the last visited goals.
In practice, the user is shown a batch of achieved goals and has to select the "best" one, which will be used to train the policy ⇡ .
M ETHOD 4. O FF - POLICY T EMPORAL R EGRESSION
To overcome the on-policy bottleneck, and achieve better sample efﬁciency with off-policy data, we propose to condition the dynamical distance on the action value and learn the function d ✓ ( s i , a i , s j ) from a replay buffer by minimizing the following loss:
[Equation]
where ✓ old are previous distance estimates from a frozen network, updated at a slower frequency for stability reasons, and U [ · ] can indicate any measure of uncertainty associated to the network given the input ( s i , a i , s j ) .
Like this, we can train d ✓ from suboptimal off-policy trajectories, as the supervision signal is the minimum value between the upper bound of the true distance, given by the empirical distance ( j  i ) , and a frozen estimate of the dynamical distance.
Intuitively, thanks to the generalization capabilities of powerful approximators such as neural networks, d ✓ old may provide reasonable lower distance estimates, which can be used as a target in the loss, rather than some temporal overestimate ( j  i ) coming from suboptimal trajectories at the beginning of training.
We explain in further detail our choice of uncertainty measure in section 4.2.
4. U NCERTAINTY WITH C OUNTS
We introduce an uncertainty term in eq.
1 so that the min operator discards distance estimates with high uncertainty,
and it starts to consider proposals from the distance network otherwise.
Our formulation of the uncertainty term is inspired by a recent work on pessimistic initialisation of
Q-functions by Rashid et al.
(2020) with count models.
We adopt the same model based on localitysensitive hashing (LSH) (Tang et al., 2017), but we are not bound to any particular counting scheme.
We deﬁne the count model as N ( s i , a i ) and use it to pessimistically initialise the distance in the following way:
[Equation]
where M, C > 0 are hyperparameters.
This penalty term polynomially decays with the number of counts, so that the distance gets updated with the off-policy estimate d ✓ ( s i , a i , s j ) only if that state-action combination has been seen enough times.
This pessimistic term is also used to avoid the trivial zero solution of eq. 1.
4. L OCAL C ONNECTIVITY AND T RIANGULAR L OSS
When the goal is not part of the replay buffer and we are in the presence of obstacles, the learned distance underestimates the shortest path to the goal.
Let us consider the example in ﬁg. 1.
In this case, the initial states distribution ⇢ 0 samples the starting points from the lower left room.
On the other hand, the goals are always sampled from the upper right room.
Unless the agent sees a sufﬁcient amount of paths that lead to the goal, it will not be able to learn the correct shortest path distance.
Contrary to the Floyd-Warshall algorithm, where all the distances are initialized with inﬁnity and then “relaxed" with the triangle
[Figure]
inequality, we consider an alternative approach and start with underestimates caused by the off-policy temporal regression, which we artiﬁcially increase whenever we sample another point in the buffer that is connected to the goal.
In order to express this connectivity information, we employ the count based term already introduced in the previous section to estimate uncertainty.
The count term N ( s i , a i , s g ) then indicates how many times in total
the goal
s g was reached from state s i after taking the ﬁrst action a i .
For instance, let us consider the case in which the goal s g was never reached from state s i , but there is a path in the buffer D from an intermediate state s j to the goal state s g (ﬁg. 1).
In this case, we increase the value of d ✓ ( s i , a i , s g ) as:
[Equation]
The distance d ?✓ ( s i , a i , s g ) is then a penalized hypotenuse, that will converge to the true shortest path distance and thus mitigate local minima.
A LGORITHM S UMMARY
The ﬁnal formulation for the loss comprises the sum of these two terms (from eq. 1 and 3):
[Equation]
( d ✓ old ( s i , a i , s j ) + d ✓ old ( s j , a j , s g )) ◆  (4b) where µ ( s g ) =
{ ( s, a ) | ( s, a ) 2 D last ^ N ( s, a, s g )
= 0 } and µ ( s g )
=
D \ µ ( s g ) where D last is the replay buffer containing the most recently collected rollouts.
Notice that in the triangular loss we can either increase or decrease the value of the hypotenuse.
In one case the hypotenuse is penalized, in the other one it is regularized (relaxed, as in Floyd-Warshall) because a better path was found.
Moreover, in this loss, s g always indicates the desired goal of the trajectory which s i is sampled from.
Thus, we are only correcting the paths that should have led to the goal.
We underline that this procedure does not search for the ( s j , a j ) tuple minimizing the triangle inequality, it purely samples ( s j , a j ) from the points actually present in the replay buffer D .
E XPERIMENTS
Our method is tested on four environments illustrated in ﬁg. 2: 2D navigation A point mass in two dimensions has to navigate a playground with internal walls to reach a goal.
In the setting we present, using the Euclidean goal-distance as a heuristic cost would fail the task because of the presence of local optima created by the walls.
3D reaching task
A 7DoF robot arm has to reach a goal.
In the F ETCH R EACH W ITH W ALL variation, the robot arm is initialized in one half of the table while the goal is placed on the other half, on the ground, with a wall separating the table in two halves.
2 do
D
COLLECT DATA WITH
+ ] // eq.
4a
UPDATE TARGET NETWORK
3D manipulation task
A 7DoF robot arm has to fetch, pick and place a box to a target location.
The goal can be sampled either above or on the table with 50% probability.
In the F ETCH P ICK AND
P LACE W ITH W ALL , the box is always placed on one half of the table, and the goal in the other, so the the only solution is lifting it.
Claw Manipulation (Ahn et al., 2020)
A 9-DOF “claw”-like robot is required to turn a valve to various positions.
The state space includes the positions of each joint of each claw (3 joints onclaws) and embeds the current angle of the valve in Cartesian coordinate ( ✓ !
(sin ✓ , cos ✓ )) .
The robot is controlled via joint angle control.
The goal space consists only of the claw angle, which is sampled uniformly from the unit circle.
For all the baselines and ablations we use HER (Andrychowicz et al., 2017) while learning a policy with Soft Actor Critic Haarnoja et al.
(2018), while for the F ETCH P ICK AND P LACE environment we use HER combined with DDPG (Lillicrap et al., 2016) due to better empirical performance.
All the methods make use of the future strategy presented in Andrychowicz et al.
(2017), each of them with the best k selected by grid search.
The parameters used in the experiments are reported in the appendix.
We consider 4 kinds of reward coupled with HER:
Sparse+HER (baseline) Standard binary reward with hindsight relabeling of the buffer with the default SAC hyperparameters.
DDL+HER (baseline) We use the original HER relabeling scheme, and rewards given by the negative temporal loss (on-policy).
This method is slightly different from (Hartikainen et al., 2019) as it does not make use of the original "user preferences" to propose goals.
Off-policy DDL+HER (ablation)
We augment DDL+HER by introducing the off-policy temporal regression loss term of sec. 4.1. Ours+HER Our complete method, including the off-policy temporal regression loss (eq. 4a) and the bootstrapped triangular loss (eq. 4b), as described in algorithm 2.
F ETCH R EACH
F ETCH R EACH W ITH W ALL
C LAW T URN
P OINT M AZE
F ETCH P ICK AND P LACE
F ETCH P ICK AND P LACE W ITH W ALL
Ours+HER
Off-policy DDL+HER
DDL+HER
Sparse+HER
[Figure]
[Figure]
Sparse+HER DDL+HER Off-Policy DDL+HER
Ours+HER
6. L OCAL O PTIMA
The results in ﬁg. 3 show how our method is robust even when we consider environments with obstacles or goals that are hard to reach.
The off-policy temporal loss can ﬁnd shortcuts as suggested by the function approximator, and introducing the triangular loss further improves performance In particular, the F ETCH P ICK AND P LACE task is an eloquent example.
In this case we opt to augment the goal space to also consider the end effector position: the goal is to have both box and gripper at the target state.
This simple addition already improves the sample efﬁciency of Sparse+HER by an order of magnitude.
However, we can also notice how Sparse+HER quickly reaches a 50% success rate, equivalent to reaching all the goals placed on the table, but cannot generalize so easily to the goals in the air.
In fact, given the augmented goal state, it often reaches a suboptimal solution by placing the end effector in between the current box position and its target position.
While our method takes more time to catch up, it eventually overtakes both baselines.
In ﬁg. 5 we show the effect of the triangular loss for the F ETCH P ICK AND P LACE task.
In particular, we computed the rates at which the distances to the goal get increased (merging rate) or regularized (relaxations rate).
In the F ETCH P ICK AND P LACE WITH W ALL task, it is even more clear that the triangular loss consistently helps avoiding local minima.
The robot arm is able to lift the box and bring it to the other side in most cases, without the need to specify an auxiliary reward to avoid the wall.
More details in Appendix B.1.
6. S AMPLE E FFICIENCY
Being off-policy, our method can make better use of the transitions present in the replay buffer.
In ﬁg. 6 we see how, even if all methods are trained with a ﬁxed exploration noise and a max-entropy policy, only our approach is able to consistently explore the goal distribution (top right room).
In the speciﬁc example of the P OINT M AZE , Off-Policy DDL+HER fails to solve the task on its own as it underestimates the distance to the goal.
In ﬁg. 4 we represent the Q-function learned by SAC for
DDL+HER, and for Ours+HER.
The maximum value of the Q-function (its negative value, in the plot) is not centered around the goal for DDL+HER.
We presume this might be due to estimation errors induced by the (off-policy) exploration noise.
We also achieve better sample efﬁciency for the F ETCH P ICK AND P LACE task.
As shown in ﬁg.(a), we manage to get a maximum reward above the threshold after circa 11k rollouts, employing 45% less samples than DDL+HER, and achieving the same performance of Sparse+HER with less than half of the rollouts.
In the variant W ITH W ALL , the efﬁciency gain is 5x higher.
C ONCLUSIONS
Learning distances or reward functions in a self-supervised fashion is a generally difﬁcult problem, as it relies on bootstrapped estimates, which can produce strongly biased and high-variance solutions.
Solving it, leads to increased sample efﬁciency during learning which is highly relevant for reinforcement learning applied to real robotic tasks.
In addition, learning the distance function does not require any domain knowledge nor expert demonstrations and is thus of general interest for goal-conditioned tasks.
The proposed neural APSP algorithm shows promising empirical results which indicate that learning these distances, and the corresponding goal-conditioned policies, is not only possible, but also robust to local optima and sample efﬁcient.
A CKNOWLEDGEMENTS
The authors thank Marco Bagatella, Cansu Sancaktar and Anselm Paulus for their useful feedback.
This work is funded by the Max Planck ETH Center for Learning Systems and by the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s Excellence Strategy - EXC number 2064/1 - Project number 390727645.