Conservative Data Sharing for Multi-Task Offline Reinforcement Learning

1 Introduction . Recent advances in offline reinforcement learning ( RL ) make it possible to train policies for real-world scenarios , such as robotics [ 32 , 60 , 33 ] and healthcare [ 24 , 67 , 35 ] , entirely from previously collected data . Many realistic settings where we might want to apply offline RL are inherently multi-task problems , where we want to solve multiple tasks using all of the data available . For example , if our goal is to enable robots to acquire a range of different behaviors , it is more practical to collect a modest amount of data for each desired behavior , resulting in a large but heterogeneous dataset , rather than requiring a large dataset for every individual skill . Indeed , many existing datasets in robotics [ 17 , 11 , 66 ] and offline RL [ 19 ] include data collected in precisely this way . Unfortunately , leveraging such heterogeneous datasets leaves us with two unenviable choices . We could train each task only on data collected for that task , but such small datasets may be inadequate for good performance . Alternatively , we could combine all of the data together and use data relabeled from other tasks to improve offline training , but this naïve data sharing approach can actually often degrade performance over simple single-task training in practice [ 33 ] . In this paper , we aim to understand how data sharing affects RL performance in the offline setting and develop a reliable and effective method for selectively sharing data across tasks . A number of prior works have studied multi-task RL in the online setting , confirming that multitasking can often lead to performance that is worse than training tasks individually [ 56 , 62 , 90 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . These prior works focus on mitigating optimization challenges that are aggravated by the online data generation process [ 64 , 89 , 88 ] . As we will find in Section 4 , multi-task RL remains a challenging problem in the offline setting when sharing data across tasks , even when exploration is not an issue . While prior works have developed heuristic methods for reweighting and relabeling data [ 3 , 16 , 44 , 33 ] , they do not yet provide a principled explanation for why data sharing can hurt performance in the offline setting , nor do they provide a robust and general approach for selective data sharing that alleviates these issues while preserving the efficiency benefits of sharing experience across tasks . In this paper , we hypothesize that data sharing can be harmful or brittle in the offline setting because it can exacerbate the distribution shift between the policy represented in the data and the policy being learned . We analyze the effect of data sharing in the offline multi-task RL setting , and present evidence to support this hypothesis . Based on this analysis , we then propose an approach for selective data sharing that aims to minimize distributional shift , by sharing only data that is particularly relevant to each task . Instantiating a method based on this principle requires some care , since we do not know a priori which data is most relevant for a given task before we ’ ve learned a good policy for that task . To provide a practical instantiation , we propose the conservative data sharing ( CDS ) algorithm . CDS reduces distributional shift by sharing data based on a learned conservative estimate of the Q-values that penalizes Q-values on out-of-distribution actions . Specifically , CDS relabels transitions when the conservative Q-value of the added transitions exceeds the expected conservative Q-values on the target task data . We visualize how CDS works in Figure 1 . The main contributions of this work are an analysis of data sharing in offline multi-task RL and a new algorithm , conservative data sharing ( CDS ) , for multi-task offline RL problems . CDS relabels a transition into a given task only when it is expected to improve performance based on a conservative estimate of the Q-function . After data sharing , similarly to prior offline RL methods , CDS applies a standard conservative offline RL algorithm , such as CQL [ 39 ] , that learns a conservative value function or BRAC [ 82 ] , a policy-constraint offline RL algorithm . Further , we theoretically analyze CDS and characterize scenarios under which it provides safe policy improvement guarantees . Finally , we conduct extensive empirical analysis of CDS on multi-task locomotion , multi-task robotic manipulation with sparse rewards , multi-task navigation , and multi-task imaged-based robotic manipulation . We compare CDS to vanilla offline multi-task RL without sharing data , to naïvely sharing data for all tasks , and to existing data relabeling schemes for multi-task RL . CDS is the only method to attain good performance across all of these benchmarks , often significantly outperforming the best domain-specific method , improving over the next best method on each domain by 17.5 % on average . 2 Related Work . Offline RL . Offline RL [ 14 , 61 , 40 , 43 ] has shown promise in domains such as robotic manipulation [ 32 , 52 , 60 , 70 , 33 ] , NLP [ 29 , 30 ] , recommender systems & advertising [ 72 , 22 , 7 , 78 , 79 ] , and healthcare [ 67 , 80 ] . The major challenge in offline RL is distribution shift [ 20 , 37 , 39 ] , where the learned policy might generate out-of-distribution actions , resulting in erroneous value backups . Prior offline RL methods address this issue by regularizing the learned policy to be “ close “ to the behavior policy [ 20 , 50 , 29 , 82 , 93 , 37 , 68 , 57 ] , through variants of importance sampling [ 59 , 74 , 49 , 75 , 54 ] , via uncertainty quantification on Q-values [ 2 , 37 , 82 , 43 ] , by learning conservative Q-functions [ 39 , 36 ] , and with model-based training with a penalty on out-of-distribution states [ 34 , 91 , 53 , 4 , 76 , 60 , 42 , 92 ] . While current benchmarks in offline RL [ 19 , 25 ] contain datasets that involve multi-task structure , existing offline RL methods do not leverage the shared structure of multiple tasks and instead train each individual task from scratch . In this paper , we exploit the shared structure in the offline multi-task setting and train a general policy that can acquire multiple skills . Multi-task RL algorithms . Multi-task RL algorithms [ 81 , 56 , 77 , 15 , 27 , 89 , 85 , 88 , 33 , 71 ] focus on solving multiple tasks jointly in an efficient way . While multi-task RL methods seem to provide a promising way to build general-purpose agents [ 33 ] , prior works have observed major challenges in multi-task RL , in particular , the optimization challenge [ 27 , 64 , 89 ] . Beyond the optimization challenge , how to perform effective representation learning via weight sharing is another major challenge in multi-task RL . Prior works have considered distilling per-task policies into a single policy that solves all tasks [ 62 , 77 , 23 , 85 ] , separate shared and task-specific modules with theoretical guarantees [ 13 ] , and incorporating additional supervision [ 71 ] . Finally , sharing data across tasks emerges as a challenge in multi-task RL , especially in the off-policy setting , as naïvely sharing data across all tasks turns out to hurt performance in certain scenarios [ 33 ] . Unlike most of these prior works , we focus on the offline setting where the challenges in data sharing are most relevant . Methods that study optimization and representation learning issues are complementary and can be readily combined with our approach . Data sharing in multi-task RL . Prior works [ 3 , 31 , 58 , 63 , 16 , 44 , 33 , 8 ] have found it effective to reuse data across tasks by recomputing the rewards of data collected for one task and using such relabeled data for other tasks , which effectively augments the amount of data available for learning each task and boosts performance . These methods perform relabeling either uniformly [ 33 ] or based on metrics such as estimated Q-values [ 16 , 44 ] , domain knowledge [ 33 ] , the distance to states or images in goal-conditioned settings [ 3 , 58 , 55 , 48 , 73 , 47 , 28 , 51 , 87 , 8 ] , and metric learning for robust inference in the offline meta-RL setting [ 45 ] . All of these methods either require online data collection and do not consider data sharing in a fully offline setting , or only consider offline goal-conditioned or meta-RL problems [ 8 , 45 ] . While these prior works empirically find that data sharing helps , we believe that our analysis in Section 4 provides the first analytical understanding of why and when data sharing can help in multi-task offline RL and why it hurts in some cases . Specifically , our analysis reveals the effect of distributional shift introduced during data sharing , which is not taken into account by these prior works . Our proposed approach , CDS , tackles the challenge of distributional shift in data sharing by intelligently sharing data across tasks and improves multi-task performance by effectively trading off between the benefits of data sharing and the harms of excessive distributional shift . 3 Preliminaries and Problem Statement . Multi-task offline RL . The goal in multi-task RL is to find a policy that maximizes expected return in a multi-task Markov decision process ( MDP ) , defined asM = ( S , A , P , γ , { Ri , i } Ni=1 ) , with state space S , action space A , dynamics P ( s′|s , a ) , a discount factor γ ∈ [ 0 , 1 ) , and a finite set of task indices 1 , · · · , N with corresponding reward functions R1 , · · · , RN . Each task i presents a different reward function Ri , but we assume that the dynamics P are shared across tasks . While this setting is not fully general , there are a wide variety of practical problem settings for which only the reward changes including various goal navigation tasks [ 19 ] , distinct object manipulation objectives [ 83 ] , and different user preferences [ 10 ] . In this work , we focus on learning a policy π ( a|s , i ) , which in practice could be modelled as independent policies { π1 ( a|s ) , · · · , πN ( a|s ) } that do not share any parameters , or as a single task-conditioned policy , π ( a|s , i ) with parameter sharing . Our goal in this paper is to analyze and devise methods for data sharing and the choice of parameter sharing is orthogonal , and can be made independently . We formulate the policy optimization problem as finding a policy that maximizes expected return over all the tasks : π∗ ( a|s , · ) : = argmaxπ Ei∼ [ N ] Eπ ( ·|· , i ) [ ∑ t γ tRi ( st , at ) ] . The Q-function , Qπ ( s , a , i ) , of a policy π ( ·|· , i ) is the long-term discounted reward obtained in task i by executing action a at state s and following policy π thereafter . Standard offline RL is concerned with learning policies π ( a|s ) using only a given static dataset of transitions D = { ( sj , aj , s′j , rj ) } Nj=1 , collected by a behavior policy πβ ( a|s ) , without any additional environment interaction . In the multi-task offline RL setting , the dataset D is partitioned into per-task subsets , D = ∪Ni=1Di , where Di consists of experience from task i . While algorithms can choose to train the policy for task i ( i.e. , π ( ·|· , i ) ) only on Di , in this paper , we are interested in data-sharing schemes that correspond to relabeling data from a different task , j ̸= i with the reward function ri , and learn π ( ·|· , i ) on the combined data . To be able to do so , we assume access to the functional form of the reward ri , a common assumption in goal-conditioned RL [ 3 , 16 ] , and which often holds in robotics applications through the use of learned classifiers [ 83 , 32 ] , and discriminators [ 18 , 9 ] . We assume that relabeling data Dj from task j to task i generates a dataset Dj→i , which is then additionally used to train on task i . Thus , the effective dataset for task i after relabeling is given by Deffi : = Di ∪ ( ∪j ̸=iDj→i ) . This notation simply formalizes data sharing and relabeling strategies explored in prior work [ 16 , 33 ] . Our aim in this paper will be to improve on this naïve strategy , which we will show leads to significantly better results . Offline RL algorithms . A central challenge in offline RL is distributional shift : differences between the learned policy and the behavior policy can lead to erroneous target values , where the Q-function is queried at actions a ∼ π ( a|s ) that are far from the actions it is trained on , leading to massive overestimation [ 43 , 37 ] . A number of offline RL algorithms use some kind of regularization on either the policy [ 37 , 20 , 82 , 29 , 68 , 57 ] or on the learned Q-function [ 39 , 36 ] to ensure that the learned policy does not deviate too far from the behavior policy . For our analysis in this work , we will abstract these algorithms into a generic constrained policy optimization problem [ 39 ] : π∗ ( a|s ) : = argmax π JD ( π ) − αD ( π , πβ ) . ( 1 ) JD ( π ) denotes the average return of policy π in the empirical MDP induced by the transitions in the dataset , and D ( π , πβ ) denotes a divergence measure ( e.g. , KL-divergence [ 29 , 82 ] , MMD distance [ 37 ] or DCQL [ 39 ] ) between the learned policy π and the behavior policy πβ . In the multitask offline RL setting with data-sharing , the generic optimization problem in Equation 1 for a task i utilizes the effective dataset Deffi . In addition , we define πeffβ ( a|s , i ) as the effective behavior policy for task i and it is given by : πeffβ ( a|s , i ) : = |Deffi ( s , a ) |/|Deffi ( s ) | . Hence , the counterpart of Equation 1 in the multi-task offline RL setting with data sharing is given by : ∀i ∈ [ N ] , π∗ ( a|s , i ) : = argmax π JDeffi ( π ) − αD ( π , π eff β ) . ( 2 ) We will utilize this generic optimization problem to motivate our method in Section 5 .

Summary: This paper proposes a data-sharing strategy called CDS to tackle the distributional shift of the offline RL algorithm in the multitask setting. The authors reveal the problem of naive sharing strategies experimentally and theoretically. CDS achieves an overall advantage against naive sharing strategies in various multitask offline environments.