Curriculum Reinforcement Learning using Optimal Transport via Gradual Domain Adaptation
Abstract
Curriculum Reinforcement Learning (CRL) aims to create a sequence of tasks, starting from easy ones and gradually learning towards difficult tasks.
In this work, we focus on the idea of framing CRL as interpolations between a source (auxiliary) and a target task distribution.
Although existing studies have shown the great potential of this idea, it remains unclear how to formally quantify and generate the movement between task distributions.
Inspired by the insights from gradual domain adaptation in semi-supervised learning, we create a natural curriculum by breaking down the potentially large task distributional shift in CRL into smaller shifts.
We propose GRADIENT, which formulates CRL as an optimal transport problem with a tailored distance metric between tasks.
Specifically, we generate a sequence of task distributions as a geodesic interpolation between the source and target distributions, which are actually the Wasserstein barycenter.
Different from many existing methods, our algorithm considers a task-dependent contextual distance metric and is capable of handling nonparametric distributions in both continuous and discrete context settings.
In addition, we theoretically show that GRADIENT enables smooth transfer between subsequent stages in the curriculum under certain conditions.
We conduct extensive experiments in locomotion and manipulation tasks and show that our proposed GRADIENT achieves higher performance than baselines in terms of learning efficiency and asymptotic performance.
Introduction
Reinforcement Learning (RL) [1] has demonstrated great potential in solving complex decisionmaking tasks [2], including but not limited to video games [3], chess [4] and robotic manipulation [5].
Among them, various prior works highlight daunting challenges resulting from sparse rewards.
For example, in the maze navigation task, since the agent needs to navigate from the initial position to the goal to receive a positive reward, the task requires a large amount of randomized exploration.
One solution to address this issue is Curriculum Reinforcement Learning (CRL) [6, 7] of which the objective is to create a sequence of environments to facilitate the learning of some difficult tasks.
Although there are different interpretations of CRL, we focus on the one that views a curriculum as a sequence of task distributions that interpolate between a source (auxiliary) task distribution and a target task distribution [8, 9, 10, 11, 12, 13].
This interpretation allows more general presentations of tasks and a wider range of objectives such as generalization (uniform distribution over the task collection) or learning to complete difficult tasks (a subset of the task collection).
Again we use the maze navigation as an example.
Given a fixed maze layout and goal, a task is then defined by a start position, and the task distribution is a categorical distribution over all possible start positions.
With a target distribution putting mass over start positions far away from the goal position, a natural
[Figure]
curriculum to accelerate the learning is putting start positions close to goal position first moving
However, most of the existing methods, that are interpreting the curriculum as shifting distributions, use Kullback–Leibler (KL) divergence to measure the distance between distributions.
This setting imposes several restrictions.
First, due to either problem formulations or the computational feasibility, existing methods often require the distribution to be parameterized, e.g., Gaussian [8, 9, 10, 13], which limits the usage in practice.
Second, most of the existing algorithms using KL divergence implicitly assume an l 2 Euclidean space which ignores the manifold structure when parameterizing
RL environments [14].
In light of the aforementioned issues with the existing CRL method, we propose GRADIENT, an algorithm that creates a sequence of task distributions gradually morphing from the source to the target distribution using Optimal Transport (OT).
GRADIENT approaches CRL from a gradual domain adaptation (GDA) perspective, breaking the potentially large domain shift between the source and the target into smaller shifts to enable efficient and smooth policy transfer.
In this work, we first define a distance metric between the individual task, then we can find a series of task distributions that interpolate between the easy and the difficult task distribution.
GRADIENT is able to deal with both discrete and continuous environment parameter space, and nonparametric distributions (represented either by explicit categorical distributions or implicit empirical distributions of particles).
Under some conditions [15], GRADIENT provably ensures a smooth adaptation from one stage to the next.
We summarize our main contributions as follows:
We propose GRADIENT, a novel CRL framework based on optimal transport to generate gradually morphing intermediate task distributions.
As a result, GRADIENT requires little effort to transfer between subsequent stages and therefore improves the learning efficiency towards difficult tasks.
We develop π -contextual-distance to measure the task similarity and compute the Wasserstein barycenters as intermediate task distributions.
Our proposed method is able to deal with both continuous and discrete context space as well nonparametric distributions.
We also prove the theoretical bound of policy transfer which leads to practical insights.
We demonstrate empirically that our method improves the learning efficiency and asymptotic performance in a wide range of tasks compared with state-of-the-art CRL baselines.
Related Work
Curriculum reinforcement learning.
Curriculum reinforcement learning (CRL) [6, 16] focuses on the generation of training environments for RL agents.
There are several objectives in CRL: improving learning efficiency towards difficult tasks ( time-to-threshold ), maximum return ( asymptotic performance ), or transfer performance in unseen tasks ( generalization ).
From a domain randomization perspective, Active Domain Randomization [5, 17] uses curricula to diversify the physical parameters of the simulator to facilitate the generalization in sim-to-real transfer.
From a game-theoretical perspective, adversarial training is also developed to improve the robustness of RL agents in unseen environments [18, 19, 20, 21].
From a “intrinsic motivation” perspective, methods have been proposed to create curricula even in the absence of a target task to be accomplished [22, 13, 23].
CRL as an interpolation of distributions.
In this work, we focus on another stream of works that interprets CRL as an explicit interpolation between an auxiliary task distribution and a difficult task distribution [8, 9, 10, 11].
Self-Paced Reinforcement Learning (SPRL) [8] is proposed to generate intermediate distributions by measuring the task distribution similarity using Kullback–Leibler (KL) divergence.
However, as we will show in this paper, KL divergence brings several shortcomings that may impede the usage of the algorithms.
First, although the formulation of [8, 9, 10, 11] does not restrict the distribution class, the algorithm realization requires the explicit computation of KL divergence, which is analytically tractably only for a restricted family of distributions.
Second, using
KL divergence implicitly assumes an l 2 Euclidean space which ignores the manifold structure when parameterizing RL environments.
In this work, we use Wasserstein distance instead of KL divergence to measure the distance between distributions.
Unlike KL divergence, Wasserstein distance considers the ground metric information and opens up a wide variety of task distance measures.
CRL using Optimal Transport.
Hindsight Goal Generation (HGG) [24] aims to solve the poor exploration problem in Hindsight Experience Replay (HER).
HGG computes 2-Wasserstein barycenter approximately to guide the hindsight goals towards the target distribution in an implicit curriculum.
Concurrent to our work, CURROT [25] also uses optimal transport to generate intermediate tasks.
CURROT formulate CRL as a constrained optimization with 2-Wasserstein distance to measure the distance between distribution.
The main difference is that we propose task-dependent contextual distance metrics and directly treat the interpolation as the geodesic from the source to the target distribution.
Gradual domain adaptation in semi-supervised learning Gradual domain adaptation (GDA) [26, 27, 28, 29, 30, 31, 32, 33] considers the problem of transferring a classifier trained in a source domain with labeled data, to a target domain with unlabelled data, given only unlabelled data that gradually shifts from the source to the target domain.
GDA solves this problem by designing a sequence of learning tasks.
The classifier is retrained with the pseudolabels created by the classifier from the last stage in the sequence.
Most of the existing literature assumes that there exist intermediate domains.
However, there are a few works aiming to tackle the problem when intermediate domains, or the index (i.e., stage in the curriculum), are not readily available.
A coarse-to-fine framework is proposed to sort and index intermediate domain data [33].
Another work [32] proposes to create virtual samples from intermediate distributions by interpolating representations of examples from source and target domains and suggests using the optimal transport map to create interpolate data in semi-supervised learning.
A recent work [27] demonstrates theoretically that the optimal path of samples is the geodesic interpolation defined by the optimal transport map.
Our work is inspired by the “divide and conquer” methodology in GDA and also uses geodesics as our curriculum plan (although in a different task space).
Preliminary 3.1 Contextual Markov Decision Process
A contextual Markov decision process (CMDP) extends the standard single-task MDP to a multitask setting.
In this work, we consider discounted infinite-horizon CMDPs, represented as a tuple
M = ( S , C , A , R, P, p 0 , ρ, γ ) .
Here, S is the state space, C is the context space,
A is the action space, R : S ×
A × C (cid:55)→ R is the context-dependent reward function, P : S × A × C (cid:55)→ ∆( S ) is the context-dependent transition function, p 0 : C
(cid:55)→ ∆( S ) is the context-dependent initial state distribution, ρ ∈ ∆( C ) is the context distribution and γ ∈ (0 , 1) is the discount factor.
Note that goal-conditioned reinforcement learning [12] can be considered as a special case of the CMDP.
To sample a trajectory τ := { s t , a t , r t } ∞ t =0 in CMDPs, the context c ∼ ρ is randomly generated by the environment at the beginning of each episode.
With the initial state s 0 ∼ p 0 ( · | c ) , at each time step t , the agent follows a policy π to select an action a t ∼ π ( s t ) and receives a reward R ( s t , a t , c ) .
Then the environment transits to the next state s t +1 ∼ P ( · | s t , a t , c ) .
Contextual reinforcement learning naturally extends the original RL objective to include the context distribution ρ .
To find the optimal policy, we need to solve the following optimization problem:
[Equation]
3.2 Optimal Transport
Wasserstein distance.
The Kantorovich problem [34], a classic problem in optimal transport [35], aims to find the optimal coupling θ ∗ which minimizes the transportation cost between measures µ, ν ∈
M ( C ) .
Therefore, the Wasserstein distance defines the distance between probability distribution:
[Equation]
subject to
[Equation]
where C is the support space, Θ( µ, ν ) is the set of all couplings between µ and ν , d ( · , · ) : C × C (cid:55)→ R ≥ 0 is a distance function, γ C is the projection from C × C onto C , and T # P generally denotes the push-forward measure of P by a map T [36, 37, 38].
This optimization is well-defined and the optimal θ always exists under mild conditions [35].
Wasserstein Geodesics and Wasserstein Barycenter.
To construct a curriculum that allows the agent to efficiently solve the difficult task distribution, we follow the Wasserstein geodesic [39], the shortest path [27] between the source and the target distributions, under the Wasserstein distance.
While the original Wasserstein barycenter [40, 41] problem concerns the Frèchet mean of multiple probability measures in a space endowed with the Wasserstein metric, we consider only two distributions µ 0 and µ 1 .
The set of barycenters between µ 0 and µ 1 is the geodesic curve given by
McCann’s interpolation [42, 43].
Thus, the interpolation between two given distributions µ 0 and µis defined as:
[Equation]
where α ∈ [0 , 1] specifies one unique interpolation distribution on the geodesic.
While the computational cost of Wasserstein distance objective eq.(2) could be a potential obstacle, we follow the entropic optimal transport and utilize the celebrated Sinkhorn’s algorithm [44].
Moreover, we adopt smoothing bias to solve for scalable debiased Sinkhorn barycenters [45].
Curriculum Reinforcement Learning using Optimal Transport
We first formulate the curriculum generation as an interpolation problem of probability distributions in Section 4.1.
Then we propose a distance measure between contexts in Section 4.2 in order to compute the Wasserstein barycenter.
Next, we show our main algorithm, GRADIENT, in Section 4.
and an associated theorem that provides practical insights in Section 4.4.
4.1 Problem Formulation
Formally, given a target task distribution ν ( c ) ∈ ∆( C ) , we aim to automatically generate a curriculum of task distributions, ρ 0 ( c ) , ρ 1 ( c ) , . .
, ρ K ( c ) with K stages, that enables the agent to gradually adapt from the auxiliary task distribution µ ( c ) to the target ν
( c ) , i.e., ρ 0 ( c ) → µ ( c ) and ρ K ( c ) → ν ( c ) .
If the context space C is discrete (sometimes called fixed-support in OT [35]) with cardinality |C| , the task distribution is represented by a categorical distribution, e.g., µ ( c ) =
[ p ( c 1 ) , p ( c 2 ) , ..., p ( c |C| )] , where p ( c i ) ≥ 0 , (cid:80) i p ( c i ) = 1 .
Whereas if the context space C is continuous (sometimes called free-support in OT), the task distribution is then approximated by a set of particles sampled from the distribution, e.g., µ ( c ) ≈ ˆ µ ( c ) = 1 n s (cid:80) n s i =1 δ c si ( c ) , where δ c si ( c ) is a Dirac delta at c si .
This highlights the capability of dealing with both parametric and nonparametric distributions in our formulation and algorithms, in contrast to existing algorithms [8, 13, 8, 9, 10, 11].
4.2 Contextual Distance Metrics
To define the distance between contexts, we start by defining the distance between states.
Bisimulation metrics [46, 47, 48] measure the "behaviorally equivalence" between states by defining a recursive relationship between states.
Recent literature in RL [49, 14] uses the bisimulation concept to train state encoders that facilitate multi-tasking.
[50] uses this bisimulation concept to enforce the policy to visit state-action pairs close to the support of logged transitions in offline RL.
However, the bisimulation metric is inherently “pessimistic” because they consider equivalence under all actions, even including the actions that never lead to positive outcomes for the agent [48].
Even if V ∗ ( s i )
=
V ∗ ( s j ) , the bisimulation metric could be arbitrarily large.
To address this issue, [48] proposes on-policy π -bisimulation that considers the dynamics induced by the policy π .
Similarly, we extend this notion to the CMDP settings and define the contextual π -bisimulation metric :
[Equation]
where R πs,c := (cid:80) a π ( a | s )
R ( s, a, c ) and P πs,c := (cid:80) a P ( · | s, a, c ) π ( a | s ) .
With the definition of contextual π -bisimulation metric, we are ready to propose the π -contextual-distance to measure the distance between two contexts:
Definition 4.1 ( π -contextual-distance)
Given a CMDP M =
( S , C , A , R, P, p 0 , ρ, γ, ) , the distance between two context d
π ( c i , c j ) under the policy π is defined as
[Equation]
Conceptually, the π -contextual-distance approximately measures the performance difference between two contexts
c i and c j under π .
The algorithm to compute a simplified version of this metric under some conditions is detailed in the Appendix C.2.
Note that it is difficult to compute this metric precisely in general.
In practice, we can design and compute some surrogate contextual distance metrics, depending on the specific tasks.
There are situations where it is reasonable to use l 2 as a surrogate metric when the contextual distance resembles the Euclidean space, such as in some goal-conditioned continuous environments [24].
In addition, although the contextual distance is not a strict metric, we can still use it as a ground metric in the OT computation.
With the concepts of a contextual distance metric, we now introduce the algorithms to generate intermediate task distribution to enable the agent to gradually transfer from the source to the target tasks distribution.
4.3 Algorithms
We present our main algorithm in Algorithm 1.
We introduce an interpolation factor ∆ α to decide the “difference” between two subsequent task distributions in the curriculum.
Smaller ∆ α means a smaller difference between subsequent task distributions.
In this work, we treat ∆ α as a constant for simplicity and show empirically that constant ∆ α already performs well in our experiments.
We leave the adaptive ∆ α to future works.
Note that the ∆ α is analogical to the KL divergence constraint in [9, 10, 8].
The main difference is that we use Wasserstein distance instead of KL divergence.
At the beginning of each stage in the curriculum, we add a ∆ α to the previous α (starting from 0).
Then we pass the α into the function ComputeBarycenter
(Algorithm 2) to generate the intermediate task distribution.
Note that when α = 0 and 1 , the generated distribution are the source and target distribution respectively.
In the ComputeBarycenter function, the computation method differs depending on whether the context is discrete or continuous.
After generating the intermediate task distribution, we optimize the agent in this task distribution until the accumulative return G reaches the threshold ¯ G .
Then the curriculum enters the next stage and repeats the process until α = 1 .
In other words, the path of the intermediate task distributions is the OT geodesic on the manifold defined by Wasserstein distance.
4.4 Theoretical Analysis
The proposed algorithm benefits from breaking the difficulty of learning in the target domain into multiple small challenges by designing a sequence of K stages that gradually morph towards the target.
This motivates us to theoretically answer the following questions:
Can we achieve smooth transfer to a new stage based on the optimal policy of the previous stage?
Algorithm for curriculum re I nforc E ment lear N ing via optimal T ransport ( GRADIENT )
Input: Source task distribution µ ( c ) , target task distribution ν ( c ) , interpolation factor ∆ α , distance metric d , reward threshold ¯ G , maximum number of stages K
Initialize the agent policy π for k in 0 , 1 , 2 , ...,
K do α ← min( k ∗ ∆ α, 1) ρ ( c ) ←
ComputeBarycenter ( µ, ν, α, d ) do G ← Optimize π in the task distribution ρ ( c ) and return training reward until G > ¯ G ;
Output: Agent policy π
Algorithm
Input: Source and target µ , ν , interpolation constant α , distance metric d w 1 ← α, w 2 ← 1 − α , ( b 1 , b 2 ) , l ←if C is discrete then λ 1 ← µ ( c ) , λ 2 ← ν ( c )
Cost matrix: C :=
d ( · , · ) : C × C (cid:55)→ R else λ 1 = 1 n s (cid:80) n s i c si , λ 2 = 1 n t (cid:80) n t j c tj
Cost matrix C : C ij =
d ( c si , c tj )
K = exp ( − C /ϵ ) //
Sinkhorn param ϵ while not converge do a 1 ←− ( λ 1 / K b 1 ) , a 2 ←− ( λ 2 / K b 2 ) ρ ←− l ⊙ (cid:81) 2 m =1 ( K ⊤ a m ) w m b 1 ←− (cid:0) ρ/ K ⊤ a 1 (cid:1) , b 2 ←− (cid:0) ρ/ K ⊤ a 2 (cid:1) l ←− (cid:112) l ⊙ ( ρ/ K l )
Output: ρ
Towards this, motivated by the maze navigation example shown in Figure 1, throughout this section, we focus on the collection of CMDPs which obeys the following assumption:
Assumption 4.1 (Consistency and homogeneity)
The initial state distribution is consistent with the associated context distribution, namely p
( s | c )
= 1 ( s = c ) , which indicates S ⊆ C .
In addition, the transition kernels and rewards of the CMDP are homogeneous with respect to the context.
Specifically,
[Equation]
To continue, for k -th stage associated with the context distribution ρ k ∈ ∆( C ) , we denote the optimal policy in k -th stage as π ⋆k :=
arg max π V π ( ρ k ) .
Armed with the above notations and assumptions, we are positioned to introduce our main theorem.
Theorem 4.1 Let m be a problem-dependent positive constant.
Consider any MDP under the assumption 4.1 and the situations further specified in Appendix A.1.
For any two subsequent stages k, k + 1 ∈ [ K ] , the optimal policies π ⋆k , π ⋆k +1 obey
[Equation]
Here, d ⋆ is the distance metric defined on the metric space ( C , d ⋆ ) which obeys ( U ( C ) denote the uniform distribution over C )
[Equation]
The proof can be found in Appendix A.4.
In addition, the following corollary can be directly verified by Theorem 4.1.
Corollary 4.1.1 By constant speed geodesics (7.2 of [51]), the stages generated by GRADIENT obey
[Equation]
Remark 4.1 Theorem 4.1 combined with Corollary 4.1.1 indicates that with tailored stages by
GRADIENT (sufficiently small ∆ α ), it is easy to adapt the learned optimal policy π ⋆k in the current stage to the optimal policy π
⋆k +1 in the next stage defined by ρ k +1 , since the performance gap is small.
In particular, after we obtain the optimal policy π ⋆k associated with the k -th stage in the curriculum and transfer this policy to the next stage, the performance gap between it and the optimal policy in the ( k + 1) -th stage can be controlled by the Wasserstein distance between two subsequent context distributions and therefore a fraction of the Wasserstein distance between the source and the target (which can be seen as a constant).
By properly controlling the Wasserstein distance using GRADIENT, we can achieve gradual domain adaptation from the source to the target task distribution.
Experiments
Evaluation Metric.
To quantify the benefit of CRL, we compare the learning progress of an agent evaluated on the target task distribution 1 .
More concretely, the metric we use is time-to-threshold [6], which shows how fast an agent can learn a policy that achieves a satisfying return the target task if it transfers knowledge.
We also the agent learner
Environments.
The first Maze environment is to investigate the performance of GRADIENT and the effect of ∆ α .
Maze has a discrete context space and therefore the exact context distance metric is available.
The next two environments are to investigate the benefits of GRADIENT when the context space is continuous and we use l 2 distance as a surrogate metric.
In the PointMass environment, the source and the target distribution are both Gaussian distributions, which are in line with the settings of existing algorithms.
In the FetchPush environment, we want to investigate the case where the source and the target are not Gaussian distributions.
We benchmark our algorithm with six baselines including four CRL algorithms:
No-Curriculum : sampling contexts according to the target task distribution.
Domain-Randomization : uniformly sampling the contexts from the context space.
Linear-Interpolation : sampling contexts based on a weighted average between two distributions,
i
e., (1 − α ) µ + αν . 4.
ALP-GMM [13]: the state-of-the-art intrinsically motivated CRL algorithm that maximizes the absolute learning progress of the agent.
Goal-GAN [12]: using GAN to generate contexts of intermediate difficulty and automatically explore the goal space.
SPDL [8]
and WB-SPDL [55]: interpreting the curriculum generation as a variational inference problem to progressively approach the target task distribution.
WB-SPDL extends SPDL by using
Wasserstein distance instead of KL divergence.
5.1 Can GRADIENT Handle Discrete Context and What is the Effect of ∆ α ?
In the Maze task, we fix the layout of the maze and the goal.
The state and action space are discrete.
The action space includes moving in four directions.
The context is the initial position, and therefore the context space overlaps with the state space.
The observation is the flattened value representation of the maze, including the goal, the current position, and the layout.
We visualize the optimal π ∗ -contextual-distance metric in Figure 3a (computed using the optimal policy and normalized to make the maximum value be 1).
The distance metric can be represented by a symmetrical matrix.
We then generate curricula using GRADIENT with ∆ α = 0 .
From Figure 3b, our proposed method outperforms the baselines significantly in terms of the time-to-threshold evaluated in the target distribution.
With a smaller ∆ α , GRADIENT learns slightly slower; nevertheless, all choices achieve good performance.
In this simple environment, DomainRandomization also improves the learning efficiency compared with No-Curriculum (which cannot learn at all).
Interestingly, Domain-Randomization even outperforms the
Linear-Interpolation , which demonstrates the caveat of bad intermediate task distribution.
5.2 How Does GRADIENT Perform When the Distributions Are Gaussian?
In PointMass [8], the agent needs to navigate a pointmass through a wall with a small gap at an off-center position to reach the goal.
The context is a 2-dimension vector representing the width
[Figure]
and position of the gap on the wall.
Following the setting in the original SPDL paper [8], the target distribution is an isotropic Gaussian distribution centered at [2 . 5 , 0 . 5] with a negligible variance [4 × 10 − 3 , 3 .
75 × 10 − 3 ] (which is effectively a point as shown in Figure 5a).
The source distribution is an isotropic Gaussian distribution centered at [0 , 4 . 25] with a variance of [2 , 1 . 875] .
From Figure 4a, we observe that both GRADIENT and SPDL significantly outperform other baselines in terms of the time-to-threshold and achieve the same asymptotic performance.
It is reasonable since these two methods are capable of specifying the target distribution to guide curriculum generation.
We present the curricula visualization of both methods in Figure 5a and show that the intermediate task distributions gradually move from the source to the target distribution.
However, the other three baselines lack the flexibility to leverage the target distribution.
This experiment highlights the importance of being able to specify the target distribution.
5.3 How Does GRADIENT Perform When the Distributions Are NOT Gaussian?
In FetchPush [56], the objective is to use the gripper to push the box to a goal position.
The observation space is a 28-dimension vector, including information about the goal.
The context is a 2-dimension vector representing the goal position on a surface.
The target distribution is a uniform distribution over the circumference of a half-circle (Figure 5b).
The source distribution is a uniform distribution over a square region centered at the box position, excluding the region within a certain radius of the object.
We use this experiment to highlight the importance of the capability to handle arbitrary distributions rather than only the parametric Gaussian distributions.
Since SPDL can only deal with parametric Gaussian distribution, we first fit the target and the source distribution with two
Gaussian distributions and feed them into the baseline algorithms.
[Figure] [Figure]
From Figure 4b, our proposed algorithm outperforms all the baseline methods in terms of the timeto-threshold in the target distribution.
For ALP-GMM and Goal-GAN , since these two methods do not consider the target distribution, they achieve lower performance than the target-aware baseline such as SPDL and GRADIENT .
We further compare these two target-aware methods and find that
GRADIENT outperforms the SPDL in terms of learning efficiency and asymptotic performance.
The reason can be explained by visualizing the context distribution in Figure 5b.
SPDL indeed shows the behavior of moving the context distribution towards the target.
However, since the target distribution cannot be fit with a Gaussian well, many sampled contexts are in fact far away from the target distribution, which may not improve the learning efficiency in the target domain.
Although
Domain-Randomization learns slightly lower than the CRL methods, it achieves even higher asymptotic performance than ALP-GMM and Goal-GAN .
The potential reason could be that it does not bias towards a certain area while ALP-GMM and Goal-GAN are possible.
5.4 Beyond Euclidean Distance Metric: Using Distance Embedding for Continuous Spaces
The red ball represents the the goal position
In practice, it is not always appropriate to use Euclidean distance as the surrogate contextual distance.
Due to the obstacle in the middle.
Though there are many existing OT algorithms based on Euclidean distance metric, solving for free-support Wasserstein barycenters on non-Euclidean space is non-trivial [57, 58].
One solution to by-pass this issue is to embed contexts to a latent space such that the Euclidean interpolation in the latent space approximates the interpolation in the original non-Euclidean space after reconstruction.
For the fixed-support setting, Deep Wasserstein Embedding (DWE) proposed in [59] uses Siamese networks [60] to learn a mapping function such that the Euclidean distance in the latent space approximates the Wasserstein distance in the original space.
Motivated by this idea, we instead develop an embedding mechanism for free supports, i.e., continuous contexts.
We leverage an encoder to embed contexts to latent space in which the Euclidean distance mimics the original distance metric.
We train the encoder and decoder based on a reconstruction loss and a distance embedding loss.
We include the details of distance embedding in Algorithm 4 and a modified version of
GRADIENT in Algorithm 3 in Appendix C.3.
[Figure]
We show an example of using Algorithm 3.
Here we consider a classical U-shaped maze with continuous spaces in Figure 6.
We assume that the source and target distributions are two Gaussian distributions at the two ends of the maze.
Due to the existence of the obstacle in the middle, it is not appropriate to use the l 2 distance as the contextual distance.
In this case, we use d π ( c 1 , c 2 ) :
=
| J π ( c 1 ) − J π ( c 2 ) | as a surrogate metric for the π -contextual-distance, where J π ( c ) represents the expected episodic reward of context c under the policy π .
The interpolation results are shown in Figure 7.
At the first few stages, since the agent does not have a good policy or a good estimate of the episodic reward, the potential geodesic interpolation tries to go through the obstacle in the middle.
Fortunately, since α is small at first, the selected intermediate barycenters (yellow circles) are relatively close to the source distribution, so GRADIENT will not generate unreasonable tasks.
With the learning progressing, the estimates become better and better and therefore produces better interpolation results.
Conclusion and Limitation
In this work, we propose a novel curriculum reinforcement learning framework as an optimal transport problem.
We also develop a distance metric to measure the difference between tasks.
We show that our proposed method is able to achieve smooth transfer between subsequent stages in the generated curriculum.
Finally, the empirical results highlight the three features of our methods: considering metric information, being aware of the target, and being able to handle nonparametric distributions.
For future work, it could be interesting to investigate how to use the “intrinsic curiosity” driven type of method as in [22, 13, 23] to automatically generate the source distribution.
In addition, wider choices of the context distance can be explored to provide more insights into the OT-based CRL methods.
From a theoretical perspective, OT also provides powerful mathematical tools for future researches to investigate rigorously why CRL methods can help reduce the sample complexity.
There are two main limitations of GRADIENT.
First, due to the limitation of OT computation complexity, we only examine the settings that are either discrete or low-dimensional.
Such experiment design choice is common in the related works and baselines [13, 12, 8].
It comes from both the fact of many tasks can be represented by a low-dimensional vector, such as goal positioned or physical parameters of robots.
We believe how to deal with high-dimensional contexts, such as images, is a promising future direction.
Second, our analysis is conducted based on a restrictive subset of contextual MDPs (Assumption 4.1)where the context controls only the initial state distribution.
It would be beneficial to relax this assumption by imposing assumptions on the specific manner in which the context controls the dynamics or reward.
Acknowledgments
We gratefully acknowledge support from the National Science Foundation under grant CAREER
CNS-2047454.
Fei Fang was supported in part by NSF grant IIS-2046640 (CAREER).
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Proof of Theorem 4.
In this section, we shall provide the proof for Theorem 4.1.
The rest of this section is organized as follows.
Appendix A.1 provides additional useful notations and definitions including but not limited to
CMDPs, value function, and distance metrics.
Appendix A.2 introduces further assumptions, and A. introduces the preliminary of optimal transport.
Then, we provide the proof pipeline of Theorem 4. in Appendix A.4 and postpone the auxiliary proof to Appendix A.5.
Finally, we summarize all the useful notations in Table 1.
A.1 Additional notation and definitions used in the proof
Before starting, let’s introduce some additional notations useful throughout the theoretical analysis.
For any discrete set X , we will denote the set of probability measures on X by ∆( X ) .
For any vector x ∈ R SA (resp. x ∈ R S or x ∈ R C ) that constitutes certain values for each state-action pair (resp.
state or context), we use x ( s, a ) (resp.
x ( s ) or x ( c ) ) to denote the entry associated with the ( s, a ) pair (resp state s or context c ).
We denote supp ( ρ ) as the support of any distribution ρ .
Target CMDPs.
Throughout the proof, we shall focus on the set of CMDPs introduced in Assumption 4.1.
Moreover, without loss of generality, let the state space S
= { 1 , 2 , · · · , S } (resp.
action space A = { 1 , 2 , · · · , A } ) to be with finite cardinality S (resp.
A ).
Without loss of generality, let
the immediate reward R
( s, a, s ) ∈ [0 , 1 − γ ] for all ( s, a ) ∈ S × A .
First, we introduce the following occupancy distributions associated with policy π and any CMDP with the initial state distribution c ∼ ρ :
[Equation]
With this in mind, we introduce two important definitions in the proof,:
the covering state-action (resp.
state) space by executing π with the initial state distribution c ∼
ρ defined as
[Equation]
The covering space represent the area that is possible to be visited by the policy π given the initial state distribution.
Value function.
Different from the value function defined by context distribution ρ , with abuse of notation, we denote the value function V s associated with a state s , a policy π and the target CMDPs as
[Equation]
For convenience, for any s ∈ S , we define ρ s ( c ) :
= 1 ( c = s ) .
Recalling that the target MDPs are defined by p 0 ( s | c )
= 1 ( s = c ) , it is easily verified that for any policy π , the value function satisfies the following properties:
[Equation]
Distance metrics.
Furthermore, for the specified target CMDPs, plugging in p
( s | c )
= 1 ( s = c ) , for any two contexts c
i , c j ∈ C , the π -contextual-distance metric defined in context space C obeys
[Equation]
where s i (resp.
s j ) is the corresponding initial state when the context are c i (resp.
c j ).
With this observation, we highlight that for the target CMDPs, the contextual π -bisimulation metric degrades to π -bisimulation d π s ( · , · ) [48] such that:
[Equation] [Equation]
The above result directly leads to that for any ρ, ρ ′ ∈ ∆( C ) ⊆
∆( S ) ,
[Equation] [Equation]
A.2 Additional Assumptions in the proof
Besides the key properties of the target CMDPs introduced in Assumption 4.1, we shall introduce some auxiliary assumptions for convenience as follows:
• Bounded environment.
For the optimal policy π ⋆ and any state s ∈ S
without loss of generality, the minimum and maximum π -bisimulation distance between s and other states satisfies
[Equation]
for some universal positive constant D max
[Equation]
• Deterministic environment.
Without loss of generality, we suppose the transition kernel is deterministic.
Moreover, at any state s ∈ supp ( ρ k +1 ) \ supp ( ρ k ) , there are a fixed portion of actions lead to the decreasing of the distance min s ′ ∈
ρ k d ⋆ s ( s ′ , s ) , i
[Equation]
e., for any transition sample ( s t , a t , s t +1 ) ,
[Equation]
with probability p
[Equation]
• Closeness of the subsequent stages.
Starting from any state s 0 = t ∈ supp ( ρ k +1 ) \ supp ( ρ k ) , if the sample trajectory shall enter C s ( ρ k , π ⋆k ) before arriving at somewhere s obeying d ⋆ s ( s, t ) =
max s ∈S d ⋆ s ( s, t ) .
Then we suppose that the hitting point s is not far from the initial distribution ρ k of the previous stage k , namely
[Equation]
for some universal constant C 1 > 0 .
We would like to note that these additional assumptions are useful for a concise proof to show the main intuition and idea of our proposed algorithm GRADIENT .
These are some general assumptions without specific limitations.
It is interesting to extend our main Theorem 4.1 to more general cases which shall be the further work.
A.3 Preliminaries: optimal transport and random walk
In this subsection, we introduce several key properties/facts of optimal transport, value function, and random walk, which play a crucial role in the proof of Theorem 4.1.
The proofs for this subsection are deferred to Appendix A.5.
Preliminaries of optimal transport.
We first introduce an important fact of the optimization problem in optimal transport.
Suppose there exists two discrete sets X , Y (possibly different).
As it is well-known, for any two distributions µ ∈ ∆( X ) , ν ∈ ∆( Y ) , 1 -Wasserstein distance with any metric d between µ and ν can be expressed as the optimal solution of the following linear programming (Kantorovich duality) problem:
[Equation]
Note that, when X = Y = S , the resulting claim in (19) is equivalent to [48, 61, 62].
Controlling performance gap.
Inspired by the above problem formulation, the difference of the value function conditioned on two different states/state distribution can be controlled in the following lemma:
Lemma A.1 For any s, t ∈ S and any policy π , we have
[Equation]
In addition, for any two context distribution ρ 1 , ρ 2 ∈ ∆( C ) ∈ ∆( S ) , the performance gap associated with any policy π obeys
[Equation]
This fact connects the performance gap in RL and the Wasserstein distance between different state distribution/contextual distribution, taking advantage of the Kantorovich duality form of the optimal transport problem.
Definition and properties of random walk.
Finally, we describe the following useful lemma which is essential in proving the main part of Theorem 4.1 when the starting state is not in the covering area of the optimal policy of the previous training stage
Lemma A.
The process { S n : n ≥ 1 } is called a random walk if { X i : i ≥ 1 } are iid Bernoulli distribution with p < 1 and S n = (cid:80) ni
=1 X i .
Starting from 0 , the expectation of the stopping time N of hitting any d > 0 or − D max and the probability p α of hitting d > 0 obeys:
[Equation]
A.4 Proof of Theorem 4.
With above preliminaries in hand, we are ready to embark on the proof for Theorem 4.1, which is divided into multiple steps as follows.
Step 1: decomposing the performance gap of interest.
To begin with, we decompose the term of interest as follows:
[Equation]
where (i) holds by the equivalence verified in (12b), (ii) follows from the value function of the optimal policy π ⋆ (covering the entire state-action space) is the same as that of π ⋆k +1 in the region
C ( ρ k +1 , π ⋆k +1 ) , and (iii) comes from applying Lemma A.1 and the fact that π ⋆ achieves the same value function as π ⋆k in the region C ( ρ k , π ⋆k ) .
Step 2: controlling the second term in (23) in two cases.
To proceed, it is observed that
[Equation]
can directly leads to that ( u, v )
= (cid:0) V
π
⋆k s ,
V π ⋆k s (cid:1) is a feasible solution to the Wasserstein distance
W d ⋆ ( ρ k , ρ k +1 ) dual formulation in (19), which yields
[Equation]
As a result, we turn to show (24) instead of controlling V
π ⋆k s ( ρ k )
−
V π ⋆k s ( ρ k +1 ) .
Towards this, we start from considering V
π ⋆k s ( s )
− V
π ⋆k s ( t ) for any s ∈ ρ k and t ∈ ρ k +1 in two different cases: (i) t ∈ C s ( ρ k , π ⋆k ) ; (ii) otherwise.
In the first case when t ∈ C s ( ρ k , π ⋆k ) , since π ⋆k is the same as the optimal policy π ⋆ in C s ( ρ k , π ⋆k ) , invoking Lemma A.1 directly yields
[Equation]
Step 3: focusing on case (ii).
Then, we shall focus on the other case when
[Equation]
Invoking Lemma A.2, without loss of generality, since in the stage k , we can’t visit the region outside of C s ( ρ k , π ⋆k ) , we can define π ⋆k ( · | s ) to be uniformly random in the unseen region, i.e.,
[Equation]
Then we define p g as the probability of the benign events B when we can hit some point s b at the boundary of the region C s ( ρ k , π ⋆k ) visited by the previous stage k before arriving at some limit point s limit
[Equation]
s
[Equation]
t, s ′ ) .
Invoking the assumption of the target CMDPs in (17), combind with the policy defined in (28), we observe that applying Lemma A.2 leads to
[Equation]
where the last inequality holds by the assumption in (16).
To continue, we express the value function at state t as
[Equation]
Similarly, we have
[Equation]
We shall control the performance gap in two cases separately:
• When V π ⋆k s ( s b ) ≥
V π
⋆k s ( s N ) .
we have V π
⋆k s (
[Equation]
• When V π ⋆k s ( s b ) <
V π ⋆k s ( s N ) , we have
V
[Equation]
Summing up the above two cases yields, for all s ∈ C ( ρ k , π ⋆k ) , t ∈ supp ( ρ k +1 ) \ C ( ρ k , π ⋆k ) ,
[Equation]
where (i) follows from Lemma A.
[Equation]
and the last inequality holds by choosing
[Equation]
as a large enough problem-dependent parameter
[Equation]
Step 4: summing up the results.
Combining the results (26) and (34) in two cases, we directly arrive at, for all s ∈ C ( ρ k , π ⋆k ) , t ∈ supp ( ρ k +1 ) , the following fact is satisfied
[Equation]
We observe that ( u, v ) =
(cid:0) 1 c model
V π
⋆k s , 1 c model
V π
⋆k s (cid:1) is a feasible solution to the Wasserstein distance
W d ⋆ s ( ρ k , ρ k +1 ) dual formulation in (19), we achieve
[Equation]
Finally, plugging (36) into (23) complete the proof by
[Equation]
A.5 Proof of auxiliary lemmas
Proof of Lemma A.1.
As it is well known that V π s ( s ) is the fixed point of the Bellman operator
T ( V, π ) :
= E
a
[Equation]
·| s )
[ R
[Equation]
s, a, s )
+ γ
[Equation]
s ′ P
( s ′ | s, a, s ) V ] .
In addition, d π s is the fixed point of the operator [48]
[Equation]
Armed with above facts, initializing V π s , 0 = 0 and d π s , 0 = 0 , the update rules of V π s ,n +1 and d π s ,n +1 at the ( n + 1) -th iteration are defined as
[Equation]
[Equation]
With this in mind, for any s, t ∈ S
[Equation]
we arrive at,
[Equation]
40) where (i) holds by ( u, v )
= (cid:0) V π s ,n , V π s ,n (cid:1) is a feasible solution to the Wasserstein distance
W d π s ,n (cid:0) E a ∼ π ( ·| s )
[ P ( s ′ | s, a, s )] , E a ∼ π ( ·| t )
[ P ( s ′ | t, a, t )] (cid:1) dual formulation in (19).
As a result, we have
[Equation]
which directly yields
[Equation]
The above fact indicates that for any two distribution ρ, ρ ′ ∈ ∆( S ) ,
[Equation]
where the penultimate inequality arises from ( u, v ) = (cid:0) V π s , V π s (cid:1) is a feasible solution to the Wasserstein distance W d π s ( ρ, ρ ′ ) dual formulation in (19), and the last equality holds by (15).
Proof of Lemma A.2 Before starting, we denote p α as the probability that stopping at d and N as the stopping time.
Then we consider the terms of interest in two cases.
In the first case when p = 12 , invoking the facts in [63] directly leads to
[Equation]
So the remainder of the proof will focus on the other case when p ̸ = 12 .
Following the proof in [63], by basic calculus, it is easily verified that the following two centered terms are all martingales
[Equation]
In continue
[Equation]
it can also be verified that
[Equation]
which leads to
[Equation]
Observing that
[Equation]
yields
[Equation]
and then implies
[Equation]
Plugging in (50) into (47) leads to
[Equation]
which follows from (16) in Appendix A.2.
Finally
[Table]
summing up the two cases
[Table]
we arrive at
[Equation]
A.6 Table of notations
We summarize useful notations in the proof here.
[Table]
B Experiment Details
In this section
[Table]
we discuss details that could not be included in the main paper due to space limitations.
This includes hyperaparameters of the algorithms, additional details about the environment
[Table]
and visualizations that assist the qualitative analysis.
The experiments were conducted on a desktop computer with an Intel Core i7-8700K CPU @
3.70GHz 12-Core Processor, an Nvidia RTX 2080Ti graphics card and 64GB of RAM.
B.1 Maze
We show the observation, action and context space in Table 2.
The state, action and context space are discrete.
For this toy-like example, we fix the maze layout throughout.
The context is the initial position.
The observation is the flattened value representation of the maze, including the goal, the current position of the agent, and the layout.
For this task, we use the PPO implementation in the StableBaseline3 library [54].
We list the hyperparameters in the Table 3 (set to the default value if not mentioned in the table).
We conduct hyperparameter grid search mainly on the following hyperparameters: ( learning _ rate, ent _ coef, n _ steps ) ∈ { 0 . 0001 , 0 . 0003 , 0 .
001 } × { 0 . 001 , 0 . 01 , 0 . 1 } × { 10 , 100 , 200 }
We visualize the intermediate task distributions generated by GRADIENT in Figure 8 with a ∆ α = 0 .
For ∆ α = 0 . 1 and 0.2 are in the same figure as well (just with different gaps between two consecutive α ).
The π -contextual-distance is computed using the A* path finding algorithm.
Although
[Table] [Table]
it is not always realistic to have access to the optimal policy, this toy example is just to show the effectiveness of an appropriate contextual distance metric.
More choices of the contextual distance could be explored in the future work.
B.2 PointMass
In PointMass, the agent needs to navigate a pointmass through a wall with a small gap at an off-center position to reach the goal.
The context is a 2-dimension vector representing the width and position of the gap on the wall.
Following the setting in the original SPDL paper [8], the target distribution is an isotropic Gaussian distribution centered at [2 . 5 , 0 . 5] with a negligible variance [4 × 10 − 3 , 3 .
75 × 10 − 3 ]
[Figure]
(which is effectively a point as shown in Figure 5a).
The source distribution is an isotropic Gaussian distribution centered at [0 , 4 . 25] with a variance of [2 , 1 . 875] .
We show the observation, action and context space in Table 4.
For the baseline implantation, we use and modify the code base provided in [8].
We also use the best hyperparameters found by [8].
For ALP-GMM, they conduct hyperparameter grid search over ( p RAND , n ROLLOUT , s BUFFER ) ∈ { 0 . 1 , 0 .
2 , 0 . 3 } × { 50 , 100 , 200 } × { 500 , 1000 , 2000 } .
For
Goal-GAN, they conduct grid search over ( δ NOISE , n ROLLOUT , p SUCCESS ) ∈ { 0 . 025 , 0 . 05 , 0 .
1 } × { 50 , 100 , 200 } × { 0 . 1 , 0 . 2 , 0 . 3 } .
Due to unknown issue, the scale of episodic reward we get (about 0-100) is different from what is shown in [8] (0-10), nevertheless the trend of the training curves is very similar.
We visualize the context distribution of ALP-GMM and Goal-GAN in Figure 9.
From this figure, we can reason why they underform GRADIENT and SPDL: since they cannot specify the target distribution, it is very difficult for them to learn to navigate through a narrow door at a specific position.
[Table]
Context Space gate position [ − 4 , 4] gate
[Table]
B.3 FetchPush
In FetchPush [56], the objective is to use the gripper to push the box to a goal position.
The observation space is a 28-dimension vector, including information about the goal.
The context is a 2-dimension vector representing the goal position on a surface.
The target distribution is a uniform distribution over
[Table]
the circumference of a half-circle (Figure 5b).
The source distribution is a uniform distribution over a square region centered at the box position, excluding the region within a certain radius of the object.
We use this experiment to highlight the importance of the capability to handle arbitrary distributions rather than only the parametric Gaussian distributions.
Since SPDL can only deal with parametric
Gaussian distribution, we first fit the target and the source distribution with two Gaussian distributions and feed them into the baseline algorithms.
The source Gaussian is [1 . 14655655 , 0 . 74819359] with a variance array of [[0 . 0141083 , 0 . 00055327] , [0 . 00055327 , 0 . 0149638]]
[Table]
and the target Gaussian is [1 . 33561676 , 0 . 74819359] with a variance array of [[8 .
e − 03 , − 1 .
34467507 e − 18] , [ − 1 .
34467507 e − 18 , 4 .
e − 02]] .
We visualize the context distribution of ALP-GMM and Goal-GAN in Figure
C Algorithm Details C.1 Computation of π -contextual-distance from Samples
Computing contextual π -bisimulation from samples.
We start from estimating the contextual πbisimulation metric .
We first define an iterative procedure for computing it from a data buffer D .
Let
d πc i ,c j , 0 ≡ 0 be the everywhere-zero metric.
At step n , let τ = { s i , s j , a i , a j , s ′ i , s ′ j , R i , R j ,
c i , c j } ∈
T be a sample from D , define d π c i ,c j ,n as:
[Equation]
[Table]
Learning an approximation for contextual π -bisimulation
[Table]
Let ϕ : S × C (cid:55)→ R k be a kdimensional represetation of the state and context space and let ψ ω :
R 2 k
(cid:55)→ R be a neural network parameterized by ω
[Table]
that receives a concatenation of two state representations such that ψ ω ( ϕ ( s i , c i ) , ϕ ( s j , c j ))
≈
d πc i ,c j
( s i , s j ) .
Following the conventional practice in [3]
[Table]
we use online parameters ω and target parameters ω −
[Table]
where the online parameters ω are updated at each iteration while the target parameters are updated very multiple iterations.
We define the target objective
T ω − ( τ ) as:
[Equation]
Then we can define our loss L ω − ( τ ) as E ( T ω − ( τ ) − ψ ω
( ϕ (( s i , c i ) , ϕ ( s j , c j )) 2 .
This loss is defined for a single pair of transitions from the buffer, but the mini-batch version will be very similar.
The computation can be carried out efficiently on GPUs.
To obtain π -contextual-distance, given a set of s i ∈ s sampled from p 0 ( · | c i ) and s j ∈ t from p 0 ( · | c j ) ,
[Equation]
Note that with approximation, these distance metric are not guaranteed to be a proper pseudometric: it can violate the identity of indiscernibles, symmetry, and subadditivity conditions.
C.2 Exact Computation of π -contextual-distance for Maze
For the maze, we compute the exact π -contextual-distance using dynamic programming.
We iteratively compute the π -contextual-distance (represented by a matrix) until the maximum difference in metric estimate between successive iterations is smaller than a tolerance.
Due the Assumption 4.1, the computation of π -contextual-distance is largely simplified (since the context space is essentially overlapped with the state space).
We present the pseudocode as follows:
def compute_exact_metric(tolerance=0.001): """Compute the exact pi-contextual-distance metric up to the specified tolerance.
Args:
tolerance: float, maximum difference in metric estimate between successive iterations.
Once this threshold is reached,
[Table]
gamma
[Table]
policy = compute_optimal_policy() n = get_context_size()
self.contextual_distance = np.zeros((n, n))
metric_difference
= tolerance * 2.
i =
exact_metric_differences
=
[] while metric_difference > tolerance: new_metric = np.zeros((n, n)) for s1 in range(n): for s2 in range(n): # get the action from policy action1 = policy[s1] action2 = policy[s2] # environment step next_state1, rew1 = step(s1, action1) next_state2, rew2 = step(s2, action2)
# apply iterative operator new_metric[s1, s2]
=
( abs(rew1 - rew2) + gamma*self.contextual_distance[next_state1, next_state2])
metric_difference
= np.max(abs(new_metric - self.contextual_distance))
[Figure]
exact_metric_differences.append(metric_difference)
self.contextual_distance=
np.copy(new_metric) i+= # adding very small offset to the off-diagonal entries # to avoid computation issue offset = 0.1*np.ones(self.contextual_distance.shape) np.fill_diagonal(offset, 0) self.contextual_distance += offset # normalize to maximum value ofself.contextual_distance = self.contextual_distance/np.max(self.contextual_distance) return self.contextual_distance
C.3 Learning Embeddings to Encode Non-Euclidean Distance
Due to the computation complexity of the free-support Wasserstein barycenters, it is generally difficult to compute them, especially for non-euclidean cost metric.
There are some attempts to use neural networks and stochastic gradient descent to solve for the barycenters approximately [58].
Another possible route is to learn embeddings such that the euclidean interpolation in the latent space assembles the interpolation in the non-euclidean original space.
In the fixed-support setting,
Deep Wasserstein Embedding (DWE) [59] uses siamese networks to learn an latent space where the euclidean distance approximates the Wasserstein distance in the original space.
We could adopt a similar method but for the free-support, i.e., continuous context space.
Let the encoder be ϵ ( · ) and decoder be δ ( · ) , given pairs of contexts { c i 1 , c i 2 } i =1 ,...,N and their contextual distance { d ( c i 1 , c i 2 ) } i =1 ,...,N , the global objective funtion is to minimize
[Equation]
The first term is the loss for distance embedding, and the second term is for reconstruction.
Other regularization could be added to improve the robustness.
After obtaining the trained encoder and decoder, we can first encode the source and target contexts to the latent space, compute the Wasserstein barycenter in the latent space and finally decode the latent barycenter back to the original context space.
The main algorithm GRADIENT with distance embedding is shown in Algorithm 3.
Algorithm
Input: Source task distribution µ ( c ) , target task distribution ν ( c ) , interpolation factor ∆ α , reward threshold ¯ G
Initialize the agent policy π α ←for k in 0 , 1 , 2 , ..., K do if k ==
then ρ ( c ) ← µ ( c ) else ρ ( c ) ← δ ( ComputeBarycenter ( ϵ ( µ ) , ϵ ( ν ) , α, l 2 )) { c i , R i } i =1 ,...,M ← Optimize π in the task distribution ρ ( c ) until G > ¯ G , and return recent
M sampled context and corresponding episodic rewards
Estimate J π ( c ) from { c i , R i } i =1 ,...,M using Gaussian Process //
Use the absolute difference between episodic reward to define the distance metric as an example
Define d π ( c i , c j ) :=
| J π ( c i ) − J π
( c j ) |
Encoder ϵ ,
Decoder δ ← EmbedDistanceMetric ( d π ) //
Algorithmα ← α + ∆ α µ ( c ) ← ρ ( c )
Output: Agent policy π
Input: contextual distance metric d
Initialize encoder ϵ and decoder δ
Uniformly sample pairs of contexts from the context range { c i 1 , c i 2 } i =1 ,...,N //
This step is cheap since there is no interaction with environment required.
Compute { d ( c i 1 , c i 2 ) } i =1 ,...,N
Train encoder and decoder by minimizing (56)
Output: Encoder ϵ , Decoder δ
We show an example of using Algorithm 3.
Here we consider a classical U-shaped maze with continuous state, action and context space.
We assume that the source and target distributions are two Gaussian distributions at the two ends of the maze.
Due to the existence of the obstacle in the middle, it is not appropriate to use the l 2 distance as the contextual distance.
In this case, we use d π ( c 1 , c 2 ) :
=
| J π ( c 1 ) − J π ( c 2 ) | .
The interpolation results are shown in Figure 11.
At the first a few stages, since the agent does not have a good policy as well as a good estimate of the episodic reward, the interpolation results are not very good.
However, since α is small, the barycenters are relatively close to the source distribution, so
GRADIENT does not generate too unreasonable tasks.
With the learning progressing, the estimates become better and better and therefore produces much better interpolation results.
[Figure]