Factored Adaptation for Non-Stationary Reinforcement Learning
Abstract
Dealing with non-stationarity in environments (i.e., transition dynamics) and objectives (i.e., reward functions) is a challenging problem that is crucial in real-worldapplications of reinforcement learning (RL).
While most current approaches modelthe changes as a single shared embedding vector, we leverage insights from therecent causality literature to model non-stationarity in terms of individual latentchange factors and causal graphs across different environments.
In particular, wepropose Factored Adaptation for Non-Stationary RL (FANS-RL), a factored adaption approach that learns jointly the causal structure in terms of a factored MDP,and a factored representation of the individual time-varying change factors.
Weprove that under standard assumptions, we can completely recover the causal graphrepresenting the factored transition and reward function and a partial structurebetween the individual change factors and the state components.
Through ourgeneral framework, we can consider general non-stationary scenarios with different changing function types and changing frequency, including changes acrossepisodes and within episodes.
Experimental results demonstrate that FANS-RLoutperforms existing approaches in terms of rewards, compactness of the latentstate representation, and robustness to varying degrees of non-stationarity.
Introduction
Learning a stable policy under non-stationary environments is a long-standing challenge in Reinforcement learning (RL) [1, 2, 3].
While most RL approaches assume stationarity, in many real-worldapplications of RL there can be changes in the dynamics or the reward function, both across different episodes and within each episode.
Recently, several works adapted Meta-RL methods tolearn sequences of non-stationary tasks [4, 5].
However, the continuous MAML [6] adaptation fornon-stationary RL [4] does not explicitly model temporal changing components while TRIO [5]needs to meta-train the model on a set of non-stationary tasks.
LILAC [7] and ZeUS [8] leveragelatent variable models to directly model the change factors in the environment in a shared embeddingspace.
In particular, they consider families of MDPs indexed by a single latent parameter.
In thispaper, we argue that disentangling the changes as separate latent parameters and modeling the processwith a factored representation improves the efficiency of adapting to non-stationarity.
In particular, we leverage insights from recent causality literature [9] that model non-stationarity interms of individual latent change factors and causal graphs across different environments.
We propose
Factored Adaptation for Non-Stationary RL (FANS-RL), a factored adaptation framework that jointlylearns the causal structure of the MDP and a factored representation of the individual change factors,allowing for changes at discrete timepoints and continuously varying environments.
While we providea specific architecture (FN-VAE), the theoretical framework of FANS-RL can be implemented withdifferent architectures and combined with various RL algorithms.
We formalize our setting as a
Factored Non-stationary MDP (FN-MDP), which combines a Factored-MDP [10, 11, 12], for whichwe learn the structure, with latent change factors that evolve in time following a Markov process.
We build upon the AdaRL framework [13], a recently proposed fast adaptation approach.
AdaRLlearns a factored representation that explicitly models changes (i.e., domain-specific components)in observation, dynamics and reward functions across a set of source domains.
An optimal policylearnt on the source domains can then be adapted to a new target domain simply by identifying alow-dimensional change factor, without any additional finetuning.
FANS-RL extends AdaRL fromthe stationary case with constant change factors to a general non-stationary framework.
Specifically,
FANS-RL learns the low-dimensional and time-evolving representations θ st and θ rt that fully capturethe non-stationarity of dynamics and rewards, allowing for continuous and discrete changing functions,both within-episode and across-episode .
Our main contributions can be summarized as:
• We formalize FN-MDPs, a unified factored framework that can handle many non-stationary
settings, including discrete and continuous changes, both within and across episodes.
We
prove that, under standard assumptions, the causal graph of the transition and reward function
is identifiable, while we can recover a partial structure for the change factors.
• We introduce Factored Adaptation for Non-Stationary RL (FANS-RL), a general non-
stationary RL approach that interleaves model estimation of an FN-MDP and policy opti-
mization.
We also describe FN-VAE, an example architecture for learning FN-MDPs.
• We evaluate FANS-RL on simulated benchmarks for continuous control and robotic manip-
ulation tasks and show it outperforms the state of the art on rewards, compactness of the
latent space representation and robustness to varying degrees of non-stationarity.
Factored Non-stationary MDPs
To model different types of non-stationarity in a unified and factored way, we propose Factored
Non-stationary Markov Decision Processes (FN-MDPs).
FN-MDPs are an augmented form of afactored MDPs [10, 11, 12] with latent change factors that evolve over time following a Markovprocess.
Since the change factors are latent, FN-MDPs are partially observed.
We define them as:
Definition 1.
A Factored Non-stationary Markov Decision Process (FN-MDP) is a tuple( S , A , Θ s , Θ r , γ, G , P s , R , , P θ r , P θ s ) , where S is the state space, A the action space, Θ s the space ofthe change factors for the dynamics, Θ r the space of the reward change factors and γ the discount factor.
We assume G is a Dynamic Bayesian Network over { s 1 ,t , ... , s d,t , a 1 ,t , ... , a m,t , r t , θ s 1 ,t , ... , θ s p,t ,θ r 1 ,t , ... , θ rq,t } , where d , m , p , and q are the dimensions of states, action, change factors on dynamicsand reward, respectively.
We define the factored state transition distribution
P s as:
[Equation]
where pa ( s i,t ) denotes the causal parents of s i,t in G , which are a subset of the dimensions of s t − 1 ,a t − 1 and θ st .
Note that the action a t − 1 is a vector of m dimensions in our setting.
We assume agiven initial state distribution P s ( s 0 ) .
Similarly, we define the reward function R as a function ofthe parents of r t in G , i.e., R ( s t , a t , θ r t ) =
R ( pa ( r t )) , where pa ( r t ) are a subset of dimensions ofs t , a t , and θ r t .
We define the factored latent change factors transition distributions P θ s and P θ r as:
[Equation]
where pa ( θ sj,t ) are a subset of the dimensions of θ s t − 1 , while pa ( θ rk,t ) are a subset of dimensions ofθ r t − 1 .
We assume the initial distributions P θ s ( θ s 0 ) and P θ r ( θ r 0 ) are given.
We show an example DBN representing the graph G of an FN-MDP in Fig. 1(a).
Since we areinterested in learning the graphical structure of the FN-MDP, as well as identifying the values of thelatent change factors from data, we describe a generative process of an FN-MDP environment.
Inparticular, we will learn a set of binary masks c · (cid:41) · and C · (cid:41) · that are the indicators for edges in G .
Generative environment model.
We adapt the generative model in AdaRL [13] across k differentdomains to a time-varying setting on a single domain.
We assume the generative process of theenvironment at timestep t in terms of the transition function for each dimension i = 1 , ... , d of s t is:
[Equation]
s 1, t − a 1, t − r t θ s 1, t s 2, t − s 3, t − a 2, t − s 1, t s 2, t s 3, t θ s 2, t θ r 1, t θ r 2, t a 1, t a 2, t θ s 1, t
+ θ s 2, t
+ θ r 1, t + θ r 2, t + (a) (b)
CF inference network CF dynamics network Transition decoder Reward decoder s t a t r t q sϕ h s t θ rt p α̂ s t + ̂ r t +1 p
β̂ r t p βp γ s p γ r s t a t ̂
s t p α 1 a t − s t
C θ s →
θ s C θ r
→ θ r c s → r c a → r s t +1 a t + a t s t −
C s → s C a → s C θ s → s θ s t q rϕ s t +1 a t +1 r t + q sϕ a t +1 s t +1 θ rt + p α̂ s t + ̂ r t +
p β̂ r t +1 p βp
α 1 s t a t
C s → s C a → s C θ s → s c s → r c a → r s t +2 a t + s t +1 a t + ̂ s t
+ θ s t +1 q rϕ h
rt where ⊙ is the element-wise product, f i are non-linear functions and ϵ si,t is an i.i.d.
random noise.
The binary mask c s (cid:41) s i ∈ { 0 , 1 } d represents which state components s j,t − 1 are used in the transition 84 function of s i,t .
Similarly, c a (cid:41) s i ∈ { 0 , 1 }
m indicates whether the action directly affects s i,t .
Thechange factor θ s t ∈ R p encodes any change in the dynamics.
The binary mask c θ t (cid:41) s i ∈ { 0 , 1 } p 86 represents which of the components of θ s t influence s i,t .
We model the reward function as:
[Equation]
where c s (cid:41) r ∈ { 0 , 1 } d , c a (cid:41) s ∈ { 0 , 1 } m , and ϵ rt is an i.i.d.
random noise.
The change factor θ rt ∈ R qencodes any change in the reward function.
The binary masks c · (cid:41) · can be seen as indicators of edgesin the DBN G .
In AdaRL, all change factors are assumed to be constant in each domain.
Since in thispaper we allow the change parameters to evolve in time, we introduce two additional equations:
[Equation]
for i = 1 , ... , d , j = 1 , ... , p , k = 1 , ... , q , and g s , and g r are non-linear functions.
We assume thebinary masks c · (cid:41) · are stationary across timesteps and so are the ϵ si,t , ϵ rt , ϵ θ s t and ϵ θ r t , the i.i.d.
randomnoises.
Although c · (cid:41) · and ϵ are stationary, we model the changes in the functions and some changesin the graph structure through θ .
For example a certain value of θ rt can switch off the contribution ofsome of the state or action dimensions in the reward function, or in other words nullify the effectof some edges in G .
Similarly the contribution of the noise distribution to each function can bemodulated via the change factors.
On the other hand, this setup does not allow adding edges that arenot captured by the binary masks c · (cid:41) · .
We group the binary masks in the matrices C s (cid:41) s :=
[ c s (cid:41) s i ] di =1 ,
C θ s (cid:41) s :=
[ c θ s (cid:41) s i ] di =1 , and C a (cid:41) s :=
[ c a (cid:41) s i ] di =1 .
Similarly, we also group the binary vectors in thedynamics of the latent change factors in matrices C θ s (cid:41) θ s :=
[ c θ s (cid:41) θ s j ] pj =1 and C θ r (cid:41) θ r :=
[ c θ r (cid:41) θ r k ] qk =1 .
Since latent change factors θ s and θ r follow a Markov process based on g s and g r , we can considerdifferent types of changes by varying the form of g s and g r , generalizing the approaches in literature.
We can also model concurrent changes in dynamics and reward, including different types of changes,e.g. a continuous g s and a piecewise-constant g r .
Compact representations.
Huang et al.
[13] show that the only dimensions of the state and changefactors useful for policy learning are those that eventually affect the reward.
These dimensions arecalled compact representations and are defined as the dimensions of the state and change factors witha path (i.e. a sequence of edges → ) to the present or future reward r t
+ τ for τ ≥ 0 in the DBN G
:s i,t ∈ s min ⇐⇒ s i,t → ...
→ r t + τ for τ ≥ 0 , and θ i ∈ θ min ⇐⇒ θ i → ... → r t + τ for τ ≥
Continuous changes.
If g s and g r are continuous, then they can model smooth changes in theenvironment, including across episodes.
While the functions in Eq.
1-3 allow us to model withinepisode changes, i.e. changes that can happen only before t
= H where H is the horizon, we also want to model across-episode changes.
We use a separate time index ˜ t that models the agent’s lifetime.
Initially ˜ t = t for t ≤ H , but while we reset t = 0 afterwards, ˜ t continues to grow indefinitely.
Discrete changes.
We assume discrete changes happen at specific timesteps and can be representedwith a piecewise constant function.
In particular we denote change timesteps as ˜ t =
(˜ t 1 , . .
. , ˜ t M )where ˜ t i describes a specific timepoint in the agent’s lifetime time index ˜ t .
This allows us tomodel within-episode changes .
In this case, we assume that the change happens always at the samesteps ( t 1 , ... , t m ) in each episode, i.e., we assume that ˜ t
= ( t 1 , ... , t m , H + t 1 , ... , H + t m , 2 H +t 1 , ... , 2 H + t m , ... ) , where H is the horizon.
We can also model across-episode changes , when weassume the change points only occur at the end of each episode, i.e., ˜ t
= ( H, 2 H, 3 H, ... ) .
We extend the result on identifiability of factored MDPs in AdaRL [13] to non-stationary environments.
We first assume that we observe the change factors and show we can identify the true causal graph G :
Proposition 1 (Full identifiability with observed change factors) .
Suppose the generative processfollows Eq.
1-3 and all change factors θ s t and θ r t are observed, i.e., Eq. 1-3 is an MDP.
Underthe Markov and faithfulness assumptions, i.e. conditional independences correspond exactly tod-separations, all binary masks C · (cid:41) · are identifiable, i.e., we can fully recover the causal graph G .
We provide all proofs and a detailed explanation in Appendix B. If we do not observe the changefactors θ s t and θ r t , we cannot identify their dimensions, and we cannot fully recover the causal graph
G .
On the other hand, we can still identify the partial causal graph over the state variables s t , rewardvariable r t , and action variable a t .
We can also identify which dimensions in s i,t have changes, i.e.,we can identify C θ s (cid:41) s .
We formalize this idea in the following (proof in Appendix B):
Proposition 2 (Partial Identifiability with latent change factors) .
Suppose the generative processfollows Eq.
1-3 and the change factors θ s t and θ r t are unobserved.
Under the Markov and faithfulnessassumptions, the binary masks C s (cid:41) s , C a (cid:41) s , c s (cid:41) r and c a (cid:41) r are identifiable.
Moreover, we can identifywhich state dimensions are affected by θ s t and whether the reward function changes.
This means that even in the most general case, we can learn most of the true causal graph G in an
FN-MDP, with the exception of the transition structure of the latent change factors.
In the following,we show a variational autoencoder setup to learn the generative process in FN-MDPs.
Learning the generative process in FN-MDPs
There are many possible architectures to learn FN-MDPs through Eq.
We propose FN-VAE , avariational autoencoder architecture described in Fig. 1(b).
In FN-VAE, we jointly learn the structuralrelationships, state transition function, reward function, and transition function of the latent changefactors, as described in detail in Appendix Alg.
A1.
An FN-VAE has four types of components:change factor (CF) inference networks that reconstruct the latent change factors, change factor(CF) dynamics networks that model their dynamics with an LSTM [14], transition decoders thatreconstruct the state dynamics at the time t and predict one step further at t + 1 , and reward decodersthat reconstruct the reward at t and predict the future reward at t
+ 1 .
We now describe them in detail.
CF inference networks (blue boxes in Fig. 1(b)).
The two inference models for latent changefactors q ϕ s ( θ st | s t , a t ) and q ϕ r ( θ rt | s t , a t , r t ) are parameterised by ϕ s and ϕ r , respectively.
Tomodel the time-dependence of θ st and θ rt , we use LSTMs [14] as inference networks.
At timestep t ,the dynamics change factor LSTM infers q ϕ s (cid:0)
θ st | s t , a t , r t , h st − 1 (cid:1) , where h st − 1 ∈ R L is the hiddenstate in the LSTM.
Thus we can obtain µ ϕ s ( τ 0: t ) and σ 2 ϕ s ( τ 0: t ) using q ϕ s , and sample the latentchanging factor θ s t ∼ N ( µ ϕ s ( τ 0: t ) , σ 2 ϕ s ( τ 0: t )) , where τ 0:
t = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 2 , . . .
, s t , a t , r t ) .
Similarly, the reward change factor LSTM infers q ϕ r ( θ rt | s t , a t , r t , h rt − 1 ) , where h rt − 1 ∈ R L is thehidden state, such that we can sample θ rt ∼ N ( µ ϕ r ( τ 0: t ) , σ 2 ϕ r ( τ 0: t )) .
CF dynamics network (orange boxes in Fig. 1(b)).
We model the dynamics of latent change factorswith
p γ s (cid:0) θ s t +1 | θ s t , C θ s (cid:41) θ s (cid:1) and p γ r (cid:0) θ r t +1 | θ r t , C θ r (cid:41) θ r (cid:1) .
To ensure the Markovianity of θ s t andθ rt , we define a loss L KL that helps minimize the KL-divergence between q ϕ and p γ .
[Equation]
If we assume that the change between θ s t and θ s t +1 , and similarly θ r t , is smooth, we can add asmoothness loss L smooth .
We provide a smooth loss for discrete changes in Appendix D.2.
[Equation]
(5) 162 Transition decoders (purple boxes in Fig. 1(b)).
We learn an approximation of the transitiondynamics in Eq.
1 by learning a reconstruction, parameterized by α 1 , and a prediction encoder,parametrized by α 2 .
To simplify the formulas, we define C · (cid:41) s :=
( C s (cid:41) s , C a (cid:41) s ,
C θ s (cid:41) s ) .
At timestept , the reconstruction encoder p α 1 ( s t | s t − 1 , a t − 1 , θ st ; C · (cid:41) s ) reconstructs the state from currentstate s t with sampled θ s t .
The one-step prediction encoder p α 2 ( s t +1 | s t , a t , θ s t ) instead tries toapproximate the next state s t +1 .
We do not use the prediction loss, when the one-step predictionis not smooth.
In particular, we do not use it for the last time-step in episode i if there is a changehappening at the first step in episode ( i + 1) , since the states in new episodes will be randomlyinitiated.
We also do not use it in the case of discrete changes at the timesteps (˜ t 1 − 1 , . . .
, ˜ t M − 1) .
The loss functions are:
[Equation]
6) 173 Reward decoders (green boxes in Fig. 1(b)).
Similarly, we use a reconstruction encoder p β 1 ( r t | s t , a t , θ rt , c s (cid:41) r , c a (cid:41) r ) , parameterized by β 1 , and a one-step prediction encoderp β 2 ( r t +1 | s t
+1 , a t +1 , θ rt ) , parametrized by β 2 , to approximate the reward function.
Similarly 176 to transition decoders, we do not use the one-step prediction loss, if it is not smooth.
The losses are:
[Equation]
Sparsity loss.
We encourage sparsity in the binary masks C · (cid:41) · to improve identifiability, by usingfollowing loss with adjustable hyperparameters ( w 1 , ... , w 7 ) , which we learn through grid search.
[Equation]
The total loss is
[Equation]
-rw ) − k 3 L KL − k 4 L sparse − k 5 L smooth , 180 where ( k 1 , ... , k 5 ) : hyper-parameters, which we learn with an automatic weighting method [15].
Learning from raw pixels.
Our framework can be easily extended to image inputs by adding anencoder ϕ o to learn the latent state variables from pixels, similar to other works [16, 17, 13].
In thiscase, our identifiability results do not hold anymore, since we cannot guarantee that we identify thetrue causal variables.
We describe this component in Appendix D.3.
FANS-RL: online model estimation and policy optimization
We propose Factored Adaptation for Non-Stationary RL (FANS-RL), a general algorithm thatinterleaves model estimation and policy optimization, as shown in Alg.
After we estimate theinitial FN-MDP with the FN-VAE, we can identify compact representations s min and θ min following
AdaRL [13].
In particular, the only dimensions of the state and change factors that are useful forpolicy learning are those that have a directed path to the reward in the graph G .
The online policyπ ψ (cid:0) a t | s mint , q ϕ ( θ mint | τ 0: t ) (cid:1) can be learned end-to-end, including learning the FN-VAE, as shownin Alg.
We use SAC [18] as our policy learning model, so the policy parameters are ψ = ( π, Q ) .
Continuous changes.
In Alg.
1 we describe our framework in case of continuous changes that canspan across episodes.
We start by collecting a few trajectories τ
and then learn our initial FN-VAE(Lines 3-4).
We can use the graphical structure of the initial FN-VAE to identify the compactrepresentations (Line 5).
During the online model estimation and policy learning stage, we estimatethe latent change factors θ st and θ rt using the CF dynamics networks γ s and γ r (Lines 8-20).
Sincein this case, we assume the dynamics of the change factors are smooth across episodes, at time t = 0 do 8:
10: θ s ← θ s θ old 11:
12: θ s ← θ st − θ r ← θ rt − 1 13: end
14: for j = s, r do 15:
Infer mean µ γ j ( θ j ) and variance σ j ( of the parameter jt via p j 16:
Sample θ jt ∼ N (cid:16) µ γ j ( θ j ) , σ 2 γ j ( θ j ) (cid:17) 17: end for 18: if t = H − 1 then 19: θ s old ← θ st and θ r old ← θ st ; 20: end if 21:
Generate a t ∼ π ψ ( a t
| s mint , θ mint and receive reward r n,t from
Add s to
D ; 23:
Extract a trajectory with length k from replay buffer D ; 24:
Learn FN-VAE (Appendix Alg.
A1) with updateG=
False (i.e. with fixed masks C · (cid:41) · ); 25:
Sample a batch of data from replay buffer D and update policy network parameters ψ ; 26: end for 27: end for we will use the last timestep ( H − 1) of the previous episode as a prior on the change factors (Line10).
Otherwise, we will estimate the change factors using their values in the previous timestep t − 1 .
We use the estimated latent factors θ t and observed state s t to generate a t using π ψ and receive areward r t (Line 21).
We add ( s t , a t , r t , θ st , θ rt ) to the replay buffer (Line 22).
We now update ourestimation of the FN-VAE, but we keep the graph G fixed (Lines 23-24).
Finally, we sample a batchof trajectories in the replay buffer and update the policy network ψ (Line 25).
Discrete changes.
Since we assume discrete changes happen at specific timestep ˜ t =
(˜ t 1 , . .
. , ˜ t M ) ,we can easily modify Alg.
1 for discrete changes, both within-episode and across-episode, by changing
Lines 9-20 to only update the change parameters at the timesteps in ˜t , as shown in Appendix Alg.
A2.
Evaluation
We evaluate our approach on four well-established benchmarks, including Half-Cheetah-V3 from
MuJoCo [19, 20], Sawyer-Reaching and Sawyer-Peg from Sawyer [21, 17], and Minitaur [22].
Wemodified these tasks to test several non-stationary RL scenarios with continuous and discrete changes.
The results suggest that FANS-RL can (1) obtain high rewards, (2) learn meaningful mappings tothe change factors with compact latent vectors, and (3) be robust to different non-stationary levels,different functional forms (e.g. piecewise linear, sine, damped sine) and multiple concurrent changes.
For space limits, we only highlight a subset of results, and report the full results in Appendix C.4.
Half-Cheetah-v3 in MuJoCo.
In this task, the agent is moving forward using the joint legs and theobjective is to achieve the target velocity v g .
We consider both the changes in the dynamics (changeof the wind forces f w , change of gravity) and reward functions (change of target velocity v g ).
Thereward function is r t = − ∥ v ot − v gt ∥ 2 − 0 .
∥ a t ∥ 2 , where v o and a t are the agent’s velocity and 220 action, respectively, at timestep t .
The number of time steps in each episode is 50 .
For dynamics,
we
[Figure]
TRIO
[Figure] [Figure]
[Table]
[Figure]
D)
A-EP(R+D)
A-EP(R+2D) 0 .0 .0 .1 .
R e w a r d s
Oursw/ostructure w/osmooth w/ominsets w/osparsity w/orewardpred w/ostatepred w/odisentangled change the wind forces f w in the environment.
Moreover, in terms of the reward functions, we changethe target velocity v g to be a time-dependent variable.
The change function in the dynamics f w canbe either continuous or discrete, and discrete changes can happen both within and across episodes.
Similarly to LILAC [7], we choose different functions (piecewise linear, sine and damped sine),besides allowing the change at specified intervals, we also allow it to change smoothly.
The change inthe reward function v g is not generally stable in the continuous case, so we only consider the discreteand across episode change functions v g for the reward.
We also design a scenario where dynamic andreward functions change concurrently.
We report all equations for g s and g r in Appendix C.1.
Sawyer.
We consider two robotic manipulation tasks, Sawyer-Reaching and Sawyer-Peg.
In SawyerReaching, the sawyer arm is trained to reach a target position s g .
The reward r t is the differencebetween the current position s t and the target position r t = − ∥ s t − s g ∥ 2 .
In this task, we cannot 232 directly modify the dynamics in the simulator, so consider a reward-varying scenario, where thetarget location changes across each episode following a periodic function.
In Sawyer-Peg, therobot arm is trained to insert a peg into a designed target location s g .
The reward function isr t
= I ( ∥ s t − s g ∥ 2 ≤ 0 . 05) .
In this task, we cannot directly modify the dynamics in the simulator, soconsider a reward-varying scenario, where the target location changes across each episode followinga periodic function.
In order to compare with similar approaches, e.g., ZeUS [8], CADM [23],
Hyperdynamics [24]and Meld [17], we evaluate our method on raw pixels.
Following [8], we
[Table]
consider discrete across-episode changes, where the target location can change in each episode, andis randomly sampled from a small interval.
We describe the non-stationary settings in Appendix C.2.
Minitaur.
A minitaur robot is a simulated quadruped robot with eight direct-drive actuators.
Theminitaur is trained to move at the target speed s t,v .
The reward is r t = 0 .
3 − | 0 .
3 − s t,v | − 0 .
01 ·∥ a t − 2 a t − 1 + a t − 2 ∥ 1 .
We modify (1) the mass m (dynamics) and (2) target speed s t,v (reward) of 244 the minitaur.
We consider continuous, discrete across-episode and discrete within-episode changes forthe dynamics, and across-episode changes for the reward.
We describe the settings in Appendix C.3.
Baselines.
We compare our approach with a meta-RL approach for stationary RL, VariBAD [25],a meta-RL approach for non-stationary RL, TRIO [5], as well as with two representative taskembedding approaches, LILAC [7] and ZeUS [8].
The details on the meta-learning setups are givenin Appendix D.4.
We also compare with stationary RL method, SAC [18], which will be our lowerbound, and compare with an oracle agent that has access to the full information of non-stationarity(e.g.
, the wind forces) and can use it to learn a policy, which will be our upper-bound.
For allbaselines, we use SAC for policy learning.
We compare with the baselines on average rewards andthe compactness of the latent space in varying degrees and functional forms of non-stationarity.
Experimental results and ablation studies.
Fig.
2(a)-(c) shows the smoothed curves of averagerewards across timesteps in a subsection of Half-Cheetah, Sawyer-Reaching, and Minitaur experiments.
Smoothed curves of other experiments are given in Appendix Fig.
A2 and A3.
For continuouschanges, we only compare with LILAC [7] since other approaches are not applicable.
We smooth thelearning curves by randomly selecting a few data points for readability.
The original learning curvesare in Appendix Fig.
A4-A6.
Table 1 shows the results on Sawyer-Peg experiments using raw pixelsas input, based on the reported results in [8].
For a fair comparison, we indicate the best and finalaverage rewards for each baseline.
The learning curve is given in Appendix Fig.
A12.
Full results aregiven in Appendix C.4, including significance tests for Wilcoxon signed-rank test at α = 0 .
05 showing that FANS-RL is significantly better than baselines.
Finally, we conduct ablation studies on eachcomponent of FANS-RL and report some key results in Fig. 2(d).
Full results are in Appendix C.5.
The ablation studies verify the effectiveness of all components, including binary masks/structure,smoothness loss, sparsity loss, reward prediction or state prediction.
The results show that the largestgain is provided by the factored representation , validating our original hypothesis, followed by stateprediction.
As expected, reward prediction is also important when there is a nonstationary reward,while smoothness is important for within-episode changes.
The disentangled design of CF inferencenetworks is valuable when there are changes on both dynamics and reward functions.
Visualization on the learned θ .
To verify that the learned θ can capture the true change factors,we compute the pairwise distance between learned θ r at different time steps.
We randomly sample10 time steps from the Half-Cheetah experiment with across-episode changes on reward functions.
Fig. 2(e) gives the pairwise distance of θ r among 5 time steps (from episode 148 , 279 , 155 , 159 , 230 )with different target speed values ( 0 . 04 , 0 . 48 , 0 . 89 , 2 . 06 , 2 . 85 ), respectively.
We can find that there isa positive correlation between the distance of learned θ r and values of change factors.
Meanwhile,we also sample 5 time steps from episodes 31 , 188 , 234 , 408 with target values around 1 . 67 .
Fig.
2(f)shows that the distance of θ r among these 5 time steps is very small, indicating that the learned θ rare almost the same for the similar values of change factors at different time steps.
The visualizationsuggests that the learned θ can capture meaningful mappings from the time-varying factors.
Varying latent dimensions, non-stationary levels and functional forms Fig. 2(g) shows thenormalized rewards versus the number of two latent features on Half-Cheetah.
Our frameworkcan learn a better policy with relatively smaller feature dimensions in the latent space than otherapproaches.
As we show in Appendix C.4, Saywer and Minitaur have the similar trend, where welearn a better policy than the baselines with fewer latent features.
We also vary the non-stationarylevels in Half-Cheetah with discrete across-episode changes on both dynamics and rewards.
A highernon-stationary degree indicates a faster change rate.
Fig. 2(h) shows that FANS-RL achieves thehighest rewards across all tested non-stationary degrees and that the gap increases with the non- stationarity.
We also test FANS-RL together with all baselines on different non-stationary functions,including piecewise linear, damping-like and sinusoid waves.
Fig.
2(i) displays the results, whichindicate that FANS-RL can generally outperform the baselines on diverse non-stationary functionforms.
Detailed function equations and experimental setups can be referred to Appendix C.
Multiple change factors.
We consider different numbers and types of changes to verify the benefitsfrom the factored structure in FANS-RL.
We conduct experiments with 1) only change wind forces(1D); 2) change wind forces and gravity concurrently (2D); 3) change wind force and target speed(1D+1R); and 4) change wind force, gravity, and target speed together (2D+1R) in an across-episodeway in Half-Cheetah.
From Fig. 2(j), we find that, thanks to the factored representation, FANS-RLperforms better in those more complicated scenarios with multiple numbers and types of changes.
Related work
Early works in non-stationary RL [26, 27] only detect changes that have already happened insteadof anticipating changes.
If the evolution of non-stationary environments is a (Semi)-Markov chain,one can deal with non-stationarity with HM-MDPs [28] or HS3MDPs [29].
Several recent methodslearn to anticipate changes in non-stationary deep RL.
Chandak et al.
[30] propose to maximizefuture rewards without explicitly modeling the non-stationary environments.
MBCD [31] uses highconfidence change-point detection to decide whether the agent should learn a novel policy or reusepreviously trained policies.
Al-Shedivat et al.
[4] extend MAML [6] for the non-stationary setting,but do not explicitly model the temporal changes.
TRIO [5] tracks the non-stationarity by inferringthe evolution of latent parameters, which captures the temporal change factors during the meta-testingphase.
ReBAL and GrBAL [32] meta-train the dynamic prior, which adapts to the local contextsefficiently.
However, these methods have to meta-train the model on a set of non-stationary tasks,which may not be accessible in real-world applications.
Another line of research directly learns thelatent representation to capture the non-stationary components.
In particular, LILAC [7] and ZeUS [8]leverage latent variable models to directly model the change factors in environments.
Our approachfits in this line of work; however, as opposed to these methods, which model all temporal changesusing a shared embedding space or a mixture models of dynamics, we use a factored representation.
Among works on factored representations in stationary settings, Hallak et al.
[33] learn factored MDPs[10, 11, 12] to improve sample efficiency in model-based off-policy RL.
Balaji et al.
[34] employknown factored MDP to improve both model-free and model-based RL algorithms.
Working memorygraphs [35] learn the factored observation space using Transformers.
NeverNet [36] factorizes thestate-action space through graph neural networks.
Zholus et al.
[37] factorize the visual states intoactor-related, object of manipulation, and the latent influence factor between these two states toachieve sparse interaction in robotic control tasks.
AFaR [38] learns factored value functions viaattention modules to improve sample efficiency and generalization across tasks.
In an MDP setting,
Mutti et al.
[39] learn a causal structure that generalizes across a family of MDPs.
Another related lineof work is based on independent causal mechanisms [40, 41].
Recurrent independent mechanisms(RIMs)
[42] learn the independent transition dynamics in RL with sparse communication among thelatent states.
Meta-RIMs [43] leverage meta-learning and soft attention to learn a set of RIMs withcompetition and communication.
As opposed to these works, we do not assume that the mechanismsare independent and we learn the factored structure among all components in MDPs with a DBN.
Conclusions, Limitations and Future work
We describe Factored Adaptation for Non-Stationary RL (FANS-RL), a framework that learns afactored representation for non-stationarity that can be combined with any RL algorithm.
Weformalize our problem as a Factored Non-stationary MDP (FN-MDP), augmenting a factored MDPwith latent change factors evolving as a Markov process.
FN-MDPs do not model a family of MDPs,but instead include the dynamics of change factors analogously to the dynamics of the states.
Thisallows us to capture different non-stationarities, e.g., continuous and discrete changes, both withinand across different episodes.
To learn FN-MDPs we propose FN-VAEs, which we integrated in
FANS-RL, an online model estimation and policy evaluation approach.
We evaluate FANS-RL onbenchmarks for continuous control and robotic manipulation, also with pixel inputs, and show itoutperforms the state of the art on rewards and robustness to varying degrees of non-stationarity.
Learning the graph in model estimation is computationally expensive, which limits the scalability ofour approach.
In future work, we plan to meta-learn the graphs among different tasks to improve thescalability of our approach and its applicability to complex RL problems, e.g., multi-agent RL.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] There are theoretical claims backed by two theoretical
results in Prop. 1 and Prop. 2.
The empirical claims are backed by an extensive
evaluation in Section 5.
(b) Did you describe the limitations of your work?
[Yes] We did provide a discussion on
the limitations and potential future work in Section
(c) Did you discuss any potential negative societal impacts of your work?
[Yes]
We
discussed the broader societal impact in Appendix A.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] The assump-
tions are stated in Prop. 1 and Prop. 2.
(b) Did you include complete proofs of all theoretical results?
[Yes] The complete proofs
are in the Appendix B
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] We will include
the code and instructions in the supplementary files.
(b) Did you specify all the training details (e g., data splits, hyperparameters,
how they
were chosen)?
[Yes] The details are in the Appendix D.4.
(c) Did you report error bars (e g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] The learning curves (Fig. 2 (a-c), Appendix Fig. A2-A6)
include error bars with 10 runs.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] The information of computational
resource is given in Appendix E.1.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] Yes, we have cited
the OpenAI Gym and Mujoco libraries in our paper.
(b) Did you mention the license of the assets?
[Yes] Yes, the licenses are given in Appendix E.2.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
Appendix
for Factored Adaptation for
Non-Stationary
Reinforcement LearningA Broader impact
Our work is a first exploration in leveraging factored representations for non-stationary RL in orderto improve adaptation.
One of the limitations of our current approach, and similarly to a few othernon-stationary RL approaches, is that we do not provide theoretical guarantees in terms of adaptingto non-stationarity.
This limits the applicability of these approaches in safety-critical applications, e.g.self-driving cars, or in adversarial environment.
One of the future directions of this work would be toprovide theoretical guarantees under reasonable assumptions, similarly to generalization bounds forfactored representations for fast domain adaptation [13].
Based on the application, there might be different assumptions or inductive biases that might beconsidered reasonable.
Since our work leverages insights from recent causality literature [9], wealso inherit the same inductive bias in terms of assuming there is an underlying causal structurethat is time-invariant throughout the non-stationarity.
In our current method, this causal structureis estimated in the model estimation phase as one of the first steps of the algorithm.
After thisestimation, there can still be changes in the functional dependencies between the various components,but we assume there are no new edges/causal relations between components that were previouslydisconnected, or new forms of non-stationarity, in terms of connections between the change factorcomponents and the state dimensions or reward.
For example, if we consider Halfcheetah-v3 with achange in gravity and estimate a model that can handle this type of non-stationarity, our method willnot be able to perform well under a new type of non-stationarity (e.g. change of wind forces) thatwas not observed during model estimation.
An interesting extension of our work would be designingways to efficiently detect changes in the causal structure and adapt the model.
As is the case with other works in causal discovery, we also make some standard assumptions torecover the causal graph from time series observational data.
In particular, we assume that there areno other unobserved confounders, except for the change factors, and that there are no instantaneouscausal effects between the state components, which is implied by our definition of an FN-MDP andits Dynamic Bayesian Network.
In practical applications, this means that we are able to measure allthe relevant causal variables and we are measuring them at a rate that is faster than their interaction.
Additionally, in our identifiability proofs, we assume the causal Markov and faithfulness assumptions[44], which provide a correspondence between conditional independences and d-separations in thegraph.
The faithfulness assumption can be violated for example in case of deterministic relations,thus requiring a careful modelling of the system.
In general, if these assumptions are violated, thecausal structure we learn might be incorrect, and therefore the factored representation might notbe beneficial to adapt to non-stationarity.
A future direction would be to relax some of the currentassumptions to provide a more realistic and flexible framework for factored non-stationary RL.
B Proofs and causality backgroundB.1 PreliminariesB.1.1 Dynamic Bayesian networks.
Dynamic Bayesian networks (DBNs) [45] are the extensions of Bayesian networks (BN), whichmodel the time-dependent relationship between nodes (See an example in Fig.
A1(b)).
The unfolded
DBNs can be represented as BNs.
The variables in DBNs are in discrete time slices and dependent onvariables from the same and previous time slices.
Hence, the DBNs can model the stationary processrepeated over the discrete time slices.
[Equation]
x t − x t − x t − x t − x tx tx tx t
[Figure]
B.1.2 Markov and faithfulness assumptions
Given a directed acyclic graph G =
( V , E ) , where V is the set of nodes and E is the set of directededges, we can define a graphical criterion that expresses a set of conditions on the paths.
Definition 2 (d-separation [46]) .
A path p is said to be blocked by a set of nodes Z ⊆ V if and onlyif (1) p contains a chain i → m → j or a fork i ← m → j such that the middle node m is in Z , 571 or (2) p contains a collider i → m ← j such that the middle node m is not in Z and such that nodescendant of m is in Z .
Let X , Y , and Z be disjunct sets of nodes.
Z is said to d-separate X from Y(denoted as X ⊥ d Y | Z ) if and only if Z blocks every path from a node in X to a node in Y .
Definition 3 (Global Markov Condition [44, 46]) .
A distribution P over V satisfies the global
Markov condition on graph G if for any partition ( X, Z, Y ) such that X is d-separated from Y given
Z , i.e. X ⊥ d Y | Z the distribution factorizes as:
[Equation]
In other words, X is conditionally independent of Y given Z , which we denote as X ⊥⊥ Y | Z .
Definition 4 (Faithfulness Assumption [44, 46]) .
There are no independencies between variablesthat are not entailed by the Markov Condition.
If we assume both of these assumptions, then we can use d-separation as a criterion to read all ofthe conditional independences from a given DAG G .
In particular, for any disjoint subset of nodes
X,Y,Z ⊆ V : X ⊥⊥ Y | Z ⇐⇒ X ⊥ d Y | Z .
B.2 AdaRL summary
Our work extends the factored representation for fast policy adaptation across domains introduced in
AdaRL [13], which we summarize here.
While Huang et al.
[13] propose a general framework thatcan be applied to both MDPs and POMDPs, in this work we focus on MDPs, so we only present thesimplified version of AdaRL for MDPs.
The simplified AdaRL setting considers n source domainsand n ′ target domains.
The state at time t is represented as s t =
( s 1 ,t , · · · , s d,t ) ⊤ ∈ S d , while 589 a t ∈ A m is the executed action and r t ∈ R is the reward signal.
The generative process of the 590 environment in the k -th domain with k = 1 , ... , n + n ′ can be described in terms of the transitionfunction for each dimension i = 1 , ... , d of s t as:
[Equation]
where ⊙ denotes the element-wise product.
The binary mask c s (cid:41) s i ∈ { 0 , 1 } d represents which of 593 the state components s j,t − 1 are used in the transition function of s i,t .
Similarly, c a (cid:41) s i ∈ { 0 , 1 } m is a 594 mask that indicates whether the action directly affects s i,t .
The change factor θ s k ∈ R p is the onlyparameter that depends on the domain k in Eq.
and it encodes any change across domains in the dynamics.
The binary mask c θ k (cid:41) s i ∈ { 0 , 1 } p represents which of the θ s k components influence the 597 s i,t .
Finally, ϵ si,t is an i.i.d.
random noise.
Similarly the reward function is modeled as:
[Equation]
A
[Equation]
where c s (cid:41) r i ∈ { 0 , 1 } d , c a (cid:41) s i ∈ { 0 , 1 } m , and ϵ rt is an i.i.d.
random noise.
The change factor θ rk ∈ R qis the only parameter that depends on the domain k in Eq.
A2 and it encodes any change in the rewardfunction.
In this simplified setting, the binary masks c · (cid:41) · can be seen as indicators of edges in a
Dynamic Bayesian Network (DBN).
Under Markov and faithfulness assumptions, i.e., assumingthe conditional independences in the data and d-separations in the true underlying graph coincide,the edges in the graph can be uniquely identified.
This means one can learn the true causal graphrepresenting jointly all of the environments, even if the change parameters are latent.
In the general AdaRL framework, the representation is learned via a combination of a state predictionnetwork (to estimate the various f i ) and a reward prediction network (to estimate h ).
All binary masksc · (cid:41) · and change factors θ · k are trainable parameters.
All change factors θ · k are assumed to be constantin each domain k .
If the inputs are pixels, another encoder is added to infer the symbolic states,forming a Multi-model Structured Sequential Variational Auto-Encoder (MiSS-VAE) .
MiSS-VAEleverages the generative modeling to learn the data generation process in RL system with multipledomains.
All change factors in MiSS-VAE are modeled by constants in each domain.
This setting isnot suitable for non-stationary RL where the change factors evolve over time.
In general, not all of the dimensions of the learned state and change factor vectors are useful inpolicy learning.
Huang et al.
[13] select a subsection of dimensions which are essential for policyoptimization.
Leveraging the learned representation as a DBN, we can select compact states and CFsas having a directed path to a reward:
[Equation]
B.3 Proofs
Proposition 3 (Identifiability with observed change factors) .
Suppose all the change factors θ s t andθ r t are observed, i.e., Eq. 1-3) is an MDP.
Under the Markov and faithfulness assumptions, all thebinary masks C · (cid:41) · are identifiable.
Proof.
We construct the graph in the Factored Non-stationary MDP as a dynamic Bayesian network(DBN)
G over the variables V MDP
=
{ s 1 ,t − 1 , . . .
, s d,t − 1 , s 1 ,t , . . . , s d,t , a 1 ,t − 1 , . .
. , a m,t − 1 , r t − 1 } , 624 and the change factors V
[Equation]
In this setting, we can always the correct causal graph under the Markov and faithfulness assumptions.
We rewrite the time index of r t − 1 as r t .
We could rewrite also the change factors time indices by shifting them back in time by one step,i.e. t → t − 1 , and this would allow us to model the whole setup as a Markov DBN without anyinstantaneous effect.
In this setup, we can leverage existing results to show that the true causal graphis asymptotically identifiable from conditional independences in a time-series without any unobservedconfounders or instantaneous effects [47].
In order to make our proof clearer, we instead show step bystep how we can recover the parts of the graph related to V MDP , to V θ and finally the connectionsbetween them.
In this setting, there are no instantaneous effects except for the change factors, and the only causalparents for a variable that is not a change factor at time t can be in the previous time-step t − 1 .
Inparticular, as in usual MDPs, these are the only allowed edges:
state dimension s i,t − 1 at time t − 1 to state dimension s j,t at time t , for i, j ∈ { 1 , . . .
, d }
(this includes the case in which i = j );
action dimension a k,t − 1 at time t − 1 to state dimension a j,t at time t , for j ∈ { 1 , . . .
, d } , k ∈{ 1 , . . .
, m } ;
state dimension s i,t − 1 at time t − 1 to reward r t at time t , for i ∈ { 1 , . . .
, d } , k ∈{ 1 , . . .
, m } ;
action dimension a k,t − 1 at time t − 1 to reward r t at time t , for k ∈ { 1 , . . .
, m } ;
and in addition we have some extra knowledge about the allowed edges to and from change factors,as expressed in our generative model in Equations (1-3):
transition change factor dimension θ si,t at time t to state dimension s j,t at time t , for
i ∈ { 1 , . . .
, p } , j ∈ { 1 , . . .
, d } ;2.
transition change factor dimension θ si,t at time t to transition change factor dimension θ sj,t +1at time t + 1 , for i, j ∈ { 1 , . .
. , p } ;3.
reward change factor dimension θ ri,t at time t to reward change factor dimension θ rj,t +1 attime t + 1 , for i, j ∈ { 1 , . .
. , q } ;
For the reward change factors, we assume they are fully connected to the reward, as shown in Eq.
Using this background knowledge, we can learn the edges from any other variable V i,t − 1 → V j,t , for 654 V i,t − 1 , V j,t ∈ V MDP , just by checking if V i,t − 1 ̸⊥⊥ V i,t | θ rt , θ st , s t − 2 , a t − 2 .
This dependence implies 655 that the variables are d-connected, under the Markov and faithfulness assumptions.
Except for adirect edge, there is no other possible path through the graph unrolled in time, since we have blockedall influence of time-step t − 2 and earlier, and the paths through future time-steps contain colliders.
So this means that V i,t − 1 and V j,t are adjacent, and in particular the edge follows the arrow of time,from t − 1 to t .
In our setting, we do not model the actions at time t , since we assume they are notcaused by any other variable.
We also never need to condition on the reward to check if two variablesare adjacent, since it’s always on a collider path.
This means we are able to learn the following binarymasks:
state dimensions to state dimensions C s (cid:41) s ;2. action dimensions to state dimensions C a (cid:41) s ;3.
state dimensions to reward c s (cid:41) r
action dimensions to reward c a (cid:41) r ;
To learn the edges between any change factor component to another change factor component, i.e.,
V i,t − 1 → V j,t for V
i,t − 1 , V j,t ∈ V θ , we can just check if V i,t − 1 ̸⊥⊥ V i,t | θ rt − 2 , θ st − 2 , since states, 669 actions and rewards can never be parents of the change factors, so we do not need to condition onthem to close any path through the earlier time-steps.
We also assumed that change factor componentsfollow a Markov process, so they do not have instantaneous effects towards each other.
This meanswe are able to learn the following masks:
transition change factor dimensions to transition change factor dimensions C θ s (cid:41) θ s ;2.
reward change factor dimensions to reward change factor dimensions C θ r (cid:41) θ r ;
Finally to learn the edges between the change factors V i,t ∈ V θ
and the other variables V j,t ∈ V MDP , 676 we can just check if V i,t ̸⊥⊥ V j,t | s t − 1 , a t − 1 , θ rt − 1 , θ st − 1 and if this is true we can learn the edge 677 V i,t → V j,t .
This means we are able to learn the mask:
transition change factor dimensions to state dimensions
C θ s (cid:41) s ;
All of these results together show that if we have the Markov and faithfulness assumption, i.e. theconditional independence tests return the true d-separations in the graph, and we observe the changefactors, we are able to completely identify the graph of the FN-MDP represented by the binarymasks.
Proposition 4 (Partial Identifiability with latent change factors) .
Suppose the generative processfollows Eq.
1-3 and the change factors θ s t and θ r t are unobserved.
Under the Markov and faithfulnessassumptions, the binary masks C s (cid:41) s , C a (cid:41) s , c s (cid:41) r and c a (cid:41) r .
Moreover, we can identify which statedimensions are affected by θ s t and if the reward function changes.
Proof.
In this case, the MDP is non-stationary, since we cannot observe the latent change factors.
We assume that we can represent the latent change factors as a smooth function of the observed timeindex t .
This assumption is called pseudo-causal sufficiency in previous work [9].
We can then usethe time index t as a surrogate variable to characterize the unobserved change factors, since at eachtime-step their value will be fixed.
We again consider a
DBN
G MDP ∪{ t } over the variables
V MDP =
{ s 1 ,t [Equation]
, s d,t [Equation] 1 , s 1 ,t , . . .
, s d,t , a 1 ,t − 1 , . .
. , a m,t − 1 , r t − 1 } and the time index t .
We rewritethe time index of r t − 1 as r t .
Note that we do not represent the change factors in this DBN, but we can capture their effect through t since they are assumed to be deterministic smooth functions of thetime index.
We can then reuse the results by Huang et al.
[9] (Theorem 1) in which under the pseudo-causalsufficiency (Assumption 1 in that paper) and the Markov and faithfulness assumption (Assumption 2),one can asymptotically identify the true causal skeleton (i.e. the adjacencies) in the graph G MDP ∪{ t } )through conditional independence tests.
In particular for any V i , V j ∈ V MDP ∪ { t } , V i and V j are 701 not adjacent if there exists a subset of the unrolled graph V k of V MDP ∪{ t } \ { V i , V j } in G MDP ∪{ t } 702 such that V
i ⊥⊥ V j | V k .
In particular, we can also focus on the tests that were used for the proof of theprevious Proposition, by just substituting the change factors with the time index t .
The skeleton isthe undirected version of the causal graph, so this result only tells us that can asymptotically get thecorrect undirected edges, but not their orientations.
Fortunately, in our setting we have some additional background knowledge that allows us to orient allthe existing edges.
In particular, as in usual MDPs, these are the only allowed edges:
state dimension s i,t − 1 at time t − 1 to state dimension s j,t at time t , for i, j ∈ { 1 , . . .
, d }(this includes the case in which i = j );
action dimension a k,t − 1 at time t − 1 to state dimension a j,t at time t , for j ∈ { 1 , . . .
, d } , k ∈{ 1 , . . .
, m } ;3.
state dimension s i,t − 1 at time t − 1 to reward r t at time t , for i ∈ { 1 , . . .
, d } , k ∈{ 1 , . . .
, m } ;4.
action dimension a k,t − 1 at time t − 1 to reward r t at time t , for k ∈ { 1 , . . .
, m } ;
This means that, for example, we cannot have a variable at time t causing a variable at time t − 1 .
Therefore if two variables V i,t − 1 and V j,t are adjacent, we already know that the direction of thatedge will be V i,t − 1 → V j,t .
This also implies that the following binary masks are identifiable (i.e. noedge remains unoriented):
state dimensions to state dimensions C s (cid:41) s ;2. action dimensions to state dimensions C a (cid:41) s ;3.
state dimensions to reward c s (cid:41) r4.
action dimensions to reward c a (cid:41) r ;
These represent all of the edges in G MDP .
We can also learn the edges from t to V MDP (byconstruction we assume the opposite direction is not possible), which will represent the effect of thechange factors, as we show in the following.
Since t inherits all of the children of the latent change factors in G , we can further show that ifs i,t ⊥⊥ t | s t − 1 , a t − 1 in G , then none of the latent change factor dimensions θ sj,t affect s i,t , i.e.,
[Equation]
= 0 .
Intuitively, this means that the distribution of s i,t only dependson s t − 1 and a t − 1 , and not on the timestep t , or in other words, this distribution is stationary.
Underthe same principle, if r t ⊥⊥ t | s t , a t , then the reward is stationary.
C Details on experimental designs and results
C.1 MuJoCo
We modify the Half-Cheetah environment into a variety of non-stationary settings.
Details on thechange factors are given as below.
Changes on dynamics.
We change the wind forces f w in the environment.
We consider thechanging functions can be both continuous and discrete.
• Continuous changes: f w t = 10 + 10 sin(0 .
005 · t ) , where t is the timestep index;•
Discrete changes: (1) Across-episode:a. Sine function: f w = 10 + 10 sin(0 .
5 · i )
b. Damping-like function: f w = 10 + 3 · (1 . 01) −⌈ i/ 10 ⌉ sin(0 .
5 · i )c.
Piecewise linear function f w = 5 + 0 .
02 · ∥ i − 1500 ∥ ; , where i is the episode index.
743 (2) Within-episode: f w = 10 + 10 sin(0 .
· ⌊ t/ 10 ⌋ ) , where t is the timestep index.
Changes on reward functions.
To make the rewards be non-stationary, we change the target speedv g in each episode.
To make the learning process stable, we only consider the discrete changesand the change points are located at the beginning of each episode.
The changing function isv g = 1 .
+ 1 .
5 sin(0 .
2 · i ) , where i denotes the episode index.
Changes on both dynamics and rewards.
We consider a more general but challenging scenario,where the changes on dynamics and rewards can happen concurrently during the lifetime of theagents.
We change the wind forces and target speed at the beginning of each episode.
At episode i ,the dynamics and reward functions are:
[Equation]
Here, w is the non-stationary degree.
We consider multiple values of w in our experiments.
In
Fig. 2(d), w = 0 .
In Fig. 2(h), w is the value of non-stationary degree.
C.2 Sawyer benchmarks
In Sawyer-Reaching, the sawyer arm is trained to reach a target position s gt .
The reward r t is thedifference between the current position s t and the target position r t = − ∥ s t − s g ∥ 2 .
In this task, 757 we cannot directly modify the dynamics in the simulator, so consider a reward-varying scenario,where the target location changes across each episode following a periodic function.
In Sawyer-Peg,the robot arm is trained to insert a peg into a designed target location s g .
The reward function isr t
= I ( ∥ s t − s g ∥ 2 ≤ 0 . 05) .
We change the target location in Sawyer reaching task.
The target location s gt is given as below:
[Equation]
where i is the episode index.
For Sawyer-Peg task, the target location s g changes at each episode.
The parameters in each dimension of s g is randomly sampled at episode i as below:
• x_range_1: (0 . 44 , 0 . 45) ;• x_range_2: (0 . 6 , 0 . 61) ;• y_range_1: ( − 0 .
08 , − 0 .
;• y_range_2: (0 . 07 , 0 . 08) ;
C.3 Minitaur benchmarks
We consider both the changes on dynamics and reward functions.
Changes on dynamics.
We change the mass of taur m in the environment.
Specifically, we considerboth the continuous and discrete changes.
• Continuous changes: m t = 1 .
+ 0 .
75 sin(0 .
005 · t ) ;• Discrete and within-episode changes: m t = 1 .
+ 0 .
75 sin(0 .
3 · ⌊ t/ 20 ⌋ )
Changes on both dynamics and reward functions.
We also consider a case where both thedynamics and reward functions change at the beginning of each episode.
We change the target speedof minitaur to introduce the non-stationarity of reward functions.
The change functions are givenbelow:
[Equation]
C.4 Full results
Fig.
A2 and A3 give the smoothed learning curves on average rewards over 10 runs versus timestepsin Half-Cheetah and Minitaur experiments.
Fig.
A4, A5, and A6 give the full results on averagerewards over 10 runs versus timesteps in Half-Cheetah, Sawyer, and Minitaur experiments.
Table C.4shows the average final rewards over 10 runs for all experiments.
Fig.
A7 demonstrates the rewards on
Half-Cheetah with different non-stationary degrees on multi-factor changing scenario.
Fig.
A8 givesaverage rewards on different benchmarks with varying numbers of latent features with all evaluatedapproaches.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
Oracle
SAC
LILAC
TRIO
VariBAD
Ours
Half
[Table]
Cheetah:
[Table]
EP
[Table]
Half
[Table]
Cheetah:
[Table]
[Figure] [Figure]
TRIO
[Figure] [Figure]
[Figure] [Figure]
R e w a r d s
Ours
[Figure]
TRIO
[Figure]
- episode, and within-episode changes, respectively.
(D) and (R) represent changes on dynamics and reward functions, respectively.
Best viewed in color.
To verify the effectiveness of each component in our proposed framework, we consider the followingablation studies:
• Without smoothness loss ( L smooth );
• Without structural relationships ( C · (cid:41) · );• Without compact representations ( s min , θ min );
• Without sparsity losses ( L sparse );
[Figure]
• Without reward or state prediction losses ( L pred-rw , L pred-dyn );
[Figure]
• Without the disentangled design of CF inference networks for dynamics ( q sϕ
) and
rewards ( q rϕ ).
[Figure]
Specifically, we use one CF inference encoder and the mixed latent space of θ s and θ r in this setting.
[Figure]
As shown in Fig.
A9, all the studied components benefit the performance.
Furthermore, FANS-RLcan still outperform the strong baselines even without some of the components.
[Figure]
C.6 Visualization on the learned change factors
[Figure]
Fig.
A10 gives the visualization on the learned θ in Half-Cheetah.
Fig.
A10(a-b) show the pairwise
Euclidean distance between learned θ r and the axes denote the values of change factors on rewards.
Similarly, Fig.
A10(c-d) displays the Euclidean distance between learned θ s and the values of changefactors on dynamics.
The results suggest that there is a positive correlation between the distance oflearned θ versus the true change factors.
This can verify that θ can capture the change factors in thesystem.
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure]
D Details on the factored adaptation framework
D.1 Algorithm pipelines of FN-VAE
Alg.
A1 gives the full pipeline of FN-VAE.
D.2 The framework dealing with discrete and across-episode changes
Alg.
A2 gives the extended framework for handling both across- and within-episode changes innon-stationary RL, respectively.
The major difference between Alg. A2 and Alg. 1 is that weonly infer θ using via CF dynamics networks at change points.
Furthermore, we also adjust theobjective functions of FN-VAE to fit the discrete changes.
At timestep t in episode n , wheret m ≤ (cid:0) ( n − 1) · H + t (cid:1) <
t m +1 , we have:
• Prediction and reconstruction losses:
[Equation] [Equation]
• KL loss:
[Equation]
Output: ϕ , α 1 , α β β 2 , γ for i = 1 , 2 , . . .
, E do
Randomly sample a batch of trajectories τ 0: k in τ # Infer the latent change factors for j = s, r do
Infer µ ϕ j ( τ 0: k ) and σ 2 ϕ j ( τ 0: k ) using q ϕ j
Infer µ γ j ( θ j 0: k ) and σ 2 γ j ( θ j 0: k ) using p γ j
Sample θ j 0: k ∼ N (cid:16) µ ϕ j ( τ 0: k ) , σ 2 ϕ j ( τ 0: k ) (cid:17) end for
Reconstruct and predict ˆ s 0: k , ˆ s 1: k , ˆ r 0: k , ˆ r 1: k using p α 1 , p α 2 , p β 1 , and p β 2 # Update the FN-VAE model
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
←
[Table]
[Equation]
[Table]
[Equation]
[Table]
) if updateG then
G ← G − λ
[Equation]
∇
[Equation]
( L
[Equation]
-dyn
[Equation]
L
[Equation]
[Table]
+
[Table]
[Equation]
+ L sparse ) end if end for
• Sparsity loss:
[Equation]
• Smoothness loss:
[Equation]
The total loss L vae = k 1 L rec + k
[Equation]
L pred − k
[Equation]
L KL − k 4 L
[Equation]
− k 5 L smooth , where k 1 , k 2 , k 3 , k 4 , and 821 k 5 are adjustable hyper-parameters to balance the objective functions.
D.3 The framework dealing with raw pixels
We augment the generative process in Eq.
1-3 with the generative process of observation.
[Equation]
where u is a non-linear function and i = 1 , .
, d . c s (cid:41) o :=
[ c s (cid:41) o i ] di =1 .
ϵ ot is an i.i.d.
random noise.
To learn the u i , we model the states as the latent variables in FN-VAE.
Fig.
A11 gives the modified
FN-VAE dealing with raw pixels, where the states are also in the latent space.
Different from theoriginal FN-VAE, we incorporate state inference networks and state dynamics networks.
Moreover,we reconstruct and predict the current and future observations using the observation decoder.
Detailedobjective functions are given below.
• Prediction and reconstruction losses
[Equation]
1: Init: Env; VAE parameters: ϕ = ( ϕ s , ϕ r ) , α = ( α 1 , α 2 ) , β = ( β 1 , β 2 ) ,
; Mask matrices: C
Policy parameters: ψ ; replay buffer: D ; Number of episodes: N ; Episode horizon: H ; Change index ˜ t
= { t 1 , . . .
, t M } ; m = 0 .
2: Output: ϕ , α 1 , α 2 , β 1 , β 2 , γ , ψ 3: # Model initialization 4: Collect multiple trajectories τ = { τ 10: k , τ 20: k , . . .
} with policy π ψ from Env; 5: Learn an initial VAE model on τ (Alg.
A1) 6:
Identify the compact representations s min and change factors θ min based on C · (cid:41) · 7:
# Model estimation & policy learning
8: for n = 0 , . . .
, N − 1 do 9: for t = 0 , . . .
, H − 1 do 10:
Observe s t from Env; 11:
# Estimating
latent change factors
12: if n · ( H − 1) + t ∈ ˜ t then 13: m ← m + 1 14: for j = s, r do 15:
Infer µ γ j ( θ jt m − 1 ) and σ 2 γ j ( θ jt m − 1 ) via p γ j 16:
Sample θ jt m ∼ N (cid:16) µ γ j ( θ jt m
[Equation]
1 ) , σ 2 γ j ( θ jt m − 1 ) (cid:17) 17: end for 18: end if 19:
Generate a t ∼ π ψ ( a t
| s mint , θ mint m ) 20:
Receive r n,t from Env 21:
Add ( s t , a t , r t , θ st m , θ rt m ) to replay buffer D ; 22:
Extract a trajectory with length k from D ; 23:
Learn VAE (Alg. A1) with updateG=
False; 24:
Sample a batch of data from D 25:
Update policy network parameters ψ 26: end for 27: end for
[Equation]
loss
[Equation]
where τ 0: t = { o 0 , r 0 , o 1 , r
, o t , r t } .
• Sparsity loss
[Equation]
• Smooth loss
[Equation]
c s → r c a → r θ rt p α̂ o t + ̂ r t +1 p β̂ r t p βp γ s ̂ o t p α
C θ r → θ r a t + c s → o θ s t
C θ s → θ s s t θ rt + p
α̂ o t + ̂ r t +2 p
β̂
r t +1 p βp αc s → o c s → r c a → r a t + a t + ̂ o t + θ s t
+1 s t
+ p γ r
C s → s , θ s t , a t p γ o a t q sϕ h s t o t a t r t q rϕ h rt q oϕ o t +1 a t +1 r t +
[Figure]
D.4 Hyper-parameter selection
D.4.1 Factored model estimation
Input with symbolic states.
In Half-Cheetah, Sawyer-Reaching, and Minitaur the symbolic statesare observable.
For the CF dynamic networks, we use 2 -layer fully connected networks.
The numberof neurons is 512 .
For CF inference networks, we use 2 -layer fully connected networks, wherethe number of neurons is 256 , followed by the LSTM networks with 256 hidden units.
The initiallearning rates for all losses are set to be 0 .
1 with a decay rate 0 .
The batch size is 256 and thelength of time steps is equal to the horizon in each task.
The number of RNN cells is 256 .
Thedecoder networks are 2 -layer fully connected networks.
The number of neurons is 512 .
Input with raw pixels.
In Saywer-Peg, we directly learn and adapt in non-stationary environmentswith raw pixels observed.
Different from other experiments, we use the architecture described in
Fig.
A11.
At timestep t , we stack 4 frames as the input o t .
A 5 -layer convolutional networks is usedto extract the features of the trajectories of observations and rewards.
The layers have 32 , 64 , 128 ,256 , and 256 filters.
And the corresponding filter sizes are 5 , 3 , 3 , 3 , 4 .
The observation decodersare the transpose of the convolutional networks.
Then the extracted features are used as the inputof LSTM networks in state inference networks.
The state inference networks and state dynamicnetworks share the same architectures with the CF inference and dynamics networks, respectively.
We use the same CF inference networks, CF dynamics networks, and reward decoders with those incases with symbolic states as input.
The number of latent features is 40 .
[Table]
Balancing parameters in losses.
For all experiments:
• All w · are set to be 0 . 1 ;
[Table]
• Weights of the reconstruction loss: k 1 = 0 . 8 ;
[Table]
• Weights of the prediction loss: k 2 = 0 . 8 ;• Weights of KL loss: k 3 = 0 .
5 ;• Weights of sparsity loss: k 4 = 0 .
1 ;• Weights of smooth loss: k 5 = 0 .
Steps (M) −−−
R e w a r d s
Ours
ZeUS (best)
ZeUS (ﬁnal)
Meld (best)
Meld (ﬁnal)
CaDM (best)
CaDM (ﬁnal)
HyperDynamics (best)
HyperDynamics (ﬁnal)
We use the automatic weighting method in [15] to learn the weights for k 1 , . .
, K 5 and grid searchfor w 1 , . . .
, w 7 .
Model initialization.
Table A2, A3, and A4 provide the settings of learning the model initialization.
CONT (D)
A-EP (D)
W-EP (D)
A-EP (R)
A-EP (R+D)
# trajectories     # steps in each episode     # episodes
Sawyer-Reaching
Sawyer-Peg
[Table]
[Table] [Table]
episodes
CONT (D)
W-EP (D)
A-EP (
[Table]
+
[Table]
[Table] [Table]
# episodes
D.4.2 Policy learning
In the Half-Cheetah, Sawyer-Reaching, and Minitaur experiments, we follow the learning ratesselection for policy networks in [7].
In Sawyer-Peg, for both actor and critic networks, we use 2 -layerfully-connected networks.
The number of neurons is 256 .
For all experiments, we use standard
Gaussian to initialize the parameters of policy networks.
The learning rate is 3 e − 4 .
The relay buffercapacity is 50 , 000 .
The number of batch size is 256 .
Details on TRIO and VariBAD For TRIO and VariBAD, we meta-train the models (batch size:5000, # epochs: 2 for all experiments) and show the learning curves of meta-testing.
The tasksparameters for meta-training are uniformly sampled from a Gaussian distribution.
For all approaches,we use the same set of hyper-parameters for policy optimization modules (i.e., SAC).
For the latentparameters in TRIO, we follow the original paper where the latent space from the inference networkis projected to a higher dimension.
The number of latent parameters for TRIO is the same as thosein other approaches (Half-Cheetah and Minitaur: 40 , Sawyer-Reaching: 20 ).
We compared with
TS-TRIO, with the kernels set as in the original implementation.
E Experimental platforms and LicensesE.1 Platforms
All methods are implemented on 8 Intel Xeon Gold 5220R and 4 NVidia V100 GPUs.
E.2 Licenses
In our code, we have used the following libraries which are covered by the corresponding licenses:
• Tensorflow (Apache License 2.0),
• Pytorch (BSD 3-Clause "New" or "Revised" License),
• OpenAI Gym (MIT License),
• OpenCV (Apache 2 License),
• Numpy (BSD 3-Clause "New" or "Revised" License)
• Keras (Apache License).
We plan to release our code under the MIT License.
Implementations can be found at this anonymous