Year: 2021

Abstract:
One practical challenge in reinforcement learning (RL) is how to make quick adaptations when faced with new environments. In this paper, we propose a principled framework for adaptive RL, called AdaRL, that adapts reliably and efficiently to changes across domains with a few samples from the target domain, even in partially observable environments. Specifically, we leverage a parsimonious graphical representation that characterizes structural relationships over variables in the RL system. Such graphical representations provide a compact way to encode what and where the changes across domains are, and furthermore inform us with a minimal set of changes that one has to consider for the purpose of policy adaptation. We show that by explicitly leveraging this compact representation to encode changes, we can efficiently adapt the policy to the target domain, in which only a few samples are needed and further policy optimization is avoided. We illustrate the efficacy of AdaRL through a series of experiments that vary factors in the observation, transition and reward functions for Cartpole and Atari games.

heading: 1 INTRODUCTION AND RELATED WORK
text: 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.

heading: 2 A COMPACT REPRESENTATION OF ENVIRONMENTAL SHIFTS
text: 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.

heading: 3 WHAT, WHERE, AND HOW TO ADAPT IN RL
text: We first assume that the environment model in Eq. 1 is known (we will explain how to learn it in Sec. 3.1), and characterize which changes have an effect on the policy transfer to the target domain. In Eq. 1, we allow the model to change across domains, including all involved functions, and we leverage θk to capture the changes in a compact way. The varying model implies that the optimal policy function may also vary across domains. How can we then characterize the changes in the optimal policy function in a compact way, as we did in the model? Interestingly, we find that the change factor θk and the latent state st are sufficient for policy learning, but not every dimension of θk or st is necessary, since they may not ever have an effect on the reward, even in future steps. We first give the definitions of compact domain-shared representations and compact domain-specific representations, according to the graph structure, and we further show that they are the minimal set of dimensions that suffice for policy learning across domains (proof in Appendix). Definition 1. Given the graphical representation of an environment model G that is encoded in the binary masks c· )·, we define recursively the representations that affect the reward in the future as:
• compact domain-shared representations smint : the latent state components si,t ∈ st that either – have an edge to the reward in the next time-step rt+1, i.e., cs )ri = 1, or – have an edge to another state component in the next time-step sj,t+1, i.e., cs )sj,i = 1, such that the same component at time t is a compact domain-shared representation, i.e., sj,t ∈ smint ; • compact domain-specific representations θmink : the latent change factors θi,k ∈ θk that either:
– have an edge to the reward in the next time-step rt+1, i.e., θi,k = θrk and cθk )r = 1, or – have an edge to a state component sj,t ∈ smint , i.e., c θk )s j,i = 1.
Proposition 1. Under the assumption that the graph G is Markov and faithful to the measured data, the union of compact domain-specific θmink and compact shared representations smint are the minimal and sufficient dimensions for policy learning across domains.
For the example in Fig. 1, smint = (s1,t, sd,t) and θ min k = {θsk, θrk}. Note that θok is never in θ min k , and thus if only the observation function changes, the optimal policy function π∗k remains the same across domains. For example in Cartpole a change of color does not affect the optimal policy. Moreover, if cθk )r = 1, then θrk ∈ θ min k , which is the case for multi-task learning.

heading: 3.1 SIMULTANEOUS ESTIMATION OF DOMAIN-VARYING MODELS
text: In this section, we give the estimation procedure of the environment model in Eq. 1 from observed sequences {{〈yt,k, at,k〉}Tt=1}nk=1 from each source domain k, where yt,k = (o>t,k, r>t,k)> are the observations and reward at time t in domain k. Instead of estimating the model in each domain separately, we estimate models from different domains simultaneously, by exploiting commonalities across domains while at the same time preserving specific information for each domain. In particular, we propose the Multi-model Structured Sequential Variational Auto-Encoder (MiSS-VAE), which contains the following three essential components. (1) "Sequential VAE" component handles the sequential data, with the underlying latent states satisfying an MDP. It is implemented by adding an LSTM (Hochreiter and Schmidhuber, 1997) to encode the sequential information in the encoder to learn the inference model qφ(st,k|st−1,k,y1:t,k, a1:t−1,k;θk). (2) "Multi-model" component handles models from different domains at the same time, using the domain index k as an input and learning the domain-specific factors θk. (3) "Structured" component: exploits the structural information that is explicitly encoded with the binary masks, i.e., c· )· in Eq. 1. Here, the joint distribution of latent states are factorized according to structures, instead of being marginally independent as in traditional VAEs (Kingma and Welling, 2013). Figure 2 gives the diagram of neural network architecture in
model training. Let y1:n1:T = {{yt,k}Tt=1}nk=1. By taking into account the above three components, we maximize the following objective function L: L(y1:n1:T ; (β1, β2, φ, γ, c·) = Lrec(y1:n1:T ; (β1, φ, c·))+Lpred(y1:n1:T ; (β2, φ))−LKL(y1:n1:T ; (φ, γ, c·))−Lreg. In particular, Lrec is the reconstruction loss for both observed images and rewards, to learn the observation and the reward function, respectively. We also consider the one-step prediction loss Lpred. Lrec =
n∑ k=1 T−2∑ t=1 Est,k∼qφ(·|θk){log pβ1(ot,k|st,k; θok, cθk )o, cs )o) + log pβ1(rt+1,k|st,k, at,k; θrk, cθk )r, cs )r, ca )r)},
Lpred = n∑ k=1 T−2∑ t=1 Est,k∼qφ(·|θk){log pβ2(ot+1,k|st,k, θok, θsk) + log pβ2(rt+2,k|st,k, at+1,k; θrk, θsk)},
where pβ1 and pβ2 denote the generative models with parameters β1 and β2, respectively, that are shared across domains, and qφ the inference model with shared parameters φ. We also use the following KL-divergence loss to constrain the latent space:
LKL = λ0 n∑ k=1 T∑ t=2 KL ( qφ(st,k|st−1,k,y1:t,k, a1:t−1,k;θk)‖pγ(st,k|st−1,k, at−1,k; θsk,Cs )s, ca )s,Cθk )s) ) , where we explicitly model the transition dynamics pγ with the parameters γ shared across domains; this is essential for establishing a Markov chain in latent space and learning a representation for long-term predictions. Moreover, the KL loss helps to constrain the latent space to (1) ensure that the disentanglement between the inferred latent factors q(si,t|·) and q(sj,t|·) for i 6= j, since we do not consider the instantaneous connections among state dimensions, and (2) ensure that the latent representations st are maximally compressive about the observed high-dimensional data. Furthermore, according to the edge-minimality property (Zhang and Spirtes, 2011) and the minimal change principle (Ghassami et al., 2018), we add sparsity constraints on structural matrices and on the change of domain-specific factors across domains, respectively, to achieve better identifiability: Lreg = λ1‖cs )o‖1 + λ2‖cs )r‖1 + λ3‖ca )r‖1 + λ4‖Cs )s‖1 + λ5‖ca )s‖1 + λ6‖Cθk )s‖1 + λ7
∑ 1≤j,k≤n |θj − θk|.
Note that besides the shared parameters {β1, β2, φ, γ}, the structural relationships (encoded in binary masks c) are also involved in the shared parameters. Each factor in pφ, pβi , and pγ is modeled with a mixture of Gaussians, because with a suitable number of Gaussians, it can approximate a wide class of continuous distributions. Moreover, in model estimation, the domain-specific factors θk = {θok, θsk, θrk} are treated as parameters; they are constant within the same domain, but may differ in different domains. We explicitly consider θk not only in the generative models pβi and pγ , but also in the inference model qφ. In this way, except for θk, all other parameters in MiSS-VAE are shared across domains, so that all we need to update in the target domain is the low-dimensional θk, which greatly improves the sample efficiency and the statistical efficiency in the target domain.

heading: 3.2 LOW-COST AND INTERPRETABLE POLICY TRANSFER
text: After identifying what and where to transfer, we show how to adapt. Instead of learning the optimal policy in each domain separately, which is time and sample inefficient, we leverage a multi-task learning strategy: policies in different domains are optimized at the same time exploiting both commonalities and differences across domains. Given the compact domain-shared smint and domainspecific representations θmink , we represent the optimal policies across domains in a unified way:
at = π ∗(smint ,θ min k ), (2)
where θmink explicitly and compactly encodes the changes in the policy function in each domain k, and all other parameters in the optimal policy function π∗ are shared across domains. In other
words, by learning π∗ in the source domains, and estimating the value of the change factor θmintarget and inferring latent states smintarget from the target domain, we can immediately derive the optimal policy in the target domain without further policy optimization by just applying Eq. 2.
The AdaRL framework answers what and where the change factors are and which change factors need to adapt across domains in an interpretable way. Moreover, AdaRL only requires a few samples to update the low-dimensional domain-specific parameters θmintarget to achieve the optimal policy in the target domain, without further policy optimization, achieving the low cost. We provide the pseudocode for the AdaRL algorithm in Alg. 1. The algorithm has three parts: (1) data collection with a random policy or any initial policy from n source domains (line 2), (2) model estimation from the n source domains with multi-task learning (lines 2-3, see Sec. 3.1 for details), and (3) learning the optimal policy π∗ with deep Q-learning, by making use of domain-specific factors and the inferred domain-shared state representations (lines 4-21). Specifically, because we do not directly observe the states st, we infer q(smint+1,k|o≤t+1,k, r≤t+1,k, a≤t,k,θ min k ) and sample s min t+1,k from its posterior, for the kth domain (lines 7 and 13). Moreover, the action-value function Q is learned by considering the averaged error over the n source domains (line 18). AdaRL can be implemented with a wide class of policy-learning algorithms, e.g., DDPG (Lillicrap et al., 2015), Q-learning (Mnih et al., 2015), and Actor-Critic methods (Schulman et al., 2015; Mnih et al., 2016). Then, in the target domain, we only need to collect a few rollouts to estimate the low-dimensional domain-specific representations θmintarget, with all other parameters being fixed (lines 22-23).
Algorithm 1 (AdaRL with Domains Shifts) 1: Initialize action-value function Q, target action-value function Q′, and replay buffer B. 2: Record multiple rollouts for each source domain k (k = 1, · · ·, n) and estimate the model in Eq.1. 3: Identify the dimension indices of smint and the values of θ min k according to the learned model.
4: for episode = 1, . . . , M do 5: for source domain k = 1, . . . , n do 6: Receive initial observations o1,k and r1,k for the k-th domain. 7: Infer the posterior q(smin1,k |o1,k, r1,k,θ min k ) and sample initial inferred state s min 1,k . 8: end for 9: for timestep t = 1, . . . , T do
10: for source domain k = 1, . . . , n do 11: Select at,k randomly with probability ; otherwise at,k = argmaxaQ(s min t,k , a,θ min k ). 12: Execute action at,k, and receive reward rt+1,k and observation ot+1,k in the kth domain. 13: Infer the posterior q(smint+1,k|o≤t+1,k, r≤t+1,k, a≤t,k,θ min k ) and sample s min t+1,k. 14: Store transition (smint,k , at,k, rt+1,k, s min t+1,k,θ min k ) in reply buffer B. 15: end for 16: Randomly sample a minibatch of N transitions (smini,j , ai,j , ri+1,j , s min i+1,j ,θ min j ) from B. 17: Set yi,j = ri+1,j + λmaxa′ Q′(smini+1,j , a ′,θminj ). 18: Update action-value function Q by minimizing the loss:
L = 1 n ∗N ∑ i,j (yi,j −Q(smini,j , ai,j ,θ min j )) 2.
19: end for 20: Update the target network Q′: Q′ = Q. 21: end for 22: Record a few rollouts from the target domain. 23: Estimate the values of θmintarget for the target domain, with all other parameters fixed.

heading: 3.3 THEORETICAL PROPERTIES
text: Below we show the conditions under which we can identify the true graph G from observational data, even when the model in Eq. 1 is unknown. Furthermore, we derive a generalization bound of the state-value function under the PAC-Bayes framework (McAllester, 1999).
Theorem 1 (Structural Identifiability). Suppose the underlying states st are observed, i.e., Eq. (1) is an MDP. Then under the Markov condition and faithfulness assumption, the structural matrices Cs )s, ca )s, cs )r, ca )r, Cθk )s, and cθk )r are identifiable.
This theorem shows that in the MDP scenario, where the underlying states are observed and cθk )o and cs )o are not considered by definition, we can uniquely determine the structural relationships over {st−1, st, at−1, rt,θk}, i.e., the Dynamic Bayesian network G, from observed data under mild conditions, without knowing the generative environment model. Even if θk is not directly observed, we can identify which state dimension changes and if there is a change in the reward function. Suppose there are n source domains, and for the kth domain, we have Sk =( (s1,k, v ∗(s1,k)), · · · , (smk,k, v∗(smk,k)) ) , where mk is the number of samples from the kth domain, s·,k is a state sampled from the kth domain, and v∗(s·,k) is its corresponding optimal state-value. For any value function hθmink (·) parameterized by θ min k , we define the loss function `(hθmink , (sk,i, v ∗(si,k))) = Ddist(hθmink (si,k), v
∗(si,k)), where Ddist is a distance function that measures the discrepancy between the learned value and the optimal value. The following theorem gives a generalization bound of the state-value function under the PAC-Bayes framework.
Theorem 2 (Generalization Bound). Let Q be an arbitrary distribution over θmink and P the prior distribution over θmink . Then for any δ ∈ (0, 1], with probability at least 1 − δ, the following inequality holds uniformly for all Q, er(Q) ≤ 1n n∑ k=1 { êr(Q, Sk) + √ 1 2(mk−1) ( DKL(Q||P) + log 2nmkδ ) + √ 1 2(n−1) ( DKL(Q||P) + log 2nδ )} ,
where er(Q) and êr(Q, Sk) are the generalization error and the training error between the estimated value and the optimal value, respectively.
Theorem 2 states that with high probability the generalization error er(Q) is upper bounded by the empirical error plus two complexity terms. Specifically, the first one is the average of the taskcomplexity terms from the observed domains, which converges to zero in the limit of samples in each domain, i.e., mk →∞. The second is an environment-complexity term, which converges to zero if infinite domains are observed, i.e., n→∞. Moreover, if assuming different dimensions of θmink are independent, then DKL(Q||P) = ∑|θmink | i=1 DKL(Qi||Pi), which indicates that a low-dimensional θmink usually has a smaller KL divergence, so does the upper bound of the generalization error.

heading: 4 EVALUATION
text: We modify the Cartpole and Atari Pong environments in OpenAI Gym (Brockman et al., 2016). Here, we present a subset of the results; see Appendix for the complete results and the detailed settings. We consider changes in the state dynamics (e.g., the change of gravity or cart mass in Cartpole, change of orientation in Pong), changes in observations (e.g., different noise levels in images or different colors in Pong), and changes in reward functions (e.g., different reward functions in Pong based on the contact point of the ball), as shown in Fig. 3 for Pong. For each of these factors, we take into account both interpolation (where the factor value in the target domain is in the support of that in source domains), and extrapolation (where it is out of the support w.r.t. the source domains). We train on n source domains based on the trajectory data generated by a random policy. In Cartpole, for each domain we collect 10000 trials with 40 steps. For Pong experiments, each domain contains 40 episodes data and each of them takes a maximum of 10000 steps. In the target domain we consider different sample sizes with Ntarget = {20, 50, 10000} to estimate θmintarget. For both games, we evaluate the POMDP case, where the inputs are high-dimensional images; note that we did not stack multiple frames, so some properties (e.g., velocity) are not observed, resulting in a POMDP. For Cartpole, we also consider the MDP case, where the true states (cart position and velocity, pole angle and angular velocity) are used as the input to the model. In Cartpole, we also experiment with multiple factors changing at the same time (e.g., gravity and mass change concurrently in the target domain).
Modified Cartpole setting The Cartpole problem consists of a cart and a vertical pendulum attached to the cart using a passive pivot joint. The task is to prevent the vertical pendulum from falling by putting a force on the cart to move it left or right. We introduce two change factors for the state dynamics θsk: varying gravity and varying mass of the cart. In terms of changes on the observation function θok, we add Gaussian noise on the images. Since θ o k does not influence the optimal policy (as shown in Prop. 1), we need it only for the model estimation, but not for policy optimization. Moreover, if θk = {θok}, the optimal policy is shared across domains.
Modified Pong setting In Pong, one of the established Atari benchmarks (Bellemare et al., 2013), the agent controls a paddle moving up and down vertically, aiming at hitting the ball. We consider changes in observation function θok, state dynamics θ s k, and reward function θ r k, as shown in Fig. 3. We consider three change factors on perceived signals θok: different image sizes, different image colors, and different noise levels. For the setting with different image colors, we use RGB images as inputs and consider source domains with varying RGB colors {original, green, red} and target domains with colors {yellow, white}, but for other settings, we convert the images to grayscale as input. To change the state dynamics, we rotate the images ω degrees clockwise. To test the changes in the reward function, we model the reward as a function of the distance between contact point and the central point of the paddle, denoted by d, as opposed to the original Pong in which it is constant. We denote by L the half-length of the paddle and formulate two groups of reward functions: (1) Linear reward: rt= αdL ; and (2) Non-linear reward: rt= αL d+3L , where α varies across domains.
Baselines In the MDP setting, we compare AdaRL with CAVIA (Zintgraf et al., 2019) and PEARL (Rakelly et al., 2019). In the POMDP setting, we compare with PNN (Rusu et al., 2016), PSM (Agarwal et al., 2021) and MQL (Fakoor et al., 2020). We also compare with AdaRL*, a version of AdaRL that does not learn the binary masks c· )· and therefore does not use any structural information. All of these methods use the same number of samples Ntarget from the target domain. We also compare with: 1) Non-t, a vanilla non-transfer baseline that pools data from all source domains and learns a fixed model; and 2) an oracle baseline, which is completely trained on the target domain with model-free exploration. For a fair comparison, we use the same policy learning algorithm, Double DQN (Van Hasselt et al., 2016), for all methods. As opposed to MAML and PNN, AdaRL only uses the Ntarget samples to estimate θmintarget, without any policy optimization.
Results We measure performance by the mean and standard deviation of the final scores over 30 trials with different random seeds, as well as testing the significance with the Wilcoxon signedrank test (Conover, 1999). As shown in Tables 1, 2 and 3 2, AdaRL consistently outperforms the baselines across most change factors in the MDP and POMDP case for modified Cartpole, and
2In Table 1-3, "•" indicates the baselines for which the improvements of AdaRL are statistically significant (via Wilcoxon signed-rank test at 5% significance level).
in the POMDP case for Pong for Ntarget = 50. As ablation studies, to see the effect of learning the graphical structure, we also compare with AdaRL∗, which does not learn the binary masks c·, but just assumes everything is fully connected. Learning the graphical structure improves the performances significantly, and without it AdaRL∗ is generally comparable to baselines. We provide results with Ntarget = {20, 50, 10000} in Appendix, showing that the performance gains are larger at smaller sample sizes. Furthermore, we consider the change of reward functions (see Table A12-14 in Appendix). More detailed experimental results are provided in Appendix, including the average score across different Ntarget, policy learning curves and an analysis of the estimated θk w.r.t. real change factor. Interestingly, in the Cartpole case, the estimated θk matches the physical quantities that are being changed across the domains. In particular, the estimated θk for gravity and noise are linear mappings of the gravity and noise level values. For the mass-varying case, the learned θsk is a nonlinear monotonic function of the mass, which matches the influence of the mass on the dynamics.

heading: 5 CONCLUSIONS AND FUTURE WORK
text: In this paper, we proposed AdaRL, a principled framework for transfer RL. AdaRL learns a latent representation with domain-shared and domain-specific components across source domains, uses it to learn an optimal policy parameterized by the domain-specific parameters, and applies it to a new target domain. It is achieved without any further policy optimization, but just by estimating the values of the domain-specific parameters in the target domain, which can be accomplished with a few target-domain data. As opposed to previous work, AdaRL can model changes in the state dynamics, observation function and reward function in an unified manner, and exploit the factorization to improve the data efficiency and adapt faster with fewer samples. Further directions include exploiting the target domain to fine-tune the policy and handling the out-of-distribution data. Moreover, exploring an alternative to the reconstruction loss, e.g., using the contrastive loss (Srinivas et al., 2020), might also improve the training efficiency. Finally, an exciting next step is to transfer knowledge across different tasks, e.g., different Atari games.
REPRODUCIBILITY STATEMENT
Proofs of all our theoretical results are given in Appendix A1, A2, and A3 with disclosure of all assumptions. More details about the model estimation are given in Appendix A4. Appendix A5 provides the complete experimental details and results. All datasets used are publicly available or instructions are provided on how to generate them in Appendix 5 and 6. A description of the hyperparameters and network architectures used is included in Appendix 6. Source code for all our experiments are attached in the supplemental material and will be made available on GitHub after deanonymisation, providing also a complete description of our experimental environment, configuration files and instructions on the reproduction of our experiments.

Review summaries:
The paper proposes the algorithm "AdaRL" a transfer method for reinforcement learning for different domains. The method is based on learning a latent representation with "domain shared" and "domain specific" components. A policy that is parameterized by "domain specific" components is learned. The transfer is done by collecting some data in target domain and estimating "domain specific" variable of the target domain. The authors show the algorithms mechanism for POMDP and MDP settings. The method is evaluated on modified versions of Cart-Pole and Pong domains. The authors test different settings where parameters of the environment as well as rewards functions are changed within source domains and the target domain. The authors also evaluate interpolation vs extrapolation cases. The method is compared with latest transfer learning methods and the results show that AdaRL performs better compared to other methods 

Instead of implicitly updating the policy using data from the source domain, learn a particularly structured latent model and the elements of variation and learn a policy that performs pretty well on some set of the elements of variation. At test time, estimate the elements of variation and provide them as input to the policy. I think this is a good paper but could benefit from some clearer writing as discussed in the clarity section below.

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.