O FFLINE R EINFORCEMENT L EARNING WITH P SEUDOMETRIC L EARNING
A BSTRACT
Ofﬂine Reinforcement Learning methods seek to learn a policy from logged transitions of an environment, without any interaction.
In the presence of function approximation, and under the assumption of limited coverage of the state-action space of the environment, it is necessary to enforce the policy to visit state-action pairs close to the support of logged transitions.
In this work, we propose an iterative procedure to learn a pseudometric from logged transitions, and use it to deﬁne this notion of closeness.
We show its convergence and extend it to the function approximation setting.
We then use this pseudometric to deﬁne a new lookup based bonus in an actor-critic algorithm: PLO FF .
This bonus encourages the actor to stay close, in terms of the deﬁned pseudometric, to the support of logged transitions.
Finally, we evaluate the method on hand manipulation and locomotion tasks.
I NTRODUCTION
Reinforcement Learning (RL) has proven its ability to solve complex problems in recent years (Silver et al., 2016; Vinyals et al., 2019).
Behind those breakthroughs, the adoption of RL in real-world systems remains challenging (Dulac-Arnold et al., 2019).
Learning a policy by trial-and-error, while operating on the system, can be detrimental to the system where it is deployed ( e.g. , user satisfaction in recommendation systems or material damage in robotics) and is not guaranteed to lead to good performance (sparse rewards problems).
Nevertheless, in the setting where experiences were previously collected in an environment ( logged transitions ), one possible way of learning a policy is to mimic the policy that generated these experiences (Pomerleau, 1991).
However, if these experiences come from different sources, with different degrees of desirability, naive imitation might lead to poor results.
On the other hand, ofﬂine
RL (or batch RL) (Lagoudakis & Parr, 2003; Ernst et al., 2005; Riedmiller, 2005; Levine et al., 2020), offers a setting where the policy is learned from the collected experiences.
As it requires no interaction with the environment, ofﬂine RL is a promising direction for learning policies that can be deployed into systems.
Still, collected experiences typically only cover a subset of the range of possibilities in an environment (not every state is present; for a given state, not every action is taken).
With function approximation, this is particularly problematic since actions not executed in the system can be assigned overly optimistic values (especially through bootstrapping), which leads to poor policies.
To limit this extrapolation error, ofﬂine RL is typically incentivized to learn policies that are plausible in light of the collected experiences.
In other words, ofﬂine RL methods need to learn policies that maximize their return, while making sure to remain close to the support of logged transitions.
This work introduces a new method for computing the closeness to the support by learning a pseudometric from collected experiences.
This pseudometric, close to bisimulation metrics (Ferns et al., 2004), computes similarity between state-action pairs based on the expected difference in rewards when following speciﬁc sequences of action.
We show theoretical properties in the dynamic programming setting for deterministic environments as well as for the sampled setting.
We further extend the learning of this pseudometric to the function approximation setting, and propose an architecture to learn it from collected experience.
We deﬁne a new ofﬂine RL actor-critic algorithm:
PLO FF ( P seudometric L earning Off line RL), which computes a bonus through this learned pseudometric and uses it to ﬁlter admissible actions in the greedy step.
Finally, we lead an empirical study on the hand manipulation and locomotion tasks of the ofﬂine RL benchmark from Fu et al.
We make the following contributions: 1) we extend neural bisimulation metrics (Castro, 2020) to state-action spaces and to the ofﬂine RL setting, 2) we exploit this pseudometric to tackle the out-of-distribution extrapolation error of ofﬂine RL by adding a simple lookup bonus to a standard actor-critic algorithm and show that it compares favorably to state-of-the art ofﬂine RL methods.
B ACKGROUND
Reinforcement Learning.
We consider the classic RL setting (Sutton & Barto, 1998), formalized with Markov Decision Processes (MDPs).
An MDP is a tuple M :=
( S , A , r, P, γ ) , with S the state space,
A the action space, r : S × A (cid:55)→ R the expected reward function, P : S × A (cid:55)→ P ( S ) the transition function which maps state-action pairs to distributions over the set of states P ( S ) and γ the discount factor for which we assume γ ∈ [0 , 1) .
A stationary deterministic policy π is a mapping from states to actions (the following can easily be extended to stochastic policies).
The value function V π of a policy π is deﬁned as the expected discounted cumulative reward from starting in a particular state and acting according to π : V π ( s ) =
E (cid:0) (cid:80) ∞
i =
γ i r ( s i , π ( s i )) | s 0 = s ) (cid:1) .
The action-value function
Q π is deﬁned as the expected cumulative reward from starting in a particular state, taking an action and then acting according to π : Q π ( s, a ) =
r ( s, a ) + γ E (cid:0) V π ( s (cid:48) ) (cid:1) .
The Bellman (1957) operator B connects an action-value function Q for state-action couple ( s, a ) to the action-value function in the subsequent state:
B π ( Q )( s, a ) :
= r ( s, a ) + γ E (cid:0) Q ( s (cid:48) , π ( s (cid:48) ))
(cid:1) .
Q π is the (unique) ﬁxed-point of this operator, and the difference between Q ( s, a ) and its image through Bellman (cid:107) Q − B π Q (cid:107) is called a temporal difference error.
An optimal policy π maximizes the value function V π for all states.
In continuous state-action spaces, actor-critic methods (Konda & Tsitsiklis, 2000) are a common paradigm to learn the optimal policy.
In this work we only consider deterministic policies; we justify this restriction by the fact that stochastic policies are desirable because of their side effect of exploration (Haarnoja et al., 2018), but in this case we want to learn a policy with near-optimal behavior without interaction with the environment.
Therefore, we use the actor-critic framework of Silver et al.
It consists in concurrently learning a parametrized policy π θ and its associated parametrized action-value function
Q ω .
Q ω minimizes a temporal difference error (cid:107) Q ω ( s, a ) − r
( s, a )
− Q ¯ ω
( s (cid:48) , π θ ( s (cid:48) ))
(cid:107) (with Q ¯ ω a target action-value function, tracking Q w ), and π θ maximizes the action-value function Q ω ( s, π θ ( s )) .
In the classical RL setting, transitions ( s, a, s (cid:48) , r ) are sampled through interactions with the environment.
In on-policy actor-critic methods (Sutton et al., 1999; Schulman et al., 2015), updates on π θ and Q ω are made as the policy gathers transitions by interacting in the environment.
In off-policy actor-critic methods (Lillicrap et al., 2016; Haarnoja et al., 2018; Fujimoto et al., 2019), the transitions gathered by the policy are stored in a replay buffer and sampled using different sampling strategies (Lin, 1992; Schaul et al., 2015).
These off-policy methods extend to the ofﬂine RL setting quite naturally.
The difference is that in the ofﬂine RL setting, transitions are not sampled through interactions from the environment, but from a ﬁxed dataset of transitions.
Throughout the paper we refer to D
= { ( s i , a i , r i , s (cid:48) i ) } 1: N as the dataset of N transitions collected in the considered environment.
To ease notations we write s ∼ D , s, a ∼ D , r ∼ D to indicate that a transition ( s, a, s (cid:48) , r ) is sampled at random from this dataset, and that we only consider the associated state s , state-action pair ( s, a ) or reward r .
Pseudometric in MDPs.
A core issue in RL is to deﬁne a meaningful metric between states or state-action pairs.
Consider for example a maze, two states could be close according to the Euclidean distance, but far away in terms of the minimal distance an agent would have to travel to join one state from the other (due to walls).
In this case, a relevant metric is the distance in the graph formed from state transitions induced by the MDP (Mahadevan & Maggioni, 2007).
We consider the following relaxed notion of metric 1 :
Deﬁnition 1 (Pseudometric) .
Given a set M , it is a function d : M × M (cid:55)→ R + such that, ∀ x, y, z ∈
M , we have d ( x ; x ) = 0 , d ( x ; y ) =
d ( y ; x ) , d ( x ; z ) ≤ d ( x ; y )
+
d ( y ; z ) .
In the context of Markov Decision Processes, bisimulation relations (Givan et al., 2003) are a form of state abstraction (Li et al., 2006), based on an equivalence relation.
They are deﬁned by the following recurrent deﬁnition: two states are bisimilar if they yield the same expected reward and transition to bisimulation equivalence classes with equal probability.
This deﬁnition is too restrictive to be useful in practice.
Ferns et al.
(2004) introduce bisimulation metrics which are pseudometrics that soften the concept of bisimulation relations.
Bisimulation metrics are deﬁned on the state space S .
Denote M
S the set of pseudometrics on S , the bisimulation metric is the (unique) ﬁxed point of the operator
F S deﬁned as:
[Equation]
where W 1 ( d ) is the 1-Wasserstein distance (Villani, 2008) with the distance between states measured according to the pseudometric d .
Therefore, the bisimulation metric is the limit of the repeated application of the sequence (cid:0) F S n ( d 0 ) (cid:1) n ∈ N for any initial d 0 ∈ M S .
Although pseudometrics in MDPs have proven to be effective in some applications (Melo & Lopes, 2010; Kim & Park, 2018; Dadashi et al., 2021), they are usually hand-crafted or learned with ad-hoc strategies.
M ETHOD
We present the overall idea of our method in Sec.
3.1: an ofﬂine RL algorithm which is incentivized to remain close to the support of collected experiences using a pseudometric-based bonus.
In Sec. 3. to 3.4, we present how to learn this pseudometric, from a theoretical motivation to a gradient-based method.
We give practical considerations to derive the bonus in Sec.
3.5 and provide the resulting algorithm: PLO FF , in Sec.
3. O FFLINE RL WITH LOOKUP BONUS
For the time being, let us assume the existence of a pseudometric d on the state-action space S × A .
We can infer a distance d D from a transition to the dataset D :
[Equation]
This distance to the dataset, also referred to as the projection distance, is simply the distance from ( s, a ) to the nearest element of D .
It is central to our work since it deﬁnes the notion of closeness to the support of transitions D .
From d D , we can infer a bonus b using a monotonously decreasing function f : R (cid:55)→ R : b ( s, a ) =
f ( d D ( s, a )) .
Note that the concept of a bonus is overloaded in RL; it typically applies to exploration strategies (Schmidhuber, 1991; Thrun, 1992; Bellemare et al., 2016).
In our case, the bonus b is opposite to exploration-based bonuses since it will encourage the policy to act similarly to existing transitions of the collected experiences D .
In other words, it prevents exploring too far from the dataset.
We adapt the actor-critic framework by adding the bonus to the actor maximization step.
We learn a parametrized policy π θ and its corresponding action-value function Q ω .
In a schematic way, we sample transitions ( s, a, r, s (cid:48) )
∈ D and minimize the two following losses:
[Equation]
This modiﬁcation of the actor-critic framework is common in ofﬂine RL (Buckman et al., 2021).
The bonus b typically consists in a measure of similarity between an estimated behavior policy that generated D and the policy π we learn (Fujimoto et al., 2018; Wu et al., 2019; Kumar et al., 2019).
3. P SEUDOMETRIC L EARNING
To deﬁne a bonus b , we ﬁrst learn a pseudometric d on the state-action space S ×
A similarly to bisimulation metrics (Ferns et al., 2004; 2011; 2006), with the difference being that we are interested in pseudometrics in state-action space S
× A rather than state space S .
We will show that the pseudometric d we are interested in is the ﬁxed point of an operator F .
In the following, we assume that the MDP is deterministic.
Let M be the set of bounded pseudometrics on S × A .
We deﬁne the operator F : M (cid:55)→ M as follows: for two state-action pairs ( s 1 , a 1 ) and ( s 2 , a 2 ) , that maps to next states s (cid:48)
and s (cid:48)
2 we have:
[Equation]
with U ( A ) the uniform distribution over actions.
This operator is of particular interest as it takes a distance d over state-action pairs as input, and outputs a distance F ( d ) which is the distance between immediate rewards, to which is added the discounted expected distance between the two transitioning states for a random action.
Notice that contrary to bisimulation metrics, we do not use a maximum over next actions for multiple reasons: the use of a maximum can be overly pessimistic when computing the similarity (Castro, 2020); in the case of continuous action spaces a maximum is hard to estimate; and ﬁnally in the presence of function approximation (Section 3.4) it can lead to instabilities.
We now establish a series of properties of the operator F , all being proven in Appendix B.
Proposition 3.1.
Let d be a pseudometric in M , then F ( d ) is a pseudometric in M .
We are now interested in the repeated application of this operator (cid:0) F n ( d 0 ) (cid:1) n ∈ N starting from the 0 -pseudometric d 0 (mapping all pairs to 0).
Proposition 3.2.
Let d be a pseudometric in
M , we note (cid:107) d (cid:107) ∞ as max s,s (cid:48) ∈S max a,a (cid:48)
∈A
d ( s, a ; s (cid:48) , a (cid:48) ) .
The operator F is a γ -contraction for (cid:107) · (cid:107) ∞ .
Since the operator F is a contraction, it follows that the sequence (cid:0) F n ( d 0 ) (cid:1) n ∈ N converges to the pseudometric of interest d ∗ (it would for any initial pseudometric, but d 0 is of particular empirical interest).
Proposition 3.3.
F has a unique ﬁxed point d ∗ in M .
Suppose d 0 ∈ M then lim n →∞ F n ( d 0 ) =
d ∗ .
The ﬁxed point of F can be thought of as the similarity between state-actions, measured by the difference in immediate rewards added to the difference in rewards in future states if the sequence of actions is selected uniformly at random.
In other words, two state-action couples will be close if 1) they yield the same immediate reward and 2) following a random walk from the resulting transiting states yields a similar return.
3. P SEUDOMETRIC LEARNING WITH SAMPLING
Now, we move to the more realistic setting where the rewards and dynamics of the MDP are not known.
Thus, the MDP is available through sampled transitions.
We assume in this section that we have a ﬁnite state-action space.
We deﬁne an operator ˆ F , which is a sampled version of F : suppose we sample a pair of transitions from the environment (ˆ s 1 , ˆ a 1 , ˆ s (cid:48) 1 , ˆ r 1 ) , (ˆ s 2 , ˆ a 2 , ˆ s (cid:48) 2 , ˆ r 2 ) , we have: ˆ F ( d )( s 1 , a 1 ; s 2 , a 2 ) = (cid:26) | ˆ r 1 − ˆ r 2 | + γ E u 1 ,u 2 ∼U ( A )
d (ˆ s (cid:48) 1 , u 1 , ˆ s (cid:48) 2 , u 2 ) if s 1 , a 1 , s 2 , a 2 = ˆ s 1 , ˆ a 1 , ˆ s 2 , ˆ a 2 , d ( s 1 , a 1 ; s 2 , a 2 ) otherwise.
Similarly to what is observed in the context of bisimulation metrics by Castro (2020), the sampled version ˆ F has similar convergence properties as F .
Proposition 3.4.
Under common assumptions of state-action coverage, the repeated application of ˆ
F converges to the ﬁxed point d ∗ of F .
If the environment is stochastic, the repeated application of ˆ F does not converge to the ﬁxed point of the operator F .
In fact, it is not even stable in the space of pseudometrics (the expectation and absolute value are not commutative).
This results appears as a limitation of our work, since it only applies to deterministic environments.
We leave to future work whether we can deﬁne a different pseudometric (ﬁxed point of another operator) that would have convergence guarantees in the sampling case.
As we only evaluate our approach on deterministic environments (Section 4), another interesting direction is whether our approach empirically extends to stochastic environments (even if it is less principled).
3. P SEUDOMETRIC LEARNING WITH APPROXIMATION
Building upon the insights of the previous sections, we derive an approximate version of the iterative scheme, to estimate a near-optimal pseudometric d ∗ .
Pairs of transitions are assumed to be sampled from a ﬁxed dataset D .
We use Siamese neural networks Φ (Bromley et al., 1994) to derive an approximate version of a pseudometric d Φ .
Using Siamese networks to learn a pseudometric (Castro, 2020) is natural since it respects the actual deﬁnition of a pseudometric by design (Deﬁnition 1).
To ease notations, we conﬂate the deﬁnition of the deep network Φ with its parameters.
We deﬁne the pseudometric d Φ as:
[Equation]
where (cid:107) · (cid:107) is the Euclidean distance.
From the ﬁxed-point iteration scheme deﬁned in Section 3.3, we want to deﬁne a loss to retrieve the ﬁxed point d ∗ .
Similarly to ﬁtted value-iteration methods (Bertsekas & Tsitsiklis, 1996; Munos & Szepesvari, 2008), which is the basis for the DQN algorithn (Mnih et al., 2015), we consider the parameters of the image of the operator ˆ F to be ﬁxed, and note it ˆ F ( d ¯Φ ) .
We thus learn d Φ by minimizing the following loss, which is exactly the temporal difference error (cid:0)
[Equation]
We introduce another pair of Siamese networks Ψ (again we conﬂate the deﬁnition of the network with its parameters) to track the bootstrapped estimate
E a (cid:48) ∈U
( A ) d ¯Φ ( s 1 , a (cid:48) , s 2 , a (cid:48) )
[Equation]
that we learn minimizing the following loss:
[Equation]
We justify this design choice in Section 3.
[Equation]
where we show that this makes the derivation of the bonus tractable.
Therefore
[Equation]
we optimize the two following losses ˆ L Ψ and ˆ L Φ :
[Equation]
[Table]
[Equation]
A visual representation of the two pairs of Siamese networks as well as their losses is provided in the
Appendix C. Remark that the proposed architecture seems to present similar caveats as naive ofﬂine
RL approaches (since we estimate quantities that might not be present in the dataset of collected experiences).
However, here, the divergence of the quantity at hand (the pseudometric learned) is unlikely since the goal is to minimize a positive quantity.
This makes the problem of learning d
Φ inherently more stable than learning an optimal policy.
A limitation of this method is that it relies on the reward function r to build similarities between state-action pairs, therefore in very sparse reward environments or with very limited coverage of the state-action space, the quality of the pseudometric learned might conﬂate state-action pairs together, and hence be less adapted to learn a meaningful measure of similarity.
3. T RACTABLE BONUS
Once the pseudometric d Φ is learned, we can deﬁne a lookup bonus introduced in section 3.1.
Given a monotonously decreasing function f , we have: b ( s, a ) =
f (cid:0) min (ˆ s, ˆ a ) ∈D
d Φ (ˆ s, ˆ a ; s, a ) (cid:1) .
This bonus has a complexity that is linear in the size of D and in the dimension of the representation Φ( s, a ) .
As we are considering datasets with large numbers of transitions ( ∼ 10 6 ), this makes each actor step in
Equation (1) computationally expensive.
Therefore we pre-compute the k -nearest neighbors of each state s ∈ D according to the Euclidean distance d Ψ induced by Ψ ; d Ψ ( s 1 , s 2 )
= (cid:107) Ψ( s 1 ) − Ψ
( s 2 ) (cid:107) .
We note: H ( s ) =
(cid:110)
(ˆ s, ˆ a ) ∈
D| ˆ s is a k -nearest neighbor of s for d Ψ (cid:111) .
We infer the approximate distance bonus
¯ b ( s, a ) =
f (cid:0) min ˆ s, ˆ a ∈H ( s )
d Φ ( s, a ; ˆ s, ˆ a ) (cid:1) .
Pre-computing the k -nearest neighbors is expensive (the brute force complexity is quadratic in the size of the dataset, and linear in the dimension of the representation Ψ ).
In our experiments, we use a kd-tree algorithm (Friedman et al., 1977) from scikit-learn (Pedregosa et al., 2011).
With multiprocessing ( ∼ 50 CPUs), precomputing nearest neighbors did not take more than a couple hours even for the largest dataset ( 2 . 10 6 transitions).
If the size of the dataset were to be larger, we can naturally scale our method with approximate nearest neighbor methods.
3. A LGORITHM
We now compile the results from this section and present the pseudocode of our method in Algorithms 1 and 2.
We refer to the combination of both as PLO FF ( P seudometric L earning Off line RL).
1: Initialize Φ , Ψ networks.
2: for step i = 1 to N do 3: Train Φ : min Φ ˆ L Φ 4: Train Ψ : min Ψ ˆ L Ψ 5:
Initialize k -nearest neighbors array H . 6: for step j = 1 to |D| do 7: Compute k
E XPERIMENTS
In this section we conduct an experimental study for the proposed approach.
We evaluate it on a series of hand manipulation tasks (Rajeswaran et al., 2018), as well as MuJoCo locomotion tasks (Todorov et al., 2012; Brockman et al., 2016) with multiple data collection strategies from Fu et al.
We ﬁrst show the details of the learning procedure of the pseudometric, before showing its performance against several baselines from Fu et al.
All implementation details can be found in Appendix C.
4. E VALUATION ENVIRONMENTS
We evaluate PLO FF on four hand manipulation tasks (Rajeswaran et al., 2018): nailing a hammer, opening a door, manipulating a pen and relocating a ball.
We also evaluate PLO FF on MuJoCo locomotion tasks (Brockman et al., 2016) where the goal is to maximize the distance traveled:
Walker2d, HalfCheetah and Hopper.
We provide visualization of the environments in Figure 1.
For each environment we consider multiple datasets D from the D4RL benchmark (Fu et al., 2020).
On hand manipulation tasks, these datasets are the following, ”human”: transitions collected by a human operator, ”cloned”: transitions collected by a policy trained with behavioral cloning interacting in the environment + the initial demonstrations, ”expert”: transitions collected by a ﬁne-tuned RL policy interacting in the environment.
On locomotion tasks, the datasets are the following, ”random”: transitions collected by a random policy, ”medium-replay” the ﬁrst 1M transitions collected by a
SAC agent (Haarnoja et al., 2018) trained from scratch on the environment, ”medium” transitions collected by a policy with suboptimal performance, ”medium-expert”: transitions collected by a near optimal policy + transitions collected by a suboptimal policy.
To have comparable range of rewards between environments, we scale ofﬂine rewards by dividing them by max D r − min D r and learn a policy on this scaled reward (which leaves the optimal policies unchanged).
4. P SEUDOMETRIC LEARNING
We concurrently learn the deep networks Φ and Ψ by minimizing the losses ˆ L Φ and ˆ L Ψ .
State-action pairs are concatenated (and states only in the case of Ψ ) and are passed to a 2-layer network of hidden layers of size (1024, 32).
Note that the concatenation step could be preceded by two disjoint layers to which the state and action are passed (thus making it more handy for visual-based obseravations).
We sample 256 actions to derive the bootstrapped estimate (loss ˆ L Ψ ).
We optimize ˆ L Φ and ˆ L Ψ using the
Adam optimizer (Kingma & Ba, 2015) with batches of state-action pairs and states of size 256.
We show in Figure 2 the decreasing learning curves for ˆ L Φ and ˆ L Ψ .
In Figure 3, we show that the distribution of the learned distance d
Φ between state-action couples and perturbated versions of themselves (with Gaussian noise either on the state or the action).
We show that the distance respects the intuition that the greater the perturbation is, the larger the distance becomes.
Finally we provide visualization of the state similarities learned by Ψ in Appendix E.
4. A GENT TRAINING WITH PSEUDOMETRIC BONUS
In this section we empirically evaluate the performance of an agent trained with a bonus based on the pseudometric.
In our experiments we use TD3 (Fujimoto et al., 2019).
It is an off-policy actor-critic algorithm (Konda & Tsitsiklis, 2000) which builds on top of DDPG (Lillicrap et al., 2016).
We use an implementation of TD3 where we load the dataset of experiences D in its replay buffer.
We concurrently learn a policy π θ and an action-value function Q ω .
The critic step is the same as the one described in Fujimoto et al.
However we modify the actor loss to incorporate the pseudometric bonus, as described in Algorithm 2: max θ E s ∼D (cid:0) Q ω ( s, π θ ( s ))
+ ¯ b ( s, π θ ( s )) (cid:17) .
We found that the bonus ¯ b ( s, a ) :
= α max(0 , Q ¯ ω ( s, a )) exp( − βd D
( s, a )) leads to strong empirical results.
This bonus form is quite natural, since it uses the action-value function to scale the value of the bonus.
We limit the inﬂuence of the bonus on positive action-values (in the case of negative actionvalues, our bonus would incentivize to maximize the distance to the dataset d D ).
Another natural way to deﬁne the bonus would be to shift the reward function to make it positive (but problematic in environments where early termination is desirable).
We ran a hyperparameter search on α ∈ { 0 . 1 , 0 . 5 , 1 . } and β ∈ { 0 . 1 , 0 .
and the number of gradient steps N ∈ { 100000 , 500000 } .
We select the best hyperparameters for each family of environments (locomotion and hand manipulation), and re-run for 10 seeds.
For each seed we evaluate the resulting policy on 10 episodes.
We report aggregated performance in Table 2 and the complete breakdown in Table 2.
We compare the method with numerous baselines:
AWR (Peng et al., 2019),
Behavioral Cloning (Pomerleau, 1991), BEAR (Kumar et al., 2019), BRAC (Wu et al., 2019), BCQ (Fujimoto et al., 2018) and CQL (Kumar et al., 2020).
[Table]
[Equation]
[Table]
Table 1 shows the performance of PLO FF on the D4RL benchmark.
It tops other methods on hand manipulation tasks, and tops all methods but CQL on locomotion tasks.
However, even if PLO FF performs well across the board, the approach does not solve the common failure cases shared by all methods (random datasets as well as datasets with human operated transitions, see Table 2).
We leave to future work whether PLO FF , combined with other methods (in particular CQL which operates on the critic side rather than the actor side) could lead to a solution to these failure cases.
C ONCLUDING R EMARKS
We introduced a new paradigm to compute a pseudometric-based bonus for ofﬂine RL.
We learn policies consistent with the behavior policy that generated the collected transitions, and hence reduce action extrapolation error.
We showed how to derive a pseudometric from logged transitions, extending existing work from
Castro (2020) from the online to the ofﬂine setting and from pseudometrics on state space to pseudometrics on state-action space.
We showed that the pseudometric we desire to learn is the ﬁxed point of an operator, and we provide a neural architecture as well as a loss to learn it.
Conceptually, our bonus introduces a larger computational cost against other approaches that reduce extrapolation errors.
We argue that this is actually a desirable direction of research for ofﬂine RL.
In the presence of a ﬁxed dataset of transitions, and since we cannot add new transitions into memory, we should insist on the other side of the well-known memory-computation trade-off.
We demonstrated in our experimental study that our method performs comparably to existing stateof-the-art methods (tops other methods on hand manipulation task, second to top on locomotion tasks).
A R ELATED W ORK
Ofﬂine Reinforcement Learning.
Ofﬂine RL (Lagoudakis & Parr, 2003; Ernst et al., 2005; Riedmiller, 2005; Pietquin et al., 2011; Lange et al.; Levine et al., 2020)
methods suffer from overestimation of state-action pairs that are not in the support of logged transitions.
A number of methods have been explored to mitigate this phenomenon, by constraining the learned policy to be close in terms of a probabilistic distance to the behavioral policy (Jaques et al., 2019; Wu et al., 2019; Kumar et al., 2019; Siegel et al., 2020; Peng et al., 2019; Fujimoto et al., 2018), or a pessimistic (Buckman et al., 2021) estimate of either the Q-value or the bootstrapped Q-value (Kumar et al., 2020; Laroche et al., 2019; Sim˜ao et al., 2020; Nadjahi et al., 2019).
Fujimoto et al.
; Kumar et al.
; Wu et al.
(2019) estimate the behavior policy using a conditional VAE (Kingma & Welling, 2014), and constrain the policy learned ofﬂine with a distance to the behavior policy using the Kullback-Leibler divergence, the Wasserstein distance and the Maximum-Mean Discrepancy respectively.
In this work we trade the problem of estimating the behavioral policy (which is specially problematic if it is multimodal) with the computational cost of a lookup; and instead of penalizing the policy with a distribution-based distance, we learn a task relevant pseudometric instead.
As our bonus is based on learned structural properties of the MDP we can also draw a connection to model-based off-policy
RL (Kidambi et al., 2020; Yu et al., 2020; Argenson & Dulac-Arnold, 2021).
State-action similarity in MDPs.
Bisimulation relations are a form of state abstraction (Li et al., 2006), introduced in the context of MDPs by Givan et al.
(2003), where states with the same rewards and transitions are aggregated.
As this deﬁnition is strict, a number of works deﬁne approximate versions of it.
For example, Dean & Givan (1997) use a bound rather than an exact equivalence.
Similarly, Ferns et al.
(2004) introduce bisimulation metrics (Ferns et al., 2011; 2006), which use the reward signal to decide the proximity between two states.
This makes bisimulation metrics close to the difference between optimal values between two states (Ferns & Precup, 2014).
Learning a bisimulation metric online (Castro, 2020; Comanici et al., 2012) has been shown to be beneﬁcial to learn controllable representations that eliminates unnecessary details of the state (Zhang et al., 2021) or can be used as an auxiliary loss which leads to improvement performance on the Atari benchmark (Gelada et al., 2019).
Castro (2020) introduces the Siamese network architecture, as well as a loss, to derive the bisimulation metric online.
Our work differs as it learns a pseudometric on state-action pairs rather than states, and is based on ofﬂine transitions.
Ravindran & Barto (2003) introduces MDP homomorphism as another form of abstraction which is state-action dependent rather than state dependent.
Again as the partitioning induced by a homomorphism is too restrictive to be useful in practice, a number of work has looked into relaxed versions (Ravindran & Barto, 2004; Wolfe & Barto, 2006; van der Pol et al., 2020; Taylor et al., 2008).
Our work is closer to bisimulation metrics since it introduces the similarity between states-actions as a difference between reward accumulated when following the same sequence of actions (except for the ﬁrst one).
B P ROOFS
Proposition 3.1.
Let d be a pseudometric in M , then F ( d ) is a pseudometric in M .
Proof.
Let d be a pseudometric in M
[Equation]
we show that F ( d ) respects all properties in Deﬁnition 1 and therefore is a pseudometric.
Let ( s 1 , a 1 ) , ( s 2 , a 2 ) , ( s 3 , a 3 ) ∈ S × A and their associated rewards r 1 , r 2 , r 3 and next states s (cid:48) 1 , s (cid:48) 2 , s (cid:48) 3 :
• the pseudo-distance of a couple to itself is null:
[Equation]
d
[Equation]
• symmetry:
[Equation]
• triangular inequality:
[Equation]
Proposition
[Equation]
Let d be a pseudometric in
M , we note (cid:107) d (cid:107) ∞ as max s,s
[Equation]
∈S max a,a
[Equation]
∈A d ( s, a ; s (cid:48) , a (cid:48) ) .
The operator F is a γ -contraction for (cid:107) · (cid:107) ∞ .
Proof.
Let d 1 , d 2 ∈ M , let ( s 1 , a 1 ) , ( s 2 , a 2 ) ∈ S × A and their associated rewards r 1 , r 2 and next states s (cid:48) 1 , s (cid:48) 2 , we have:
[Equation]
Therefore, we have:
[Equation]
| [Equation] We thus have that (cid:107)F ( d 1 ) −F ( d 2 ) (cid:107) ∞ ≤ γ (cid:107) d 1 − d 2 (cid:107) ∞ , therefore F is a γ -contraction for (cid:107)·(cid:107) ∞ .
Proposition 3.3.
F has a unique ﬁxed point d ∗ in M .
Suppose d 0 ∈ M then lim n →∞ F n ( d 0 ) =
d ∗ .
Proof.
This is a direct application of the Banach theorem (Banach, 1922).
F is a γ -contracting operator with γ ∈ [0 , 1) , in the metric space (( S ×A ) × ( S ×A ) , (cid:107)·(cid:107) ∞ ) , therefore using the Banach theorem we have that F has a unique ﬁxed point d ∗ and ∀ d 0 ∈ M , lim n →∞ F n ( d 0 ) =
d ∗ .
Proposition 3.4.
Under common assumptions of state-action coverage, the repeated application of ˆ
F converges to the ﬁxed point d ∗ of F .
Proof.
The state-action coverage assumption is the following: ∃ (cid:15) > 0 such that for any pairs of state-action pairs ( s, a ) , (ˆ s, ˆ a ) ∈ ( S ×A ) ×
( S ×A ) , ( s, a ) , (ˆ s, ˆ a ) is sampled with at least probability (cid:15) .
The repeated application of ˆ F is an asynchronous ﬁxed point iteration scheme.
The convergence to d ∗ (almost surely) is a direct application of Proposition 3 from Bertsekas & Tsitsiklis (1991).
Note that the state-action coverage assumption enables to apply this result since all pairs of state-action are visited an inﬁnite number of times (almost surely).
C I MPLEMENTATION
In this section, we provide a detailed description of the experimental study.
We are working on making the code open-source.
Ofﬂine datasets preprocessing.
We use datasets from Fu et al.
We scale the rewards by dividing them by max r ∼D r − min r ∼D r for both pseudometric learning and policy learning.
This enables to have comparable range of rewards between environments.
Pseudometric learning.
The Siamese networks Φ et Ψ have the same architecture: a two-layer
MLP of size (1024 , 32) with relu activation on top of the ﬁrst layer.
Note that Φ takes the concatenation of state and action as input, whereas Ψ only takes the state as an input.
We use a discount factor γ = 0 .
9 (which is different than the discount factor from the agent) in the experiments (we noticed some instabilities on human datasets, which contains less transitions than the rest of the datasets, if the discount factor is larger).
We minimize the losses ˆ L Φ and ˆ L Ψ using the Adam optimizer with a learning rate 10 − 3 .
The batch size used to compute the losses is 256 .
The bootstrapped estimate in ˆ L Φ is estimated with 256 actions sampled uniformly.
We train the two networks by iteratively taking a gradient step on each loss for 2 .
10 4 gradient steps.
Once the Ψ network is trained, we derive the k -nearest neighbors of each state in D for the distance induced by Ψ , with k = 50 .
The nearest neighbors are computed using the scikit-learn implementation of the kd-tree algorithm, taking advantage of multiprocessing (with 50 CPUs).
Agent training.
We re-implemented the TD3 agent from Fujimoto et al.
(2019) in JAX (Bradbury et al., 2018).
We use the default hyperparameters (and did not perform HP search) for the exception of the target noise in the bootstrapped estimate that we removed (in the hopper environment, we noticed that target noise in the bootstrapped estimate resulted in overestimation and Q-value divergence).
For the critic, we used a three-layer network with size (256 , 256 , 1) with relu activation on top of the ﬁrst two layers.
For the policy, we used a three-layer network with size (256 , 256 , |A| ) where |A| is the dimension of the action space, with relu activation on top of the ﬁrst two layers and tanh activation on top of the last layer.
We used the Adam optimizer for both the actor and the critic and used a learning rate of 3 .
10 − 4 (consistently with Fujimoto et al.
and trained them using batch of transitions of size 256 sampled uniformly in D .
We led experiments with the following bonuses ¯ b (and focused on the ﬁrst one as it led to better empirical performance):
• ¯ b ( s, a ) =
α max(0 , Q ¯ ω ( s, a )) exp( − βd D
( s, a ))
• ¯ b
( s, a )
= αQ ¯ ω ( s, a ) exp( − βd D
( s, a ))
• ¯ b
( s, a ) =
α exp( − βd D ( s, a )) • ¯ b
( s, a ) =
α (1 − exp( βd D ( s, a )))
We led a hyperparameter search for both the locomotion environments and the hand manipulation environments on α, β, N where α, β are related to the bonus and N is the number of gradient steps.
We selected α in { 0 . 1 , 0 . 5 , 1 . } , β ∈ { 0 . 1 , 0 .
and the number of gradient steps N ∈ { 100000 , 500000 } as the better combination on the average normalized performance on the tasks (averaged over 3 seeds).
We re-ran the best combination of hyperparameters for 10 seeds and report results averaged over the 10 seeds and 10 evaluation episodes per seed.
We found that the best combination for hand manipulation tasks was α = 1 , β = 0 .
25 , N = 100000 and for locomotion tasks was α = 0 .
1 , β = 0 .
5 , N = 500000 .
D B ASELINES
In this section we provide the results for simpliﬁed versions of our method.
The ﬁrst one is the performance of TD3 without any bonus that we note TD3- OFF .
The second one is similar to PLO FF although with a bonus based on the Euclidean distance between the concatenation of the state and the action rather than a learned distance, we refer to it as PLO FF -L2.
For PLO FF -L2, we used the same experimental evaluation protocol as the one described in the previous section.
We led the same hyperparameter search and we found that for both set of tasks, the best combination was α = 0 .
1 , β = 0 .
5 , N = 100000 .
The results in Table 2 show that without a bonus, the performance of policy learned is mediocre.
Also notice that PLO FF -L2 reaches strong results on hand manipulations tasks as it signiﬁcantly tops all other methods, and reaches relatively good performance on locomotion tasks as it is third to CQL and PLO FF .
This should not come as a surprise since the L2 distance has been shown to be effective in some continuous control tasks ( e.g. (Dadashi et al., 2021)).
Note however that in more complex tasks (typically vision-based environments), the Euclidean distance is a poor measure of similarity.
A l go r it h m
BC
B E A
R
BR
A C - p
BR
A C - v
A W
R
BC
Q
C Q
L
P L
O FF
P L
O FF - L
T D 3 - O
FF h a l f c h ee t a h -r a n do m - v 2 .25 .
[-.35 .15 .± 
-][Table]
± 0 .
[Table]
[Table] [Table]
[Table]
.w
a l k e r 2d -r a nd o m - v 1 .7 .- 0 .1 .
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
hopp e r-r a ndo m - v 9 .11 .11 .12 .
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
h a l f c
[Table]
h
[Table]
m e d i u m - v
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
w a l k e r 2d - m e d i u m - v 6 .
.73 .±
[Table] [Table]
[Table]
0 .±
.hopp e r- m e d i u m - v 29 .
.38 .±
15 .47 .± 23 .0 .±
.h
a l f c h ee t a h - m e d i u m -r e p l a y - v 38 .
.40 .±
.39 .±
.35 .±
3 .w
a l k e r 2d - m e d i u m -r e p l a y - v 11 .19 .- 0 .
.10 .± 5 .12 .± 4 .8 .±
4 .hopp e r- m e d i u m -r e p l a y - v 11 .33 .0 .
.48 .28 .±
1 .28 .± 1 .9 .±
9 .h a l f c h ee t a h - m e d i u m - e xp e r tv 35 .53 .44 .41 .52 .
.67 .± 10 .51 .± 13 .3 .±
1 .w
a l k e r 2d - m e d i u m - e xp e r tv 6 .40 .76 .81 .53 .57 .
.85 .± 17 .74 .±
7 .0 .±
1 .hopp e r- m
e d i u m - e xp e r tv 111 .96 .1 .
[Table]
± 11 .
[Table]
± 11 .5 .
[Table]
M ea n P e rf o r m a n ce 25 .
[Table]
± 6 .
[Table]
p e nhu m a nv 34 .- 1 .8 .
[Table]
[Table]
[Table]
h a mm e rhu m a nv 1 .0 .0 .0 .1 .
[Table]
[Table] [Table] [Table]
[Table]
doo rhu m a nv 0 .- 0 .- 0 .- 0 .0 .- 0 .9 .
[Table]
r e l o ca t ehu m a nv 0 .- 0 .- 0 .- 0 .- 0 .- 0 .0 .- 0 .
[Table]
[Table]
p e nc l on e dv 56 .26 .1 .
[Table]
h a mm e rc l on e dv 0 .0 .0 .0 .0 .
[Table]
[Table] [Table]
doo r-
c l on e d - v - 0 .- 0 .- 0 .- 0 .0 .
[Table]
r e l o ca t e -
[Table]
l on e dv
[Table]
[Table]
[Table]
p e ne xp e r tv 85 .105 .-
[Table]
[Table] [Table]
[Table]
h a mm e re xp e r tv 125 .127 .0 .
[Table]
[Table] [Table]
[Table]
doo r-
e xp e r tv 34 .103 .- 0 .-
[Table] [Table]
[Table]
[Table]
r e l o ca t ee xp e r [-tv 101 .98 .- 0 .- 
-]t -
[Table]
[Table]
[Table] [Table]
- 0 .± 0 .
M ea n p e rf o r m a n ce 36 .38 .0 .-
.41 .± 8 .45 .± 5 .- 0 .± 0 .
T a b l e 2 :
E v a l u a ti on o f P
L O
FF , P
L O
FF - L 2 , T
D 3O FF .
W e r e po r t t
h e r e s u lt s o f t
h e b a s e li n e u s i
ng p e rf
o r m a n ce r e s u lt s r e po r t e d by
F u e t a l .
( ) , w h i c
h do no t i n c o r po r a t e s t a n d a r d d e v i a ti on o f p e rf
o r m a n ce s , s i n ce t
h e nu m b e r s a r e b a s e d on 3 s ee
d s .
I n ou r ca s e w e u s e 10 s ee d s , f
o ll o w i ng r ec o mm e nd a ti on s fr o m
H e nd e r s on e t a l .
( ) , a nd e v a l u a t e on 10 e p i s od e s f
o r ea c h s ee
d b e f o r e r e po r ti ng a v e r a g e a nd s t a nd a r d d e v i a ti on s o f p e rf o r m a n ce .
E M ETRIC V ISUALIZATION
In this section, we show the state similarity learned by PLO FF ( Ψ network) and visualize it for
MuJoCo locomotion environments (we could not provide such visualizations on Adroit tasks since states cannot be retrieved from observations, which is the condition necessary to render images from observations).