E FFICIENT F ULLY -O FFLINE M ETA -R EINFORCEMENT L EARNING VIA D ISTANCE M ETRIC L EARNING AND B EHAVIOR R EGULARIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
We study the ofﬂine meta-reinforcement learning (OMRL) problem, a paradigm which enables reinforcement learning (RL) algorithms to quickly adapt to unseen tasks without any interactions with the environments, making RL truly practical in many real-world applications.
This problem is still not fully understood, for which two major challenges need to be addressed.
First, ofﬂine RL usually suffers from bootstrapping errors of out-of-distribution state-actions which leads to divergence of value functions.
Second, meta-RL requires efﬁcient and robust task inference learned jointly with control policy.
In this work, we enforce behavior regularization on learned policy as a general approach to ofﬂine RL, combined with a deterministic context encoder for efﬁcient task inference.
We propose a novel negative-power distance metric on bounded context embedding space, whose gradients propagation is detached from the Bellman backup.
We provide analysis and insight showing that some simple design choices can yield substantial improvements over recent approaches involving meta-RL and distance metric learning.
To the best of our knowledge, our method is the ﬁrst model-free and end-to-end
OMRL algorithm, which is computationally efﬁcient and demonstrated to outperform prior algorithms on several meta-RL benchmarks.
I NTRODUCTION
Applications of reinforcement learning (RL) in real-world problems have been proven successful in many domains such as games (Silver et al., 2017; Vinyals et al., 2019;
Ye et al., 2020) and robot control (Johannink et al., 2019).
However, the implementations so far usually rely on interactions with either real or simulated environments.
In other areas like healthcare (Gottesman et al., 2019), autonomous driving (Shalev-Shwartz et al., 2016) and controlled-environment agriculture (Binas et al., 2019) where RL shows promise conceptually or in theory, exploration in real environments is evidently risky, and building a high-ﬁdelity simulator can be costly.
Therefore a key step towards more practical RL algorithms is the ability to learn from ﬁxed data (batch).
Such paradigm, termed ”ofﬂine RL” or ”batch RL”, would enable better generalization by incorporating diverse prior experience.
Moreover, by leveraging and reusing previously collected data, off-policy algorithms such as
SAC (Haarnoja et al., 2018) has been shown to achieve far better sample efﬁciency than on-policy methods.
The same applies to ofﬂine RL algorithms since they are by nature off-policy.
The aforementioned design principles motivated a surge of recent works on ofﬂine/batch RL (Fujimoto et al., 2019; Kumar et al., 2019; Wu et al., 2019; Siegel et al., 2020).
These papers propose remedies by regularizing the trained policy to stay close to the training dataset of logged transitions, in order to mitigate the effect of bootstrapping error (Kumar et al., 2019), where errors in evaluation of out-of-distribution state-action pairs are never corrected and hence easily diverge due to inability to collect new data samples for feedback.
There exist claims that ofﬂine RL can be implemented successfully without explicit correction for distribution mismatch given sufﬁciently large and diverse training data (Agarwal et al., 2020).
However, we ﬁnd such assumption unrealistic in many practices as well as in our experiments.
In this paper, to tackle the out-of-distribution problem in ofﬂine RL in general, we follow the paradigm of behavior regularization in Wu et al.
For design of practical RL, besides the ability to learn without exploration, it’s also ideal to have an algorithm that can generalize to a variety of scenarios.
To solve real-world challenges in multi-task setting, such as treating different diseases, driving under various road conditions or growing diverse crops in autonomous greenhouses, a practical/data-efﬁcient agent is expected to quickly transfer and adapt to unseen tasks, especially when the tasks share common structures.
Meta-learning methods (Vilalta & Drissi, 2002; Thrun & Pratt, 2012;
Finn et al., 2017) address this problem by learning an inductive bias from experience collected across a distribution of tasks, which can be naturally extended to reinforcement learning context.
Under the umbrella of this so-called meta-RL, almost all current methods require on-policy data during either both meta-training and testing (Finn et al., 2017;
Wang et al., 2016; Duan et al., 2016) or at least online exploration for task inference (Rakelly et al., 2019).
An efﬁcient and robust method which incorporates both fully-ofﬂine learning and meta-learning in RL, despite few attempts (Li et al., 2019b; Dorfman & Tamar, 2020), has not been fully developed and validated.
In this paper, under the ﬁrst principle of maximizing practicality of RL algorithm, we propose an efﬁcient method that integrates task inference with RL algorithms in a fully-ofﬂine fashion.
Our fully-ofﬂine context-based actor-critic meta-RL algorithm, or FOCAL, achieves excellent sample efﬁciency and fast adaptation with only limited logged experience, on a range of deterministic continuous control meta-environments.
The primary contribution of this work is designing the ﬁrst end-to-end and model-free ofﬂine meta-RL algorithm which is computationally efﬁcient and effective without any prior knowledge of task identity or reward/transition functions.
To achieve efﬁcient task inference, we propose an inverse-power loss for effective learning and clustering of latent variables, in analogy with coulomb potential in electromagnetism, which is also unseen in previous work.
Theoretical analyses and insights are provided to shed light on the speciﬁc design choices customized for the OMRL problem.
R ELATED W ORK Meta-RL
Our work FOCAL builds upon the meta-learning (Vilalta & Drissi, 2002; Thrun &
Pratt, 2012) framework in the context of reinforcement learning.
Among all paradigms of metaRL, this paper is most related to context-based and metric-based approaches.
Context-based metaRL employs models with memory such as recurrent (Duan et al., 2016; Wang et al., 2016; Fakoor et al., 2019), recursive (Mishra et al., 2017) or probabilistic (Rakelly et al., 2019)
structures to achieve fast adaptation by aggregating experience into a latent representation on which the policy is conditioned.
The design of recent methods usually leverages the temporal or Markovian properties of RL problems.
Metric-based meta-RL focuses on learning effective task representations to facilitate task inference and conditioned control policies, by employing techniques such as distance metric learning (Yang &
Jin, 2006).
Koch et al.
(2015) proposed the ﬁrst metric-based meta-algorithm for few-shot learning, in which a Siamese network (Chopra et al., 2005) is trained with triplet loss to compare the similarity between a query and supports in the embedding space.
Many metric-based meta-RL algorithms extend these works.
(Snell et al., 2017; Sung et al., 2018;
Li et al., 2019a).
Among all aforementioned meta-learning approaches, this paper is most related to context-based
PEARL algorithm (Rakelly et al., 2019) and metric-based prototypical networks (Snell et al., 2017).
PEARL achieves SOTA performance for off-policy meta-RL by introducing a probabilistic permutation-invariant context encoder, along with a design which disentangles task inference and control by different sampling strategies.
However, it requires exploration during meta-testing.
The prototypical networks employ similar design of context encoder as well as an Euclidean distance metric on deterministic embedding space, but tackles meta-learning of classiﬁcation tasks with squared distance loss as opposed to the inverse-power loss in FOCAL to tackle the more complex
OMRL problem.
Ofﬂine/Batch RL
To address the bootstrapping error (Kumar et al., 2019) problem of ofﬂine RL, this paper adopts behavior regularization directly from Wu et al.
(2019), which provides a relatively uniﬁed framework of several recent ofﬂine or off-policy RL methods (Haarnoja et al., 2018;
Fuji- moto et al., 2019; Kumar et al., 2019).
It incorporates a divergence function between distributions over actions in the actor-critic objectives.
As with SAC (Haarnoja et al., 2018), one limitation of the algorithm is its sensitivity to reward scale and regularization strength.
In our experiments, we indeed observed wide spread of optimal hyper-parameters across different meta-RL environments, shown in Table 3.
Ofﬂine meta-RL
To the best of our knowledge, despite attracting more and more attention, the ofﬂine meta-RL problem is still under-developed.
We are aware of a few papers that tackle the same problem from different angles (Li et al., 2019b; Dorfman & Tamar, 2020).
Li et al.
(2019b) focuses on a speciﬁc scenario where biased datasets cause the task inference module to spuriously correlate only state-action pairs to the task identity, ignoring the reward information.
This problem occurs when datasets of different tasks do not have signiﬁcant overlap in their state-action visitation frequencies, and is exacerbated by sparse rewards.
Even though we observed similar effects during our experiment, in this paper, we do not restrict our algorithm design to such assumptions and instead try to robustify the algorithm in a more general setting, showing experimental results on various training and testing datasets distributions (Fig 7).
Another key difference is that our method is end-to-end and model free, whereas Li et al.
(2019b) requires training of ofﬂine BCQ (Fujimoto et al., 2019) and reward function models for each task, which are computationally demanding.
They also adopted the original PEARL strategy of probabilistic context encoder and couples learning of task inference and control by Bellman gradients, for which we chose differently for provably better performance and efﬁciency.
Dorfman & Tamar (2020) on the other hand, completely circumvents the out-of-distribution problem of ofﬂine RL by assuming sufﬁcient data samples.
They also formulate the OMRL as a Bayesian RL (Ghavamzadeh et al., 2016) problem and employs a probabilistic approach.
Therefore we consider their methodology tangential to ours.
P RELIMINARIES 3. N OTATIONS AND P ROBLEM S TATEMENT
We consider fully-observed Markov Decision Process (MDP) (Puterman, 2014) in deterministic environments such as MuJoCo (Todorov et al., 2012).
An MDP can be modeled as M = ( S , A , P, R, ρ 0 , γ ) with state space S , action space A , transition function P ( s (cid:48) | s, a ) , bounded reward function R ( s, a ) , initial state distribution ρ 0 ( s ) and discount factor γ ∈ (0 , 1) .
The goal is to ﬁnd a policy π ( a | s ) to maximize the cumulative discounted reward starting from any state.
We introduce the notion of multi-step state marginal of policy π as µ tπ ( s ) , which denotes the distribution over state space after rolling out π for t steps starting from state s .
The notation R π ( s ) denotes the expected reward at state s when following policy π : R π ( s ) =
E a ∼ π
[ R ( s, a )] .
The state-value function (a.k.a. value function) and action-value function (a.k.a Q-function) are therefore
V π ( s ) = ∞ (cid:88)
[Equation]
=
[Equation] [Equation]
[Equation]
Q-learning algorithms are implemented by iterating the Bellman optimality operator B , deﬁned as:
[Equation]
(3) when the state space is large/continuous, ˆ Q is used as a hypothesis from the set of function approximators (e.g. neural networks).
In the ofﬂine context of this work, given a distribution of tasks p ( T ) where every task is an
MDP, we study off-policy meta-learning from collections of static datasets of transitions D i
= { ( s i,t , a i,t , s (cid:48)
i,t , r i,t )
| t
= 1 , ..., N } generated by a set of behavior policies { β i ( a | s ) } associated with each task index i .
A key underlying assumption of meta-learning is that the tasks share some common structures.
Accordingly, by deﬁnition of MDP, in this paper we restrict our attention on tasks with shared state and action space, but differ in transition and reward functions.
We deﬁne the meta-optimization objective as
[Equation]
where L T i ( θ ) is the objective evaluated on transition samples drawn from task T i .
A common choice of p ( T ) is the uniform distribution on the set of given tasks {T i | i
= 1 , ..., n } .
In this case, the meta-training procedure turns into minimizing the average losses across all training tasks
[Equation]
3. B EHAVIOR R EGULARIZED A CTOR C RITIC
Similar to SAC, to constrain the out-of-distribution error in ofﬂine RL, for each individual task T i , behavior regularization (Wu et al., 2019) introduces a divergence measure between the learner π θ and the behavior policy π b in both value and target Q-functions.
For simplicity, we ignore task index in this section:
[Equation]
where ¯ Q denotes a target Q-function without gradients and ˆ D denotes a sample-based estimate of the divergence function D .
In actor-critic framework, the loss functions of Q-value and policy learning are given by, respectively,
[Equation]
[Table]
3. C ONTEXT -B ASED META -RL
Context-based meta-RL algorithms aggregate context information, typically in form of task-speciﬁc transitions, into a latent space Z .
Adaptation at test time can be viewed as a special case of RL on partially-observed MDP (Kaelbling et al., 1998) in which a latent representation z as the unobserved part of the state needs to be inferred.
Once given complete information of z and s combined as the full state, the reinforcement learning of policy π θ ( s, z ) becomes RL on regular MDP, and properties of regular RL such as the existence of optimal policy and value functions hold naturally.
We therefore formulate the context-based meta-RL problem as solving a task-augmented MDP (TA-MDP).
The formal deﬁnitions are provided in Appendix B.
M ETHOD
Based on our formulation of context-based meta-RL problem, FOCAL ﬁrst learns an effective representation of meta-training tasks on latent space Z , then solves the ofﬂine RL problem on TA-MDP with behavior regularized actor critic method.
We illustrate our training procedure in Figure 1 and describe the detailed algorithm in Appendix A.
Our method consists of three key design choices: deterministic context encoder, distance metric learning on latent space as well as decoupled training of task inference and control.
4. D ETERMINISTIC C ONTEXT E NCODER
Similar to Rakelly et al.
(2019), we introduce an inference network q φ ( z | c ) , parameterized by φ , to encode task identity from context c .
In terms of the context encoder design, recent meta-RL methods either employ recurrent neural networks (Duan et al., 2016; Wang et al., 2016) to capture the temporal property of MDP, or use probabilistic models (Rakelly et al., 2019) for reasoning of uncertainties.
These design choices are proven effective in on-policy or partially-ofﬂine off-policy algorithms.
However, since our approach aims to address the fully-ofﬂine meta-RL problem, we argue that a deterministic context encoder works better in this scenario, given a few assumptions.
First, we consider only deterministic MDP in this paper, where the transition function is a Dirac delta distribution.
We assume that all meta-learning tasks in this paper are deterministic MDPs, which is satisﬁed by common RL benchmarks such as MuJoCo.
The formal deﬁnitions are detailed in Appendix B. Second, we assume all tasks share the same state and action space, while each is characterized by a unique set of transition and reward functions.
Mathematically, this means there exists an injective function f : T → P × R , where P and R are functional spaces of transition probability P : S ×
A × S → { 0 , 1 } and bounded reward R : S ×
A → R respectively.
A stronger condition of this injective property is that for every task T ∼ p ( T ) and any state-action pair ( s, a ) , the corresponding transition and reward are unique across all tasks, which brings the Task-Transition
Correspondence assumption:
Assumption 1 (Task-Transition Correspondence)
We consider meta-RL with a task distribution p ( T ) to satisfy task-transition correspondence if and only if ∀T
[Equation] [Equation]
Under the deterministic MDP assumption, the transition probability function P ( ·| s, a ) is associated with the transition map t : S ×
A → S (Deﬁnition 3).
The task-transition correspondence suggests that, given the action-state pair ( s, a ) and task T , there exists a unique transition-reward pair ( s (cid:48) , r ) .
Based on these assumptions, one can deﬁne a task-speciﬁc map f T : S ×
A → S × R on the set of transitions D T :
[Equation]
Recall that all tasks deﬁned in this paper share the same state-action space, hence { f T |T ∼ p ( T ) } forms a function family deﬁned on the transition space S ×
A × S × R , which is also by deﬁnition the context space C .
This lends a new interpretation that as a task inference module, the context encoder q φ ( z | c ) enforces an embedding of the task-speciﬁc map f T on the latent space Z , i.e. q φ : S ×
A × S × R → Z .
Following Assumption 1, every transition { s i , a i , s (cid:48)
i , r i } corresponds to a unique task T i , which means in principle, task identity can be inferred from any single transition tuple.
This implies the context encoder should be permutation-invariant and deterministic, since the embedding of context does not depend on the order of the transitions nor involve any uncertainty.
This observation is crucial since it provides theoretical basis for few-shot learning (Snell et al., 2017;
Sung et al., 2018) in our settings.
In particular, when learning in a fully-ofﬂine fashion, any metaRL algorithm at test-time cannot perform adaptation by exploration.
The theoretical guarantee that a few randomly-chosen transitions can enable effective task inference ensures that FOCAL is feasible and efﬁcient.
4. D ISTANCE M ETRIC L EARNING (DML) OF L ATENT V ARIABLES
In light of our analysis on the context encoder design, the goal of task inference is to learn a robust and effective representation of context for better discrimination of task identities.
Unlike PEARL, which requires Bellman gradients to train the inference network, our main insight is to disentangle
[Figure]
the learning of context encoder from the learning of control policy.
As explained in previous reasoning about the deterministic encoder, the latent variable is a representation of the task properties involving only transition probability and reward, which in principle should be completely captured by the transition datasets.
Given continuous neural networks as function approximators, the learned value functions conditioned on latent variable z cannot account for the difference between tasks if the corresponding embedding vectors are too close (Appendix C).
Therefore for implementation, we formulate the latent variable learning problem as obtaining the map q φ : S ×
A × S × R → Z of transition data D i = { ( s i,t , a i,t , s (cid:48)
i,t , r i,t )
| t
= 1 , ..., N } that clusters similar data points (same task) while pushes away dissimilar data points (different tasks) on the embedding space Z , which is essentially distance metric learning (DML) (Sohn, 2016).
A common loss function in DML is contrasitive loss (Chopra et al., 2005;
Hadsell et al., 2006).
Given input data x i , x j ∈ X and label y ∈ { 1 , ..., L } , it is written as
[Equation]
where m is a constant parameter, q i is the embedding vector of x i .
For data point of different tasks/labels, contrastive loss rewards the distance between their embedding vectors by L 2 norm, which is weak when the distance is small, as in the case when z is normalized and q φ is randomly initialized.
Empirically, we observe that objectives with positive powers of distance lead to degenerate representation of tasks, forming clusters that contain embedding vectors of multiple tasks (Figure 2(a)).
Theoretically, this is due to the fact that an accumulative L 2 loss of distance between data points is proportional to the dataset variance, which may lead to degenerate distribution such as Bernoulli distribution.
This is proven in Appendix B. To build robust and efﬁcient task inference module, it’s crucial to ensure every task embedding cluster to be separated from each other.
We therefore introduce a negative-power variant of contrastive loss as follows:
[Equation]
where (cid:15) > 0 is a small coefﬁcient added to avoid division by zero, the power n can be any nonnegative number.
We experimented with 1 and 2 in this paper and compare them with the classical
L 1 , L 2 metrics in 5.2.2.
E XPERIMENTS
In our experiments, we assess the performance of FOCAL by comparing it with several baseline algorithms on meta-RL benchmarks.
Speciﬁc design choices are examined through three ablations.
All return curves are averaged over 3 random seeds.
Supplementary experiments are provided in
Appendix E.
S AMPLE E FFICIENCY AND P ERFORMANCE
We evaluate FOCAL on four continuous control meta-environments of robotic locomotion, three of which are simulated via the MuJoCo simulator (Todorov et al., 2012), with exception of a 2D navigation problem with sparse reward called Sparse-Point-Robot.
All environments require adaptation across reward functions.
Additionally, we experimented on a unique meta-environment Walker-2DParams where tasks differ by transition functions and carry out ablations in 5.2.3.
These meta-RL benchmarks were previously introduced by Finn et al.
and are described in Appendix D. We train stochastic SAC models for every single task and roll out policies saved at every checkpoint to collect trajectories.
The ofﬂine training datasets can be collected as a selection of the saved trajectories, which facilitates tuning of the performance level and state-action distributions of the datasets for each task.
For OMRL, there are two natural baselines.
The ﬁrst is by modifying PEARL to train and test from solely logged data without exploration.
The second one is Contextual BCQ introduced by Li et al.
(2019b).
By incorporating a latent space z in the state space for task inference, it can be interpreted as the task-augmented variant of ofﬂine BCQ algorithm (Fujimoto et al., 2019).
Like PEARL, the task inference module is trained using Bellman gradients.
Lastly, we include comparison with the method proposed by Li et al. (2019a) which we term ”Distilled BCQ”.
Although as discussed earlier, Distilled BCQ is a model-based, two-stage method as opposed to our model-free and end-toend approach, we consider it by far the most competitive and related OMRL algorithm to FOCAL, due to the lack of other OMRL methods.
As shown in Figure 3, we observe that FOCAL outperforms prior ofﬂine meta-RL methods across all domains in terms of both sample efﬁciency and performance.
Most notably, when compared to Batch PEARL, a vanilla version of FOCAL without behavior regularization or distance metric learning, the gap in performance is evident due to divergence of value functions as well as suboptimal representation of the context encoder (Appendix E).
We started experiments with expert-level datasets.
However, for some tasks such as Ant-Fwd-Back, we observed that a diverse training sets (generated by expert, medium, random policies) result in a better meta policy (Table 2).
We conjecture that mixed datasets, despite sub-optimal actions, provides a broader support for state-action distributions, making it easier for the context encoder to learn the correct correlation between task identity and transition tuples (i.e., transition/reward functions).
While using expert trajectories, there might be little overlap between state-action distributions across tasks (Fig 7), which may cause the agent to learn spurious correlation.
This is the exact problem Li et al.
(2019b) aims to address.
Moreover, similar to other ofﬂine methods such as classical supervised learning, the performance of learned policy is sensitive to data distribution according on our observation (Table 5).
To incorporate a more competitive baseline, we also compare our design to a variant of FOCAL where the context encoder is trained by Bellman gradients, jointly with the control policy, to demonstrate the effectiveness of our proposed decoupled training.
We discuss the ablation in 5.2.3.
[Figure]
5. A BLATIONS
Based on our previous analysis, we examine and validate three key design choices of FOCAL by the following ablations.
5.2. D ETERMINISTIC VS . P ROBABILISTIC C ONTEXT E NCODER
Despite abundance successes of probabilistic/variational inference models in previous work (Kingma & Welling, 2013; Alemi et al., 2016; Rakelly et al., 2019), we observe experimentally that a deterministic context encoder performs signiﬁcantly better on both simple (Sparse-Point-Robot) and
complex (Walker-2D-Params) tasks in fully ofﬂine settings, shown in Figure 4.
As reﬂected in our design principles, this is due to
Ofﬂine meta-RL does not require exploration.
Also when the Task-Transition Correspondence is satisﬁed, there is not need for reasoning about uncertainty during adaption.
The deterministic context encoder in FOCAL is trained with carefully designed metricbased learning objective, detached from the Bellman update, which provides better efﬁciency and stability for meta-learning.
In principle, however, the deterministic context encoder which is built upon the Task-Transition Correspondence assumption, is less effective in presence of sparse reward or stochastic environments.
In this case, a non-negligible portion of samples violate Assumption 1 and provide confounding information for task inference, effective reducing the size of the dataset.
Due to these limitations, the fully-ofﬂine algorithm may be prone to over-ﬁtting.
Indeed, we observe mild performance instability and deterioration similar to overﬁtting on Half-Cheetah-Vel (Fig 3) and Sparse-Point-Robot (Fig 4(b)).
Nevertheless, we consider Assumption 1 a sufﬁcient condition for the feasibility of our deterministic context encoder design, but not necessary for FOCAL to work well.
In fact, even on environments where Assumption 1 is not satisﬁed (Sparse-Point-Robot), FOCAL still obtains excellent performance overall (Figure 3) given no signiﬁcant distribution shift between training and testing dataset (Table 5).
This is corroborated in our full-range experiments in Fig 6.
5.2. P OWER L AW OF D ISTANCE M ETRIC L OSS
To show the effectiveness of our proposed negative-power distance metrics for OMRL problem, we tested context embedding loss with different powers of distance, from L − 2 to L 2 .
A 2D PCA projection view of the high-dimensional embedding space in Figure 2 demonstrates that, distance metric loss with negative powers are more effective in separating embedding vectors of different tasks, whereas positive powers exhibit degenerate behaviors, leading to less robust and effective conditioned policies.
By a physical analogy, the inverse-power losses provide ”repulsive forces” that drive apart all data points, regardless of the initial distribution.
In electromagnetism, consider the latent space as a 3D metal cube and the embedding vectors as positions of positive charges.
By
Gauss’s law, at equilibrium state, all ”charges” are distributed on the surface of the cube with densities positively related to the local curvature of the surface.
Indeed, we observe from the ”Inversesquare” and ”Inverse” trials that almost all vectors are located near the edge of the latent space, with higher concentration around the vertices, which have the highest local curvatures (Fig 8(a)).
[Figure]
[Table] To evaluate the effectiveness of different powers of DML loss, we deﬁne a metric called effective separation rate (ESR) which computes the percentage of embedding vector pairs of different tasks whose distance on latent space Z is larger than the expectation of randomly distributed vector pairs, i.e., (cid:112) 2 l/on ( − 1 , 1) l .
Table 1 demonstrates that DML losses of negative power are more effective in maintaining distance between embeddings of different tasks, while no signiﬁcant distinction is shown in terms of RMS distance, which is aligned with our insight that RMS or effectively classical L 2 objective, can be optimized by degenerate distributions (Lemma 1).
This is the core challenge addressed by our proposed inverse-power loss.
5.2. C ONTEXT E NCODER T RAINING S TRATEGIES
The last key design choice of FOCAL is to train the context encoder via distance-based metrics, fully decoupled from control (Bellman updates).
To illustrate the necessity of such design, we compare FOCAL to a more competitive baseline than Batch PEARL by applying the same behavior regularization on the latter.
Without DML loss, we train the encoder entirely by Bellman gradients.
Since on simple environments such as Sparse-Point-Robot, FOCAL can learn expert-level policy without behavior regularization ( α = 0 in Table 3), this ablation is effectively addressed in Fig 3.
We therefore performed an additional ablation on the unique meta-environment Walker-2D-Params, where tasks differ by transition functions.
Fig 5 shows that our proposed strategy achieves effective clustering of task context and therefore better control policy, whereas training via Bellman gradients cannot.
We conjecture that on complex tasks where behavior regularization is necessary to ensure convergence, the value functions can often grow to an enormous scale ( 10 7 for instance, see Fig 9) with large variance, causing the Bellman gradients to dominate over the contribution of task inference module.
Evetually, context embedding collapses and fails to learn anything useful.
Further analyses are provided in Appendix E.
C ONCLUSION
In this paper, we propose a novel fully-ofﬂine meta-Rl algorithm, FOCAL, in pursuit of more practical RL.
Our method involves independent distance metric learning of a deterministic context encoder for efﬁcient task inference, combined with an actor-critic apparatus with regularization to learn effective policy from static data.
By re-formulating the meta-RL tasks as task-augmented MDPs with Task-Transition Correspondence assumption, we shed light on the effectiveness of our design choices in both theory and experiments.
Our approach achieves superior results compared to existing OMRL algorithms on a diverse set of continous control meta-RL domains.
Despite the success, we also discuss the breaking down of our fundamental assumptions in extreme scenarios such as sparse reward or stochastic environments, which opens up avenues for future work of more general and sophisticated OMRL algorithms.
B D EFINITIONS AND P ROOFS
Lemma
The contrastive loss of a given dataset X = { x i
| i
= 1 , ..., N } is proportional to the variance of the random variable X ∼ X Consider the contrastive loss (cid:80) i (cid:54) =
j ( x i − x j ) 2 , which consists of N ( N − 1) pairs of different samples ( x i , x j ) drawn from X .
It can be written as
[Equation]
The variance of X ∼ X is expressed as
[Equation]
where X denotes the expectation of X .
By substituting Eqn. 18 into 14, we have
[Equation]
Deﬁnition 1 (Task-Augmented MDP)
A task-augmented Markov Decision Process (TA-MDP) can be modeled as M = ( S , Z , A , P, R, ρ 0 , γ ) where
• S : state space •
Z : contextual latent space •
A : action space • P : transition function P ( s (cid:48) , z (cid:48) | s, z, a ) • R : reward function R ( s, z, a ) • ρ 0 ( s, z ) : joint initial state and task distribution • γ ∈ (0 , 1) : discount factor
Deﬁnition
The Bellman optimality operator B z on TA-MDP is deﬁned as
[Equation]
Deﬁnition 3 (Deterministic MDP)
For a deterministic MDP, a transition map t : S ×
A → S exists such that:
[Equation]
21) where δ ( x − y ) is the Dirac delta function that is zero almost everywhere except x = y .
C I MPORTANCE OF D ISTANCE M ETRIC L EARNING FOR M ETA -RL ON T ASK -A UGMENTED MDP S
We provide an informal argument that enforcing distance metric learning (DML) is crucial for meta
RL on task-augmented MDPs (TA-MDPs).
Consider a classical continuous neural network with
L ∈ N layers, n l ∈ N many nodes at the l -th hidden layer for l = 1 , ..., L , input dimension n 0 , output dimension n L +1 and nonlinear continuous activation function σ : R → R expressed as
[Equation]
22) where A l : R n l − 1 → R n l is an afﬁne linear map deﬁned by A l ( x ) =
W l x + b l for n l × n l −dimensional weight matrix W l and n l dimensional bias vector b l and σ l : R n l → R n l is an elementwise nonlinear continuous activation map deﬁned by σ l ( z ) :
= ( σ ( z 1 ) , ..., σ ( z n l )) .
θ denotes the network parameters.
Since every afﬁne and activation map is continuous, their composition N σ is also continuous, which means by deﬁnition of continuity:
[Equation]
[Equation] [Equation]
where | · | in principle denotes any valid metric deﬁned on Euclidean space R n .
A classical example is the Euclidean distance.
Now consider N θ as the value function on TA-MDP with deterministic embedding, approximated by a neural network parameterized by θ :
[Equation]
The continuity of neural network implies that for a pair of sufﬁciently close embedding vectors ( z i , z j ) , there exists sufﬁciently small η > 0 and (cid:15) > 0 that
[Equation]
Eqn 26 implies that for a pair of different tasks ( T i , T j ) ∼ p
( T ) ,
if their embedding vectors are sufﬁciently close in the latent space Z , the mapped values of meta-learned functions approximated by continuous neural networks are suffciently close too.
Since by Eqn 25, due to different transition functions P z i ( s (cid:48) | s, a ) , P z j ( s (cid:48) | s, a ) and reward functions
R z i ( s (cid:48) | s, a ) , R z j ( s (cid:48) | s, a ) of ( T i , T j ) , the distance between the true values of two Q-functions | Q θ ( s, a, z 1 ) − Q
θ ( s, a, z 2 ) | is not guaranteed to be small.
This suggests that a meta-RL algorithm with inefﬁcient representation of context embedding z = q φ ( c ) , which fails in maintaining effective distance between two distinct tasks T i , T j , is unlikely to accurately learn the value functions (or any policy-related functions) for both tasks simultaneously.
D E XPERIMENTAL D ETAILS D. D ETAILS OF THE M AIN E XPERIMENTAL R ESULT (F IGURE 3)
The main experimental result in the paper is the comparative study of performance of FOCAL with three baseline OMRL algorithms (Batch PEARL, Contextual BCQ and Distilled BCQ), shown in
Figure 3.
Here we plot the same data as in Figure 3 for the full number of time steps taken in our experiments, in Figure 6.
Due to limited computational budget, we only performed experiments of baseline algorithms Contextual BCQ and Distilled BCQ on three meta-environments (Half-CheetahVel, Half-Cheetah-Fwd-Back and Ant-Fwd-Back) for 10 6 steps.
In addition, we provide experimental details on how we produce the result in this section.
The performance levels of the training/testing data for the experiments are given in Table 2, which are selected for the best test-time performance over four levels: expert, medium, random, mixed (containing all logged trajectories of trained SAC models from beginning (random quality) to end (expert quality).
For mixed data, the diversity of samples is optimal but the average performance level is lower than expert:
[Table]
The DML loss used in experiments in Figure 3 is inverse-squared, which gives the best performance among the four power laws we experimented with in Figure 2.
D. D ESCRIPTION OF THE M ETA E NVIRONMENTS
• Sparse-Point-Robot:
A 2D navigation problem introduced in PEARL (Rakelly et al., 2019).
Starting from the origin, each task is to guide the agent to a speciﬁc goal located on the unit circle centered at the origin.
Non-sparse reward is deﬁned as the negative distance from the current location to the goal.
In sparse-reward scenario, reward is truncated to 0 when the agent is outside a neighborhood of the goal controlled by the goal radius.
While inside the neighborhood, agent is rewarded by 1 - distance, a positive value.
• Half-Cheetah-Fwd-Back: Control a Cheetah robot to move forward or backward.
Reward function is dependent on the walking direction.
• Half-Cheetah-Vel: Control a Cheetah robot to achieve a target velocity running forward.
Reward function is dependent on the target velocity.
• Ant-Fwd-Back: Control an Ant robot to move forward or backward.
Reward function is dependent on the walking direction. • Walker-2D-Params: Agent is initialized with some system dynamics parameters randomized and must move forward, it is a unique environment compared to others since tasks differ in transition function.
Transitions function is dependent on randomized task-speciﬁc parameters such as mass, inertia and friction coefﬁcients.
D. H YPERPARAMETER S ETTINGS
The details of important hyperparameters used to produce the experimental results in the paper are presented in Table 3 and 4.
[Figure] [Figure]
(a) Compared to Half-Cheetah-Vel experiment in Table 3, the meta batch and latent space dimension were reduced to speed up computation.
Trials β (cid:15)
Inverse-Square 0.
Inverse 0.
Linear 0.
Square 0.
[Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
when the average distance of x i and x j per dimension is 0.5, a reasonable value given x ∈ ( − 1 , 1) l .
E A DDITIONAL E XPERIMENTS E. S ENSITIVITY TO D ISTRIBUTION S HIFT
Since in OMRL, all datasets are static and ﬁxed, many challenges from classical supervised learning such as as over-ﬁtting apply.
By developing FOCAL, we are also interested in its sensitivity to distribution shift for better understanding of OMRL algorithms.
Since for each task T i , our datagenerating behavior policies β i ( a | s ) are trained from random to expert level, we select three performance levels (expert, medium, random) of datasets to study how combinations of training/testing sets with different quality/distribution affect performance.
An illustration of the three quality levels on Sparse-Point-Robot is shown in Fig 7.
Table 5 shows the average return at test-time for various training and testing distributions.
Sensitivity to distribution shift is conﬁrmed since training/testing on the similar distribution of data result in relatively higher performance.
In particular, this is signiﬁcant in sparse reward scenario when Assumption 1 is no longer satisﬁed.
With severe over-ﬁtting and task mis-identiﬁcation , performance of meta-RL policy is inevitably compromised by distribution mismatch between training datasets
[Table]
E. C ONTEXT E NCODER T RAINING S TRATEGIES : 5.2.3 C ONTINUED
As discussed several times in the paper, one of the key design choices of FOCAL is to decouple learning of task inference and control in terms of gradient computation, rather than sampling strategy as in PEARL.
We present our experimental evidence on Sparse-Point-Robot and Half-Cheetah-Vel in Fig 8.
Since it was found empirically that Sparse-Point-Robot is the only environment among the four that does not require behavior regularization to work, our control group of Batch PEARL can be seen as a variant of FOCAL which learns from Bellman update (same as PEARL) instead of
DML loss.
[Figure]
E. D IVERGENCE OF Q- FUNCTIONS IN O FFLINE S ETTING
The necessity of applying behavior regularization on environment like Ant-Fwd-Back to prevent divergence of value functions is demonstrated in Figure 9.
[Figure]
F I MPLEMENTATION
We build our algorithm on top of PEARL and BRAC, both are derivatives of the SAC algorithm.
SAC is an off-policy actor-critic method with a maximum entropy RL objective which encourages exploration and learning a stochastic policy.
Although exploration is not needed in fully-ofﬂine scenarios, we found empirically that a maximum entropy augmentation is still beneﬁcial for OMRL, which is likely due to the fact that different actions result in same next state and reward in environments such as Ant.
All function approximators in FOCAL are implemented as neural networks with MLP structures.
For normalization, the last activation layer of context encoder and policy networks are invertible squashing operators (tanh), making Z a bounded Euclidean space ( − 1 , 1) l , which is reﬂected in
Figure 1(a).
Since the training of context encoder and policy are completely decoupled, the whole FOCAL pipeline involves three main objectives.
The DML loss for training the inference network q φ ( z | c ) is given by Eqn 13, for mini-batches of transitions drawn from training datasets: x i ∼ D
i , x j ∼ D j .
The embedding vector q i , q j are computed as the average embedding over x i and x j .
The actor and critic losses are the task-augmented version of Eqn 8 and 9:
[Equation]
[Equation]
[Equation]
where ¯ Q is a target network and ¯ z indicates that gradients are not being computed through it.
As discussed in (Kumar et al., 2019; Wu, Tucker, and Nachum, 2019), the divergence function ˆ D can take form of Kernel MMD (Gretton et al., 2012),
Wasserstein Divergence (Arjovsky, Chintala, and
Bottou, 2017) or f-divergences (Nowozin et al., 2016) such as KL divergence.
In this paper, we use the dual form (Nowozin, Cseke, and Tomioka, 2016) of KL divergence, which learns a discriminator g with minimax optimization to circumvent the need of a cloned policy for density estimation.
In principle, as a core design choice of PEARL, the context used to infer q φ ( z | c ) can be distinct from the data used to construct the critic loss.
In OMRL however, we found this treatment unnecessary since there is no exploration.
Therefore training of DML and actor-critic objectives share the same sampled data, which form an end-to-end algorithm described in Algorithm 1 and 2.