AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning

1 INTRODUCTION AND RELATED WORK . Over the last decades , reinforcement learning ( RL ) ( Sutton and Barto , 1998 ) has been successful in many tasks ( Mnih et al. , 2013 ; Silver et al. , 2016 ) . Most of these early successes focus on a fixed task in a fixed environment . However , in real applications we often have changing environments , and it has been demonstrated that the optimal policy learned in a specific domain may not be generalized to other domains ( Taylor and Stone , 2009 ) . In contrast , humans are usually good at transferring acquired knowledge to new environments and tasks both efficiently and effectively ( Pearl and Mackenzie , 2018 ) , thanks to the ability to understand the environments . Generally speaking , to achieve reliable , low-cost , and interpretable transfer , it is essential to understand the underlying process—which decision-making factors have changes , where the changes are , and how they change , instead of transferring blindly ( e.g. , transferring the distribution of high-dimensional images directly ) . There are roughly two research lines in transfer RL ( Taylor and Stone , 2009 ; Zhu et al. , 2020 ) : ( 1 ) finding policies that are robust to environment variations , and ( 2 ) adapting policies from the source domain to the target domain as efficiently as possible . For the first line , the focus is on learning policies that are robust to environment variations , e.g. , by maximizing a risk-sensitive objective over a distribution of environments ( Tamar et al. , 2015 ) or by extracting a set of invariant states ( Zhang et al. , 2020a ; 2021a ; Tomar et al. , 2021 ) . A more recent method encodes task-relevant invariances by putting behaviorally equivalent states together , which helps better generalization ( Agarwal et al. , 2021 ) . On the other hand , with the increase of the number of domains , the common part may get even smaller , running counter to the intention of collecting more information with more domains . Moreover , focusing only on the invariant part and disregarding domain-specific information may not be optimal ; for instance , in the context of domain adaptation , it has been demonstrated that the variable part also contains information helpful to improve prediction accuracy ( Zhang et al. , 2020b ) . In this paper , we propose a method along the second line , adapting source policies to the target . Approaches along this line adapt knowledge from source domains and reuse it in the target domain to improve data efficiency , i.e. , in order for the agent to require fewer explorations to learn the target policy . For example , an agent could use importance reweighting on samples 〈s , a , r , s′〉 from sources ( Tirinzoni et al. , 2018 ; 2019 ) or start from the optimal source policy to initialize a learner in the target domain , as a near-optimal initializer ( Taylor et al. , 2007 ; Fernández et al. , 2010 ) . Another widely-used technique is finetuning : a model is pretrained on a source domain and the output layers are finetuned via backpropagation in the target domain ( Hinton and Salakhutdinov , 2006 ; Mesnil et al. , 2012 ) . PNNs ( Rusu et al. , 2016 ) , instead , retain a pool of pretrained models and learn lateral connections from them to extract useful features for a new task . Moreover , a set of approaches focus on sim2real transfer by adapting the parameters ( Yu et al. , 2017 ; Peng et al. , 2020 ) . However , many of these approaches still require a large amount of explorations and optimization in the target domain . Recently , meta-RL approches such as MAML ( Finn et al. , 2017 ) , PEARL ( Rakelly et al. , 2019 ) , CAVIA ( Zintgraf et al. , 2019 ) , Meta-Q learning ( Fakoor et al. , 2020 ) , and others ( Mendonca et al. , 2019 ; Nagabandi et al. , 2018 ; Duan et al. , 2016 ) have been successfully applied to learn an inductive bias that accelerates the learning of a new task by training on a large number of tasks . Some of these methods ( e.g. , CAVIA and PEARL ) , as well as some prior work ( e.g. , HiMDPs ( Doshi-Velez and Konidaris , 2016 ) ) and recent follow-ups ( Zhang et al. , 2021b ) , have a similar motivation to our work : in a new environment not all parameters need to be updated , so we can force the model to only adapt a set of context parameters . However , these methods mostly focus on MDPs ( except the Block MDP assumption in Zhang et al . ( 2021b ) ) and model all changes as a black-box , which may be less efficient for adaptation , as opposed to a factorized representation of change factors . Considering these limitations , we propose AdaRL , a transfer RL approach that achieves low-cost , reliable , and interpretable transfer for partially observable Markov decision processes ( POMDPs ) , with MDPs as a special case . In contrast to state-of-the-art approaches , we learn a parsimonious graphical representation that is able to characterize structural relationships among different dimensions of states , change factors , the perception , the reward variable , and the action variable . It allows us to model changes in transition , observation , and reward functions in a component-wise way . This representation is related to Factored MDPs ( Kearns and Koller , 1999 ; Boutilier et al. , 2000 ; Strehl et al. , 2007 ) and Factored POMDPs ( Katt et al. , 2019 ) , but augmented with change factors that represent a low-dimensional embedding of the changes across domains . Our main motivation is that distribution shifts are usually localized – they are often due to the changes of only a few variables in the generative processes , so we can just adapt the distribution of a small portion of variables ( Huang et al. , 2020 ; Schölkopf et al. , 2021 ) and , furthermore , factorized according to the graph structure , each distribution module can be adapted separately ( Schölkopf , 2019 ; Zhang et al. , 2020b ) . In Fig . 1 we give a motivating example and a general description of AdaRL . In this example , we consider learning policies for Pong ( Bellemare et al. , 2013 ) that can easily generalize to different rotations ω and to images corrupted with white noise . Specifically , given data from n source domains with different rotations and noise variances , we learn a parsimonious latent state representation shared by all domains , denoted by st , and characterize the changes across domains by a two-dimensional factor θk . We identify a set of minimal sufficient representations ( smint , θ min k ) for policy transfer . For instance , here only the rotation factor ω needs adapting ( i.e. , θmink = ωk ) , since the noise factor does not affect the optimal policy . Similarly , as we will show formally in the rest of the paper , not all components si , t of the state vector st are necessary for policy transfer . For example , s2 , t 6∈ smint , since it never affects the future reward . We learn an optimal policy π∗ ( ·|θmink ) on source domains 1 . In the target domain , we only need a few samples to quickly estimate the value of the low-dimensional θmintarget , and then we can apply π ∗ ( ·|θmintarget ) directly . Our main contributions are summarized below : • We assume a generative environment model , which explicitly takes into account the structural relationships among variables in the RL system . Such graphical representations provide a compact way to encode what and where the changes across domains are . • Based on this model , we characterize a minimal set of representations that suffice for policy learning across domains , including the domain-specific change factors and domain-shared state representations . With this characterization , we adapt the policy with only a few target samples and without policy optimization in the target domain , achieving low-cost and reliable policy transfer . • By leveraging a compact way to encode the changes , we also benefit from multi-task learning in model estimation . In particular , we propose the Multi-model Structured Sequential Variational Auto-Encoder ( MiSS-VAE ) for reliable model estimation in general cases . 2 A COMPACT REPRESENTATION OF ENVIRONMENTAL SHIFTS . Suppose there are n source domains and n′ target domains . In each source domain , we observe sequences { 〈ot , at , rt〉 } Tt=1 , where ot ∈ O are the perceived signals at time t ( e.g. , images ) , at ∈ A is the executed action , and rt ∈ R is the reward signal . We denote the underlying latent states by st = ( s1 , t , · · · , sd , t ) > , where d is the dimensionality of latent states . We assume that the generative process of the environment in the k-th domain ( with k = 1 , . . . n+ n′ ) can be described in terms of the transition function for each dimension of s and the observation and reward functions as si , t = fi ( c s ) s i st−1 , ca ) si · at−1 , c θk ) s i θ s k , s i , t ) , for i = 1 , · · · , d , ot = g ( c s ) o st , cθk ) o · θok , ot ) , rt = h ( c s ) r st−1 , ca ) r · at−1 , cθk ) r · θrk , rt ) , ( 1 ) where denotes the element-wise product , the si , t , ot , rt terms are i.i.d . random noises . As explained below , c· ) · are masks ( binary vectors or scalars that represent structural relationships from one variable to the other ) , and θk = ( θsk , θ o k , θ r k ) are the change factors that have a constant value in each domain , but vary across domains in the transition , observation , and reward function , respectively . The latent states st+1 form an MDP : given st and at , st+1 is independent of previous states and actions . The perceived signals ot are generated from the underlying states st . The actions at directly influence the latent states st+1 , instead of the observed signals ot , and the reward is determined by the latent states and the action . Eq . 1 can also represent MDPs as a special case if states st are directly observed , in which case the observation function of ot is not needed . Structural relationships and graphs . Often the action variable at−1 does not influence every dimension of st , and similarly , the reward rt may not be influenced by every dimension of st−1 . Furthermore , there are structural relationships between different dimensions of st−1 and st. To characterize these constraints , we explicitly take into account the graph structure G over the variables in the system characterized by a Dynamic Bayesian Network ( Murphy , 2002 ) and encode the edges with masks c· ) · . In the first equation in Eq . 1 the transition function for the state component si , where the jth entry of cs ) si ∈ { 0 , 1 } d is 1 if and only if sj , t influences si , t+1 ( graphically represented by an edge ) , while ca ) si ∈ { 0 , 1 } is 1 if and only if the action at has any effect on si , t+1 . Similarly , the binary vector cθk ) si ∈ { 0 , 1 } p encodes which components of the change factor θ s k = ( θ s 1 , k , . . . , θ s p , k ) > affect si , t+1 . The masks in the observation function g and reward function h have similar functions . The masks and the parameters of the functions f , g , and h , are invariant ; all changes are encoded in θk . For simplicity of notation , we collect all the transition mask vectors in the matrices Cs ) s : = [ cs ) si ] d i=1 and Cθk ) s : = [ cθk ) si ] d i=1 and the scalars in the vector c a ) s : = [ ca ) si ] d i=1 . Characterization of change factors in a compact way . In practical scenarios , the environment model may change across domains . Moreover , it is often the case that given a high-dimensional input , only a few factors may change , which is known as minimal change principle ( Ghassami et al. , 2018 ) or sparse mechanism shift assumption ( Schölkopf et al. , 2021 ) . In such a case , instead of learning the distribution shift over the high-dimensional input , thanks to the parsimonious graphical representation , we introduce a low-dimensional vector θk to characterize the domain-specific information in a compact way ( Zhang et al. , 2020b ) . Specifically , θok , θ r k , and θ s k capture the change factors in the observation function , reward function , and transition dynamics , respectively ; each of them can be 1We consider optimality of the policy w.r.t . the model estimated on source domains and AdaRL assumptions . multi-dimensional and that they are constant within each domain . In general , θk can capture both the changes in the influencing strength and those in the graph structure , e.g. , some edges may appear only in some domains . Since we assume that the structural relationships in Eq . 1 are invariant across domains , this means that the masks c· ) · have to encode an edge even if it presents only in one domain , and furthermore , since θk encodes the changes , it can switch the edge off in other domains . Fig . 1 shows an example of the graphical representation of the ( estimated ) environment model . Specifically , in this example , θsk only influences s1 , t , at−1 does not have an edge to s1 , t , and among the states , only sd , t−1 has an edge to rt . In this example , we consider the case when the control signals are random , so there is no edge between st and at .

Summary: The submission proposes a method for transfer in reinforcement learning building on estimating a small set of factors describing a system and modelling the agent as a dynamic bayesian network. In detail, the method splits into domain shared and domain dependent factors which are modelled via a new form of structured sequential VAE. The VAE is trained to enable dynamics, observation and reward prediction (including reconstruction, prediction, KL and sparsity regularisation losses).  Data for VAE training is collected from random policies and model fitting is followed by policy training on source domains to be transferred to target domains by identifying domain specific features for the target domains.  The method is evaluated on variations of a cartpole and a pong domain (both from images) and evaluated against a set of recent, competitive baselines.