Factored Adaptation for Non-Stationary Reinforcement Learning
Authors: Fan Feng
Biwei Huang
Kun Zhang
Sara Magliacane
Venue: NeurIPS 2022 Accept (2022)
Keywords: non-stationarity
non-stationary RL
causal RL

Abstract: Dealing with non-stationarity in environments (e.g., in the transition dynamics) and objectives (e.g., in the reward functions) is a challenging problem that is crucial in real-world applications of reinforcement learning (RL). While most current approaches model the changes as a single shared embedding vector, we leverage insights from the recent causality literature to model non-stationarity in terms of individual latent change factors, and causal graphs across different environments. In particular, we propose Factored Adaptation for Non-Stationary RL (FANS-RL), a factored adaption approach that learns jointly both the causal structure in terms of a factored MDP, and a factored representation of the individual time-varying change factors. We prove that under standard assumptions, we can completely recover the causal graph representing the factored transition and reward function, as well as a partial structure between the individual change factors and the state components. Through our general framework, we can consider general non-stationary scenarios with different function types and changing frequency, including changes across episodes and within episodes. Experimental results demonstrate that FANS-RL outperforms existing approaches in terms of return, compactness of the latent state representation, and robustness to varying degrees of non-stationarity.


Factored Adaptation for Non-stationary Reinforcement Learning
Fan Feng', Biwei Huang", Kun Zhang?:3 Sara Magliacane
City University of Hong Kong *Carnegie Mellon University Mohamed bin Zayed University of Artificial Intelligence
'University of Amsterdam MIT-IBM Watson Al Lab
(ffeng1017, sara, aaa.gg/0...aaaaaaaaaaaaaaaeeeee mag.a com, biveih@andres. cmu.edi edu, kunz @cmu edu
Abstract
Dealing with non-stationarity in environments (e.g.. in the transition dynamics) and objectives (e.g., ., in the reward functions) is challenging problem that is crucial in real-world applications of reinforcement learning (RL). While most current approaches model the changes as single shared embedding vector, we leverage insights from the recent causality literature to model non-stationarity in terms individual latent change factors, and causal graphs across different environments. In particular, we propose Factored Adaptation for Non-Stationary RL (FANS-RLA RL),, factored adaption approach that learns jointly both the causal structure in terms of a factored MDP, and a factored representation of the individual time-varying change factors. We prove that under standard assumptions, we can completely recover the causal graph representing the factored transition and reward function, as well as a partial structure between the individual change factors and the state components. Through our general framework, we can consider general non-stationary scenarios with different function types and changing frequency, including changes across episodes and within episodes. Experimental results demonstrate that FANS-RI RI. outperforms existing approaches in terms of return, compactness the latent state representation, and robustness to varying degrees of non-stationarity.
1 Introduction
Learning stable policy under non-stationary environments is long-standing challenge in Reinforce- ment learning (RL) [1, 2, 3]. While most RL approaches assume stationarity, in many real-world applications of RL there can be changes in the dynamics or the reward function, both across dif- ferent episodes and within each episode. Recently, several works adapted Meta-RL RL methods to learn sequences of non-stationary tasks [4, 5]. However, the continuous MAML [6] adaptation for non-stationary RL [4] does not explicitly model temporal changing components, while TRIO [5] needs to meta-train the model latent variable models to directly model the change factors in the environment in a shared embedding space. In particular, they consider families of MDPs indexed by single latent parameter. In this paper, we argue that disentangling the changes as separate latent parameters and modeling the process with factored representation improves the efficiency of adapting to stasstiiiinait..
set of non-stationary tasks. LILAC [7] and ZeUS [8] leverage
In particular, we leverage insights from the causality literature [9, 10] that model non non-stationarity in terms of individual latent change factors and causal graphs across different environments. We propose Factored Adaptation for non--tationary RL (FANS-RL), a factored adaptation framework that jointly learns the causal structure of the MDP and factored representation of the individual change factors, allowing for changes at discrete timepoints and continuously varying environments. While we provide specific architecture (FN-VAE), the theoretical framework of FANS-RL can be implemented with different architectures and combined with various RL algorithms. We formalize our setting as
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Factored Non-stationary MDP (FN-MDP), which combines Factored- [11, 12, 13] with latent change factors that evolve in time following Markov process.
We build upon the AdaRL. framework [14], a recently proposed fast adaptation approach. AdaRI. learns factored representation that explicitly models changes (i.e., domain-specific components) in observation, dynamics and reward functions across set of source domains. An optimal policy learnt on the source domains can then be adapted to new target domain simply by identifying low-dimensional change factor, without any additional finetuning. FANS-RL extends AdaRL from the stationary case with constant change factors general -ta-stationary framework. Specifically, FANS-RI RL learns the Now-dimensiona and time-evolving representations 0 and 0; that fully capture the non- uonsstaiionarity of dynamies and rewards, allowing for continuous and discrete changing functions, both within -pissode and across-episse across-episode. Our main contributions can be summarized as:
We formalize FN-MIDSs, 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-MI FN-MDPs. We evaluate FANS-RI RL on simulated benchmarks for continuous control and robotic manip- ulation tasks and show it outperforms the state of the art on the return, compactness of the latent space representation and robustness to varying degrees of hon-stationarity.
2 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 a factored MDPs [11, 12, 13] with latent change factors that evolve over time following Markov process. 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, O", O" 7% G. RR, .P.)), where S is the state space, A the action space, @a the space of the change factors for the dynamics, @" the space of the reward change factors and the discount fac- tor. We assume G is a Dynamic Bayesian Network over {$1,t:...,$1,, 01,ts 0m, Tt, fi....... ... 6p. .r....... b, where d, m, p, and q are the dimensions of states, action, change factors on dynamics and reward, respectively. We define the factored state transition distribution I,, as:
d
II P. (s....lla(sss.) ==1
where pa(st.) denotes the causal parents of 81, in G, which are a subset of the dimensions of 81-1- a:-1 and 0². Note that the action a:- is a vector of m dimensions in our setting. We assume given initial state distribution (90)).. Similarly, we define the reward function R as function of the parents of rt in G. i.e., R(s,, ac, 0,) = R(pa(r))), where pa(r.) are a subset of dimensions of St, a,, and O7. We define the factored latent change factors transition distributions Fo- and Por as:
9
Po- (0,000001 = II Po (0), j=1
Po- (0-
II Ter (0%, (Pa( k=1
where pa(O; (),,) are a subset of the dimensions of 0; while pa(O%.) are subset of dimensions of 0-- We assume the initial distributions P (0)) and Pa- are given.
We show an example DBN representing the graph G of an FN-MDP in Fig. 1(a). Since interested in learning the graphical structure of the FN-MDP, as well as identifying the values of the latent change factors from data, we describe a generative process of an FN-MDP environment. In particular, we assume that the graph G is time-invariant throughout the non-station,rity, and there are no unobserved confounders and instantaneous causal effects in the system. We will learn set of binary masks et and - that are the indicators for edges in G.
Generative environment model. We adapt the generative model in AdaRL. [14] across & different domains to a time-varying setting on a single domain. We assume the generative process of the environment timestep f in terms of the transition function for each dimension i 1,... d of s; is:

---

Po,
Pau
K
C C C
PA.
PA.
Figure 1: (a). graphical representation of FN-MDP. For readability, we only illustrate subsection of dimensions of states, actions, and latent change factors. The shaded variables are unobserved; (b). The architecture of FN-VAE, which learns the generative model, explained in Sec. 3.
81, = Jelc, + 8t-1, Cr o aa-1, 8..8 o ((,,
where @ the element-wise product, f. are non- non-linear functions and c,, is an i.i.d. random noise. The binary mask c²- {0, 1 represents which state components 8j,t1 are used in the transition function of Si,t- Similarly, ct {0, 11mh indicates whether the action directly affects 81,- The change factor 0r E R2 encodes any change the dynamics. The binary mask 0.-. E {0, represents which of the components of 0r influence Si,t We model the reward function
r. = h( 82, car a, 07,
where e E (0, 1]_d e E 10 and ef is an i.i.d. random noise. The change factor 0; E R encodes any change in the reward function. The binary masks c can be seen indicators of edges in the DBN G. In AdaRL, all change factors are assumed to be constant in each domain. Since in this paper we allow the change parameters evolve in time, introduce additional equations:
8%. g" (c. 0 0.11 O 4
0% 9" (ce (c
for = 1,..,,,,3 j I,.... I,... q. and g", and g" are non-linear functions. We assume the binary masks e are stationary across timesteps and so are the ef,g- Ef, ef and ef the i.i.d. random noises. Although e and are stationary, we model the changes in the functions and some changes in the graph structure through 0. For example certain value of O switch off the contribution of some of the state or action dimensions in the reward function, or in other words nullify the effect of some edges in G. Similarly the contribution of the noise distribution modulated via the change factors. On the other hand, this setup does not allow adding edges that are not captured by the binary masks c We group the binary masks in the matrices C = [c.- " C : cef' "1 and C : [c]-=]d 1=1- Similarly, we also group the binary vectors in the dynamics of the latent change factors in matrices - := == 1= and 0 - ][=-- Since latent change factors 0* and 0" follow a Markov process based g" and g", different types of changes by varying the form of g and g", generalizing the approaches in literature. We also model concurrent changes in dynamics and reward, including different types of changes, e.g. a continuous g" and a piecewise-constant g".
each function
be
consider
Compact representations. Huang et al. [14] show that the only dimensions of the state and change factors useful for policy learning are those that eventually affect the reward. These dimensions are called compact representations and are defined as the dimensions of the state and change factors with

---

a path (i.e. a sequence of edges -) to the present or future reward Pe+r for 2 0 in the DBN g: Si,t E smmin + $i,t + r++T for 2 O, and 0; E @min + 02 + + re+r for 2 0
Continuous changes. If g" and g" are continuous, then they can model smooth changes in the environment, including across episodes. While the functions in Eq. 1-3 allow us to model within- episode changes, i.e. changes that can happen only before H where H is the horizon, we also want to model across-episode changes. We use separate time index & that models the agent's 's lifetime. Initially & for & I H. but while we reset = 0 afterwards, continues to grow indefinitely.
Discrete changes. We assume discrete changes happen at specific timesteps and can be represented with a piecewise constant function. In particular we denote change timesteps as i = (f,..... {M) where f, describes a specific timepoint in the agent's ' lifetime time index f. This allows us to model within-episode changes. In this case, we assume that the change happens always at the same steps (f1....,ftm) fm) in each episode, i.c., we assume that i = (t1,.... ..... tm H + tt,.... .... H + tm 2/1 - $1,.... 2/1 + fm. ..., where H is the horizon. We can also model across-episode changes, when we assume the change points only occur at the end of each episode, i.e., i (H, 2H, (H,,....
We extend the result on identifiability of factored MDPs in AdaRL [14] to non-stationary environments. We first assume that 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 process follows Eq. 1-3 and all change factors Of and OF are observed, i.e., Eq. 1-3 is an MDP. Under the Markov and faithfulness ussumptions, iee. conditional independences correspond exactly to d-separations, all binary masks C 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 change factors 0 and or, 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 8,, reward variable rt, and action variable a,. We can also identify which dimensions in Si,t have changes, i.e., we can identify - We formalize this idea in the following (proof in Appendix B):
Proposition 2 (Partial Identifiability with latent change factors). Suppose the generative process follows Eq. 1-3 and the change factors Of and 0e are unobserved. Under the Markov and faithfulness assumptions, the binary masks Ca-a Ca+s c. and carr are identifiable. Moreover, we can identify which state dimensions are affected by 0e 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.
3 Learning the Generative Process in FN-MDPs
There are many possible architectures to learn FN-MDPs through Eq 1-3, We propose FN- VAE, a variational autoencoder architecture described in Fig. 1(b). In FN-VAE, we jointly learn the structural relationships, state transition function, reward function, and transition function of the latent change factors, as described in detail in Appendix Alg. A1. An FN- -N- 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 [15], transition decoders that reconstruct the state dynamics at the time / and predict one step further at + 1, and reward decoders that reconstruct the reward at and predict the future reward at L + 1. We now describe them in detail.
CF inference networks (blue boxes in Fig. 1(b)). The two inference models for latent change factors 90- (0: St: az) and 9a- (0) St. ag, r:) are parameterised by d² and Φ", respectively. To model the dime-dependence of 0f and 0,, we use LSTMs [15] as inference networks. At timestep t, the dynamics change factor LSTM infers qo- (0; (e, az, rı, hī_1), where h;- E R² is the hidden state in the LSTM. Thus we can obtain Mo (To:t) and o (Toz) using go-, and sample the latent changing factor 0z ~ NN (Toz), o²,- (Tu:)),, where Te:t (8o: au, To, 81, a1, ... 82, a, Tz). Similarly, the reward change factor LSTM infers gor (0) 8t; ar, rt, hī where h E TR. L the hidden state, such that we can sample O2 ~ N N( (To:)), 22r (To.))).
CF dynamics network (orange boxes in Fig. 1(b)). We model the dynamics of latent change factors with pa, (0 02, 8 (00) and p.r (0 0f, 0 .r) To ensure the Markovianity of 0; and
4

---

0 we define loss LxL that helps minimize the KL-divergence between qs and P+-
CK Σ KL (0) 8t, ae Tt, (1
(4)
+KL (0: 8t. Or rr, h 10 (0, 0;-1; CO
we assume that the change between Of and 0+ and similarly ff, smooth, we can add smoothness loss Smoothh We provide smooth loss for discrete changes in Appendix D.2.
Smmooth Σ (110 06-11 + +1, - 07-111)
(5)
Transition decoders (purple boxes in Fig. 1(b). We learn an approximation the transition dynamics in Eq. 1 by learning a reconstruction, parameterized by and prediction encoder, parametrized by a2. To simplify the formulas, we define C := ( Cass, C) At timestep the reconstruction encoder Pa. (8t 8;-1, 02; C**) reconstructs the state from current state 8; with sampled 0r. The one-step prediction encoder Pa, ($1+1 St, ag 0+) instead tries to approximate the next state $1+1. We do not use the prediction loss, when the one-step prediction is not smooth. In particular, we do not use it for the last time-step in episode i if there is change happening at the first step in episode (i + 1), since the states new episodes will be randomly initiated. We also do not use the case of discrete changes at the timesteps (t1 .g....fM 1). The loss functions are:
Lree dyn
Raj-as log Par (8t 8 Ot 07; C
(6)
Cyred--yn
Eaj~qs log (8141 at, 03)
Reward decoders (green boxes in Fig. 1(b).. Similarly, we coder P8, (Fr St, ar 0% c carr), parameterized by 31, and a one-step prediction encoder Ps, (Fz+1 8141, 1,+1, 0(), parametrized by 82, approximate the reward function. Similarly to transition decoders, we do not use the one-step prediction loss, if it is not smooth. The losses are:
reconstruction en-
Egg log (relse. d2 0J; e caar)
Cpeed-ew
-gg 444
0411,
Sparsity loss. We encourage sparsity in the binary masks C improve identifiability, by using following loss with adjustable hyperparameters (wi,.., W7), which we learn through grid search.
L wi x w2 w2ll( || wa 0 + wa | ||1 wa ar + wil 00000
(8)
11 + wy 0 1|2
The total loss is Lvxe = k1 (Crec-dyn + Cree-rw) k2 (prod-dyn where (kI, kg): typer--arameterr, which we learn with automatic weighting method [16].
+ (KgCeeddd -pcddd))-kkka kiy Copase kis Csmooth
Learning from raw pixels. Our framework can easily extended image inputs by adding an encoder Ø² to learn the latent state variables from pixels, similar to other works [17, 18, 14]. In this case, our identifiability results do not hold anymore, since we cannot guarantee that we identify the true 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 RI. (FANS-RL), general algorithm that interleaves model estimation and policy optimization, shown in Alg. and Appendix Fig. A9.
5

---

Algorithm 1: Factored Adaptation for non-stationary RL
1: Init: Enw; VAE parameters: o= ( "1"",""),""" )²") 02), 3 (31, (2), 7, Binary masks: C'
Policy parameters: v;; replay buffer: D; Number of episodes: N; Episode horizon: H; Initial 0: 0l and aa Length of collected trajectory: k.
2: Output: VAE parameters: Ø, a. 3, ?; Policy parameters: v/
3: Collect multiple trajectories of length k T = (Tos:: T To.k.... with policy T from Env;
4: Learn FN-VAE from т. including masks C that represent the graph 9 (Appendix Alg, Al)
5: Identify the compact representations z and change factors 0 fmii based on C
6: for 0 ................................ N- do
7: for t ...
1 do
8: Observe 8, from Env;
9: if then
10:
0 Sl and 0" <- -
11:
else
0 0 and or 0--1
13:
end if
14:
for =, do
Infer mean pay (01) and variance (0i) of the change parameter 0/ via pyd Sample 0f N (1 (0)) (0i) end for
16:
18:
if H then
19:
0aa + 0; and 0% f 0::
20:
end if
21:
Generate a ~ Ta st 0m and receive reward rn, from Env; Add (st, rt 0% 0,) to replay buffer D;
22:
23:
Extract a trajectory with length & from replay buffer D;
24:
Learn FN-VAE (Appendix Alg. Al) with updateG=Falsse (i.e. with fixed masks C*));
25:
Sample batch of data from replay buffer D and update policy network parameters t):
26:
end for
27:
end for
After we estimate the initial FN-MDP with the FN-VAE, we can identify compact representations mii and gm following AdaRL [14]. In particular, the only dimensions of the state and change factors that are useful for policy learning are those that have directed path the reward in the graph G. The online policy TΦ (a, | smin go (0min Toa)) can be learned end-to-end, including learning the PN-VAE, as shown Alg. 1. We use SAC [19] as our policy learning model, so the policy parameters are v/r (T, Q).
Continuous changes. In Alg. 1 describe our framework in case of continuous changes that can span across episodes. We start by collecting a few trajectories and then Iearn our initial FN- VAE (Lines 3-4). We can use the graphical structure of the initial FN-VAL representations (Line 5). During the online model estimation and policy learning stage, we estimate the latent change factors 0 and 0r using the CF dynamics networks 7 and y" (Lines 8-20). Since in this case, we assume the dynamics of the change factors are smooth across episodes, at time t = 0 we will use the last timestep (H 1) of the previous episode as prior on the change factors (Line 10). Otherwise, we will estimate the change factors using their values in the previous timestep /- 1. We use the estimated latent factors 0, and observed state 8, to generate a, using Fy and receive reward (Line 21). We add (8t, T2, 0; 0,)) to the replay buffer (Line 22). We now update our estimation of the VAE, but we keep the graph G fixed (Lines 23-24). Finally, we sample a batch of trajectories in the replay buffer and update the policy network v (Line 25).
identify the compact
Discrete changes. Since we assume discrete changes happen at specific timestep ī = (f,,... FM)) we can easily modify Alg. for discrete changes, both within-episode and across-episode, by changing Lines 9-20 to only update the change parameters at the timesteps in é, shown in Appendix Alg. A2.

---

(8)
Return
Return
Return
-06
2.90
1005
1010
10
4. Stepss (M)
Steps (M)
*. Steps (M)
Otew
LILAC
1000
VariAA
Drazle (eveage (nghesd)
SAC leghort)
(d)
24
MCC
C33
Return
321
A42/00
9-3150
(g)
(h)
Return
0.8
Return
Return
Relurn
0.6
0.6
0.4
0.2
10 20-10
Latent features
Non-stationary degrees
Functional forms
Multiple changes
VantiADI
Figure 2: Summary of experimental results. (a)-(c). Average return (smoothed) across 10 runs. We only indicate the average of the highest result of all times for oracle and SAC. The shaded region is (J4 aa о), where is the mean and is the standard deviation. (a) Half-Cheetah-VV with continuous (sine) changes Jwi (b) Sawyer-Reaching with discrete across-episode changes s and (c) Minitaur with discrete scross-episode changes m and $t,v concurrently. (d) Ablation studies on Half-Cheetal. with across & within episode changes on dynamics and across episode changes on both dynamics and rewards. (c)-(f): Pairwise distance on learned between different time steps in
experiment with across-episode changes on rewards. (g)-()): Average (g) with within-in episode changes on wind
and normalized final return on 10 runs on forees using different number of dimensions in latent representation space; (h) with different non-stationary degrees on across-episode and multi-factor changes; (i) with different functional forms (across episode changes on dynamics); and (j) with different combinations across-episode changes.
Evaluation
We evaluate our approach on four well-established benchmarks, including Half-Cheetah-VV V3 from and Sawyer-Peg from Sawyer [22, 18], and Minitaur [23]. We modified 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 to the 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. The implementation will open-sourced at httph:tt//// /ttp:t//.ttte.e. 1y/ /3rr/o....
MuloCo [20, 21],
Half- 3aaf--heetah--33 in MuJoCo. In this task, the agent is moving forward using the joint legs and the objective is to achieve the target velocity 2º, We consider both the changes in the dynamies (change of the wind forces f change of gravity) and reward functions (change of target velocity 29), We also consider changes on the agent's mechanism, where one random joint is disabled. The reward function is respectively, at timestep The number of time steps in each episode is 50. For dynamics, we change the wind forces f in the environment. Moreover, in terms of the reward functions, we change the target velocity 6 to time-dependent variable. The change function in the dynamics f can be either continuous or discrete, and discrete changes can happen both within and across episodes.
- || e ll 0.05 0.05/|ael|2. where v² and Gr are the agent's 's velocity and action,
7

---

Methods
ZeUS
Meld
CaDM
Flyper-Dnnamie.
Ours 18.01
Best Avg. Return
12.45
6.38
4.04
3.91
Table 1: Average highest return on Sawyer-Peg. The number of random trails is 10.
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 in the reward function 6 is not generally stable in the continuous case, so we only consider the discrete and across episode change functions v for the reward. We also design a scenario where dynamic and reward functions change concurrently. We report all equations for 9 and g" in Appendix C.I.
Sawyer. We consider two robotic manipulation tasks, Sawyer- Sawyer-Reaching and Sawyer-Peg. We describe the non-stationary settings in Appendix C.2.
the sawyer arm is trained reach target position 8 The reward ra is the difference between the current position 8, and the target position r. = - || 80112- In this task, we cannot directly modify the dynamics in the simulator, consider reward-varying scenario, where the target location changes across each episode following a periodic function.
In
In Sawyer-Peg, the robot arm is trained to insert peg into designed target location 8 In this task, following [8], we consider a reward- reward-varying scenario, where the target location changes across each episode following periodic function. In order to compare with similar approaches, e.g., ZeUS [8], CADM [24], Hyperdynamics [25] and Meld [18], we evaluate our method on raw pixels. Following [8], we consider discrete across-episode changes, where the target location can change in each episode, and is randomly sampled from a small interval.
Minitaur. A minitaur robot is simulated quadruped robot with eight direct-drive actuators. The minitaur is trained to move at the target speed 8t. The reward is r, = 0.3 -00.3 355--00 8t,ul 0.01 ||a 20,-1 + a;-2ll We modify (1) the mass m (dynamics) and (2) target speed St.e (reward) of the minitaur. We consider continuous, discrete across-episode and discrete within-episode changes for the dynamics, and across- 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 [26], meta-RL approach for non-stationary RL, TRIO [5], as well as with two representative task embedding approaches, LILAC [7] and ZeUS [8]. The details on the meta-learning setups are given in Appendix D.4 D.4. We also compare with stationary RL method, SAC [19], which will be our lower- bound, 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 all baselines, we use SAC for policy learning. We compare with the baselines on average return and the compactness of the latent space in varying degrees and functional forms of non-stationarity.
Experimental results. Fig. 2(a)-(c) shows the smoothed curves of average return 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 continuous changes, we only compare with LILAC [7] since other approaches are not applicable. We smooth the learning curves by uniformly selecting a few data points for readability. Table shows the results on Sawyer- Peg experiments using raw pixels as input, based on the reported results in [8]. The learning curve is given in Appendix Fig. A11. For a fair comparison, we indicate the best and final average return for each baseline in Fig. A11. Full results of all experiments are given in Appendix C.4, including significance tests for Wilcoxon signed-rank test at a = 0.05 showing that FANS-RL is significantly better than baselines. We also visualize the learned graph in Sawyer in Appendix Sec C.7, showing that FN- recovers reasonable causal relations in this domain, where the true graph is unknown.
Ablation studies. We conduct ablation studies on each component of FANS-RL and report some key results in Fig. 2(d). 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 largest gain is provided by the fuctored representation, validating our original hypothesis, followed by state prediction. As expected, reward prediction is also important when there is a nonstationary reward, while smoothness is important for within-episode changes. The disentangied design of CF inference networks is valuable when there are changes on both dynamics
8

---

and reward functions. Full results are in Appendix C.S, showing that learning which state components are affected by change, i.e. estimating the binary masks C and C
or provides the largest gains.
Visualization on the learned 0. To verify that the learned 0 can capture the true change factors, we compute the pairwise distance between learned 0" at different time steps. We randomly sample 10 time steps from the Walf-Cheetah experiment with across-episode changes Fig. 2(e) gives the pairwise distance of 0" 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 is a positive correlation between the distance of learned 0" 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(1) shows that the distance of 0" among these 5 time steps is very small, indicating that the learned 0" almost the same for the similar values of change factors at different time steps. The visualization suggests that the learned O can capture meaningful mappings from the time-varying factors.
reward functions.
Varying latent dimensions, non-stationary levels and functional forms. Fig. 2(g) shows the normalized averaged return versus the number of two latent features on Half-Cheetah. Our framework can learn better policy with relatively smaller feature dimensions in the latent space than other approaches. As we show in Appendix C.4, Saywer and Minitaur have the similar trend, where we learn better policy than the baselines with fewer latent features. We also vary the non-stationary levels in Half-Cheetah with discrete across-episode changes on both dynamics and rewards. A higher non-stationary degree indicates faster change rate. Fig. 2(h) shows that FANS-RI. achieves the highest return 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 like and sinusoid waves. Fig. 2(i) displays the results, which indicate that FANS-RL can generally outperform the baselines on diverse non-stationary function forms. 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 benefits from the factored structure in FANS-RE RL. We conduct experiments with 1) only change wind forces (ID); 2) change wind forces and gravity concurrently (2D); 3) change wind force and target speed [I)+ and 4) change wind force, gravity, and target speed together (2D+ R) in an across-episode way in Half-Cheetah. From Fig. 2(j), we find that, thanks to the factored representation, FANS-RL performs better in those more complicated scenarios with multiple numbers and types of changes.
Related Work
Non-stationary and transfer RL. Early works in non-stationary and transfer RL [27, 28] only detect changes that have already happened instead of anticipating them. If the evolution of non-stationary environments is (Semi)-Markov chain, one can deal with non-stationarity with HM-MDPs [29] or HS3MDPs [30]. Several methods learn to anticipate changes in non-stationary deep RL. Chandak et al. [31] propose to maximize future rewards without explicitly modeling non-stationary environments. MBCD [32] uses change-point detection to decide if the agent should learn novel policy or reuse previously trained policies. Al-Shedivat 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 inferring the evolution of latent parameters, which captures the temporal change factors during the meta-testing phase. ReBAL and GrBAL [33] meta-train the dynamic prior, which adapts to the local contexts efficiently. However, these methods have to meta-train the model on a set of non- non-stationary tasks, which may not be accessible in real-world applications. Another line of research directly learns the latent representation to capture the non-stationary components. In particular, LILAC [7] and 7eUS [8] leverage latent variable models to directly model the change factors in environments, and Guo et al. [34] estimate latent vectors that describe the non-stationary or variable part of the dynamics. Perez et al. [35, 36] estimate low-dimensional latent vectors that capture the changes in dynamics and rewards and leverage these inferred change factors to facilitate transfer RL across tasks. Similarly, François-Lavet et al. [37] use compact state abstraction for recovering the sufficient low-dimensional representation of the environment, and only retrain or update some of its abstracted states for transferring to a novel but related task. Our approach fits in this line of work; however, as opposed to these methods, which model changes using a shared embedding space or a mixture model, and cannot distinguish which state components are affected, we use a factored representation.
Factored MDPs. Among works on factored representations in stationary settings, Hallak et al. [38] learn factored MDPs [11, 12, 13] to improve sample efficiency in model-based off-policy RL. Balaji

---

et al. [39] employ known factored MDP to improve both model-free and baseel-based RL algorithms. Working memory graphs [40] learn the factored observation space using Transformers. NeverNet [41] factorizes the state-action space through graph neural networks. Zholus et al. [42] factorize the visual states into actor-related, object of manipulation, and the latent influence factor between these two states to achieve sparse interaction in robotic control tasks. Differently from methods modeling the factored dynamics and rewards only, Zhou et al. [43] and Tang et al. [44] explore factored entities and actions, respectively. Zhou et al. [43] extend the factored MDP to the multi-entity environments, learning the compositional structure in tasks with multiple entities involved. Tang et al. [44] leverage the factored action space to improve the sample efficiency in healthcare applications. However, the factored structures in these two works are derived from inductive bias or domain experts.
Factored MDPs, causality and multiple en environments. Several works leverage factored MDPs to improve the sample efficiency and generalization of RI.. Most of these works focus on learning an invariant (e.g. causal) representation that fits all environments and do not support learning latent change factors. For example, AFaR [45] learns factored value functions via attention modules to improve sample efficiency and generalization across tasks. Mutti al. [46] learn causal structure that can generalize across a family of MDPs under different environments, assuming that there are no latent causal factors. Similarly, Wang et al. [47] propose to learn the factored and causal dynamics in model- baaeed RL, in which the learned causal structure is assumed to generalize to unseen states. By deriving the state abstraction based on the causal graph, it can improve both the sample efficiency and generalizability of policy learning of MBRL. While most of the previous works focus on learning a causal structure that is invariann,,,,, Pitis al. [48] learn locally causal dynamics that can vary each timestep and use it to generate counterfactual dynamics transitions in RL. While the previously described methods focus on domain or time-inveen invariant representations, in our work also focus on modelling domain or time-specific factors in the form of latent change factors. A related work, AdaRL [14] learns the factored representation and the model change factors under heterogeneous domains with varying dynamics or reward functions. However, AdaRL is designed only for the domain adaptation setting and constant change factors without considering non-stationarity.
Independent causal mechanisms. Another related line of work is based on independent causal mechanisms [49, 50J. Recurrent independent mechanisms (RIMs) [51] learn the independent transi- tion dynamics in RL with sparse communication among the latent states. Meta-RIMs [52] leverage meta-learning and soft attention to learn a set of RIMs with competition and communication. As opposed to these works, we do not assume that the mechanisms are independent and we learn the factored structure among all components in MDPs with DBN.
Conclusions, Limitations and Future Work
We describe Factored Adaptation for Non-Stationary RL (FANS-RL), a framework that learns a factored representation for non-stationarity that can be combined with any RL algorithm. We formalize our problem as a Factored Non Non-stationary MDP (FN-MDP), augmenting a factored MDP with latent change factors evolving as a Markov process. FN-MDPs do model family of MDPs, but instead include the dynamics of change factors analogously to the dynamics of the states. This allows us capture different non-stationarities, e.g., continuous and discrete changes, both within and 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 on benchmarks for continuous control and robotic manipulation, also with pixel inputs, and show it outperforms 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 of our approach. In future work, we plan to meta-learn the graphs among different tasks to improve the scalability of our approach and its applicability to complex RL problems, e.g., multi- RL.
Acknowledgments
FF would like to acknowledge the CityU High-Performann Computing (HPC) resources in Hong Kong SAR and LISA HPC from the SURFnl. SM was supported by the MIT-IBM Watson AI Lab and the Air Force Office of Scientific Research under award number FA8655-22-1-7155. BH would like to acknowledge the support of Apple Scholarship. KZ was partially supported by the National Institutes of Health (NIH) under Contract R01HL159805, by the NSF- Coverrnveeeeece Accelerator Track-D D award #21 134901, by a grant from Apple Inc., and by a grant from KDDI Research Inc.
10

---

References
[1] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[2] Gabriel Dulac-Amold, Nir Levine, Daniel J Mankowitz, Jerry Li, Cosmin Paduraru, Sven Gowal, and Todd Hester. Challenges of real- world reinforcement learning: definitions, benchmarks and analysis. Machine Learning, pages 1-50, 2021.
[3] Sindhu Padakandla. A survey of reinforcement learning algorithms for dynamically varying environments, ACM Computing Surveys (CSUR), 54(6):1-25, 2021.
[4] Maruan Al-Shedivat, Trapit Bansal, Yura Burda, Ilya Sutskever, Igor Mordatch, and Pieter Abbeel. Continuous adaptation via meta-leaming in nonstationary and competitive environments. In International Conference on Learning Representations (ICLR), 2018.
[5] Riccardo Poiani, Andrea Tirinzoni, and Marcello Restelli. Meta-reing reinforenfore learning by tracking task non-stationarity. In International Joint Conference on Artificial Intelligence (IJCAI), pages 2899-2905, 2021.
[6] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In International Conference on Machine Learning (ICML), pages 1126-1135, PMLR, 2017.
[7] Annie Xie, James Harrison, and Chelsea Finn. Deep reinforcement learning amidst lifelong non-stationarity. International Conference Machine Learning (ICMI., 2021.
[8] Shagun Sodhani, Franziska Meier, Joelle Pineau, and Amy Zhang. Block contextual mdps for continual learning. In Learming for Dynamics and Control Conference (L4DC), pages 608-623.
PMLR, 2022.
[9] Biwei Huang, Kun Zhang, Jiji Zhang, Joseph D Ramsey, Ruben Sanchez-Romero Clark Glymour, and Bernhard Schölkopf. Causal discovery from lederogeneouu/nonsattioonny data. Journal of Machine Learning Research, 21(89):612 634, 2020.
[10] Kun Zhang, Mingming Gong, Petar Stojanov, Biwei Huang, Qingsong Liu, and Clark Glymour. Domain adaptation a problem of inference on graphical models. Advances in Neural Information Processing Systems (NeurIPS), 33:4965-4976, 2020.
[11] Craig Boutilier, Richard Dearden, and Moisés Goldszmidt. Stochastic dynamic programming with factored representations. Artificial intelligence, 121 107, 2000.
[12] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In International Joint Conference on Artificial Intelligence (IJCAI), volume 16, pages 740-747,
1999.
[13] Ian Osband and Benjamin Van Roy. Near-optimal reinforcement learning in factored mdps. Advances in Neural Information Processing Systems (NeurIPS), 27:604 612, 2014.
[14] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. International Conference on Learning Representations (ICLR), 2022.
[15] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1780, 1997.
[16] Lukas Liebel and Marco Körner. Auxiliary tasks in multi-task learning. arXiv preprint
wwXiv:1805.06334 2018.
[17] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In
Advances Advances Neural Information Processing Systems (NeurIPS), pages 2455-2467, 2018.

---

[18] Zihao Zhao, Anusha Nagabandi, Kate Rakelly, Chelsea Finn, and Sergey Levine. MELD: meta-reinforcement learning from images via latent state models. In Conference on Robot Learning (CORL., pages 1246-1261. PMLR, 2020.
[19] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference Machine Learning (ICML), pages 1861-1870. PMLR, 2018.
[20] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[21] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 1606vv/606.05555 2016.
[22] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CORL), pages 1094-1100. PMLR, 2020.
[23] Jie Tan, Tingnan Zhang, Erwin Coumans, Atil Iscen, Yunfei Bai, Danijar Hafner, Steven Bohez, and Vincent Vanhoucke. Sim-to-real: Learning agile locomotion for quadruped robots. In Robotics:
Science and Systems, 2018.
[24] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference Machine Learning (ICML), pages 5757-5766. PMLR, 2020.
[25] Zhou Xian, Shamit Lal, Hsiao- Yu Tung, Emmanouil Antonios Platanios, and Katerina Fragki- adaki. Hyperdynamics: Meta-learning object and agent dynamics with hypernetworks. In International Conference on Learning Representations (ICLR), 2021.
[26] Luisa Zintgraf, Sebastian Schulze, Cong Lu, Leo Feng. Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: Variational meta-learning. Journal of Machine Learning Research, 22(289): 2021.
deep rl via
[27] Bruno C Da Silva, Eduardo W Basso, Ana LC Bazzan, and Paulo M Engel. Dealing with non-stationary environments using context detection. In International Conference on Machine Learning (ICML), pages 217-224, 2006.
[28] Richard S Sutton, Anna Koop, and David Silver. On the role of tracking in stationary environ- ments. In International Conference on Machine Learning (ICML), pages 871-878, 2007.
[29] Samuel Choi, Dit-Yan Yeung, and Nevin Zhang. An environment model for nonstationary reinforcement learning. Advances in neural information processing systems (NeurlPS), 12. 1999.
[30] Emmanuel Hadoux, Aurélie Beynier, and Paul Weng. Solving hidden-semi--ee mode markov decision problems. In International Conference on Scalable Uncertainty Management, pages 176-189. Springer, 2014.
[31] Yash Chandak, Georgios Theocharous, Shiv Shankar, Martha White, Sridhar Mahadevan, and Philip Thomas. Optimizing for the future in non-stationary mdps. In International Conference on Machine Learning (ICMI.), pages 1414-1425. PMLR, 2020.
[32] Lucas N. Alegre, Ana L. C. Bazzan, and Bruno C. da Silva. Minimum-delay adaptation in nta-statiy reinforcement learning via online high-confidence change-point detection. In Proceedings of the 20th International Conference on Autonomous Agents and Multiagent Systems (AAMAS), pages 905, 2021.
[33] Ignasi Clavera, Anusha Nagabandi, Simin Liu, Ronald S. Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta-reinforcemment learning. In International Conference on Learning Representations (ICLR),
2019.
12

---

[34] Jiaxian Guo, Mingming Gong, and Dacheng Tao. A relational intervention approach for unsupervised dynamics generalization in model-based reinforcement learnings. International Conference Learning Representations (ICIR), 2022.
[35] Christian F Perez, Felipe Petroski Such, and Theofanis Karaletsos. Efficient transfer learning and online adaptation with latent variable models for continuous control. In Continual Learning Workshop at NIPS, 2018.
[36] Christian Perez, Felipe Petroski Such, and Theofanis Karaletsos. Generalized hidden parameter mdps: Transferable model-based rl in a handful of trials. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAT), volume 34, pages 5403-5411, 2020.
[37] Vincent François-Lavet, Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforce- ment learning via abstract representations. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), volume 33, pages 3582-3589, 2019.
[38] Assaf Hallak, François Schnitzler, Timothy Mann, and Shie Mannor. Off-policy moodel-based learning under unknown factored dynamics. In International Conference on Machine Learning (ICML), pages 711-719. PMLR, 2015.
[39] Bharathan Balaji, Petros Christodoulou, Xiaoyu Lu, Byungsoo Jeon, and Jordan Bell-Masterson. Factoredrl: Leveraging factored graphs for deep reinforcement learning. NeurlPS Deep Rein- forcement Learning Workshop, 2020.
[40] Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan, and Matthew Hausknecht. Working memory graphs. In International Conference on Machine Learning (ICML), pages 6404-6414. PMLR, 2020.
[41] Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International conference on learning representations (ICLR), 2018.
[42] Artem Zholus, Yaroslav Ivchenkov, and Aleksandr Panov. Factorized world models for learning causal relationships. In ICLR Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
[43] Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compositional generalization in control. In ICML Workshop on Spurious Correlations, Invariance and Stability, 2022.
[44] Shengpu Tang, Maggie Makar, Michael Sjoding, Finale Velez, and Jenna Wiens. Leverag- ing factored action spaces for efficient offline reinforcement learning in healthcare. In Decision Awareness in Reinforcement Learning Workshop at ICML, 2022.
[45] Shagun Sodhani, Sergey Levine, and Amy Zhang, Improving generalization with approximate factored value functions. ICLR Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
[46] Mirco Mutti, Riccardo De Santi, Emanuele Rossi, Juan Felipe Calderon, Michael Bronstein, and Marcello Restelli. Provably efficient causal model-based reinforcement learning for systematic generalization. arXiv preprint 200.:2200..06555, 2022.
[47] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task-independent state abstraction. In International Conference on Machine Learning (ICML), pages 23151-23188 PMLR, 2022.
[48] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. In Decision Awareness in Reinforcement Learning Workshop at ICMI, 2022.
[49] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612-634 2021.
13

---

[50] Giambattista Parascandolo, Niki Kilbertus, Mateo Rojas-Carulle and Bernhard Schölkopf. Learning independent causal mechanisms. International Conference on Machine Learning (ICML), pages 4036 4044. PMLR, 2018.
[51] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations (ICLR), 2021.
[52] Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio. Fast and slow learning of recurrent independent mechanisms. In International Conference on Learning Representations (ICLR), 2021.
[53] Peter Spirtes, Clark Glymour, and Richard Scheines, Causation, Prediction, and Search. Apring-Verlage Lectures in Statistics, 1993.
[54] Kevin Patrick Murphy. Dynamic bayesian networks: representation, inference and learning. University of California, Berkeley, 2002.
[55] Judea Pearl, Causality: Models. Reasoning, and Inference. Cambridge University Press, Cambridge, 2000.
[56] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Causal inference time series using restricted structural equation models. In Advances in Neural Information Processing Systems (NeurlPS), pages 154-162, 2013.
14

---

Checklist
1. 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 extensive evaluation in Section
(b) Did you describe the limitations of your work? [Yes] We did provide a discussion on the limitations and potential future work in Section 7.
(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]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes] The assump- tions are stated in Prop. and Prop. 2.
(b) Did you include complete proofs of all theoretical results? [Yes] The complete proofs are in the Appendix B
3. 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 (c.g., with respect the random seed after running experi- ments multiple times)? [Yes] The learning curves (Fig. 2 (a-c), Appendix Fig. A2-A3) 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.I.
4. If you using existing assets (e.g., code, data, models) curating/releusing assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] Yes, we have cited the OpenAl Gym and Mujoco libraries our paper.
(b) Did you mention the license of the assets? [Yes] Yes, the licenses are given in Ap- pendix 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]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [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 participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "qKFJ7CFiX7v", "reviewer": "Reviewer_3UJG", "summary": "This paper proposes a novel formalism to model non-stationary environments, Factored Non-stationary Markov Decision Process (FN-MDP), that models latent factors that affect the dynamics and rewards and evolves with time as in Dynamic Bayesian Network. Then, they propose FN-VAE, which is built on top of AdaRL, and is used to learn the parameters (latent factors, dynamics, rewards, masks) of the underlying FN-MDP. The method is evaluated in several robotic environments and compared with state-of-the-art algorithms tailored to deal with non-stationarity in RL.", "questions": "Furthermore, I have the following questions and constructive criticisms:\n\n- In Definition 1, I suggest clarifying that it is assumed that the action is a vector of m dimensions. It is ambiguous as m could denote the number of discrete actions.\n\n- “Although the masks and noise are stationary, we allow the change of graph structure and noise distributions, whose changes are captured by $\\theta$ instead.” \nIf the masks are stationary, does it mean that if one of its values is 0 (denoting that one state variable does not influence some other state variable), then it will never have an effect, even if the dynamics change? In other words, to be able to capture all possible dynamics changes, would we need to have all masks with none zero elements?\n\n- The proof of Proposition 1 in the Appendix ends with: “In this setting, identifiability of the graph G is trivial [6].” However, It is not clear why this is trivial. As this is an important theoretical result of the paper, I suggest elaborating this proof.\n\n- The role of the state and reward encoders in Section 3 is not very clear. Why a single prediction network, conditioned on the masks, would be not enough to predict the next state? How are the encoders used?\n\n- How can we obtain the adjustable parameters $w_1$ … $w_7$? In the total loss, how can we obtain $k_1$ … $k_5$?\n\n- The total objective function in line 167 does not include $\\mathcal{L}_{\\text{rec-dyn}}$. Additionally, there is a typo here: the coefficient $k_5$ does not appear in the loss.\n\n- “the policy parameters = $\\psi = (\\pi, Q)$ are the actor $L_\\pi$ and critic loss $L_Q$.” I believe this sentence is not well formulated, how can a loss be a parameter?\n\n- In Algorithm 1, how is $s^{min}_t$ and $\\theta^{min}_t$ computed before following the policy in line 23?\n\n- “The horizon in each episode is 50.” The horizon in these environments is generally higher, why this value was chosen?\n\n- In Fig. 2(g-j) how were the values normalized? I suppose the value of 1.0 does not mean that the proposed method reaches the maximum possible reward, but that it is the maximum reference point used in the normalization.\n\n- Because of the causality assumptions (see Section A of the Appendix), the proposed model can not model different types of non-stationarity at the same time (e.g. wind and gravity). Are any of the related work (e.g. LILAC) able to deal with these changes? If so, perhaps an experiment in such a setting would be interesting to observe whether the proposed method could still outperform the other baselines.\n\n- For completeness of the first sentence of the Related Work section, the paper [“Minimum-Delay Adaptation in Non-Stationary Reinforcement Learning via Online High-Confidence Change-Point Detection.” In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems] is a more recent work on non-stationarity that detects changes that have already happened.\n", "limitations": "The authors appropriately discuss the limitations of the proposed approach. I would suggest, however, that the authors briefly include in the main text some of the assumptions of the method which are only discussed in Appendix A, as I believe they are very relevant and readers often do not check the Appendix.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proposed formalism is able to model, in a factored manner, non-stationary environments in which the factors that change the dynamics/reward also evolve over time. This is a novel idea not yet explored in the related literature.\n- The authors compare the proposed approach with several state-of-the-art algorithms tailored to deal with non-stationarity, showing relevant performance improvements. Several ablation experiments were also presented.\n\nWeaknesses:\n- Some mathematical and algorithmic definitions are not very clear in the current version and should be clarified in the main text.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ld7Pp-S8dXv", "reviewer": "Reviewer_T7en", "summary": "This paper introduces FANS-RL, a factored adaptation approach that aims to generalize to non-stationary scenarios including changes across episodes and within episodes. They formalize FN-MDPs, and prove the causal graph of the transition and reward function is identifiable. The experiments show FANS-RL outperforms the state of the art.", "questions": "I list a few ways in which the paper can be improved below.\n\n- The experimental chapter devotes a lot of space to the various environmental settings, and the experimental results seem to be insufficient. More detailed discussions on the experimental results in the main paper would be more appealing. \n\n-  The experiment section has compared different baselines with ablation experiments. However, a large number of results were presented in the Appendix section. From the main text, it's not very clear why the authors' approach is more effective than baselines, and ablation's analysis doesn't show what each component of FANS-RL does. It would be useful to extend the discussion in this section on why FANS-RL outperforms baseline methods.\n\n- Figure 2 (c) shows that FANS-RL's reward is comparable to Oracle, even higher than Oracle. It is useful to explain why this is happening.", "limitations": "N/A", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Originality:** Fair: The paper contributes some new ideas\n\n**Quality:**\t\tGood: The paper appears to be technically sound, but I have not carefully checked the details.\n\n**Clarity:**\t\t  Good: The paper is well organized but the analysis of the experiments' results could be improved.\n\n**Significance:** Fair:  The paper is likely to have a moderate impact within a subfield of AI.\n\n**Main Strengths:**\n\n- It is a nice synthesis of causality and RL, leading to an intuitive design.\n- It's fascinating to generalize a Factored-MDP to a FN-MDP, solving the non-stationary problem from a different perspective.\n- If the results are robust then it could be an important and useful tool.\n\n**Main Weakness:**\n\nThe paper is weak in a few ways. Mainly, the experiments section could include additional insight on the results as I will discuss in my question comments to the authors. As a minor point, the paper omits certain details such as a more thorough intro of AdaRL due to page limitations, which is understandable but would require more background knowledge from the readers. I suggest improving the readability in the revised version. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "KqoKI-tsuWC", "reviewer": "Reviewer_siKe", "summary": "In this paper, the authors considered non-stationarity which is faced when RL is deployed into the real world. To model the non-stationarity across episodes or within an episode, which seems to be general scenarios, the authors disentangled the non-stationarity as two types of latent change factors and formalized FN-MDPs. These disentangled latent factors enable the proposed method called FANS-RL to estimate changes both in the environment dynamics and the reward function. The authors assumed the generative process of FN-MDPs containing these latent factors, and FANS-RL learns the generative process of FN-MDPs during training. FANS-RL can also infer the change factors in environment dynamics and reward function at the current time-step via the learned generative process of FN-MDPs. The experimental results showed that FANS-RL outperforms baseline algorithms in various simulation tasks. Especially, the ablation study showed the performance effect of their method, which has many loss functions and variables they considered.\n\n", "questions": "FANS-RL can address non-stationarity, which is an important challenge in RL, and I think it is a great approach. However, it seems to rely on AdaRL, so I need additional explanation and experimental results to support the acceptance of the paper.\n\n1. Why did the authors formalize FN-MDPs, not N-MDPs? Although the ablation study without structure shows improvement with respect to the expected return, I need a brief and intuitive explanation for understanding the effect of factored MDP.\n\n2. As mentioned in “Weakness”,  the experimental results showed the adaptation to the limited changes in the environment dynamics. To find out the adaptation for the changes in the agent’s mechanism, I suggest additional experiments for these changes; for example, the increasing joint friction is introduced at a joint of Half-Cheetah.\n\n3. To check the effect of disentangling the latent change factor as the two ones, I suggest adding an ablation study without disentangling latent factors in figure 2-(d).\n\n4. Figure 2-(e,f) shows the distances of the learned $\\theta^r$, similarly, I suggest plotting the distances of the learned $\\theta^s$; for example, the $\\theta^s$ estimates the changing wind force. It is helpful to identify that the learned latent factor can capture the true change in environment dynamics.\n\n5. In Figure 2-(g), is the number of latent features of FANS-RL total number of dimensions of two latent features $\\theta^s$ and $\\theta^r$?\n\n6. What is the fundamental difference between the FN-VAE and the MiSS-VAE of AdaRL?\n\n", "limitations": "I already mentioned the limitations of this paper and the suggestions for the proposed method. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper proposed FANS-RL that addresses non-stationarity both in environment dynamics and reward function. \n\n1. Strengths\n\n(1-1) The authors formalized FN-MDPs that can deal with a general non-stationary RL setting including changes both across episodes and within an episode. It is a very important problem in RL.\n(1-2) FN-MDPs also contain two latent change factors for the environment dynamics and the reward function, and FANS-RL can more explicitly handle the non-stationarity of the MDP due to the two change factors. \n(1-3) FANS-RL shows better adaptation for non-stationarity than other existing algorithms in various tasks. Also, the ablation study showed the effect of each component of FANS-RL architecture with respect to the expected return.\n\n2. Weakness\n\n(2-1) The proposed method showed the adaptation for the limited changes in the environment dynamics; for example, the changes of wind force or agent’s mass or gravity. So, the reviewer wonders if FANS-RL can adapt to changes in the agent’s mechanism; for example, the disabled joint of the agent due to aging.\n\n(2-2) It seems difficult to find the fundamental difference between the architecture of FN-VAE and that of MiSS-VAE of AdaRL[1], even though they learn different MDPs. They have too similar components and loss functions.\n\n(2-3) The paper relies much on the definitions and theorems of AdaRL [1] in literature. It would be better to add an additional explanation of the shared representation $s^{min}$ and the compact domain-specific representation $\\theta^{min}$.  \nEspecially, in Line 74-84 in Appendix B.2, there exist sentences verbatim to those in  the proof of theorem1, page 18, [1].  I marked ethics alert  due to this point. Please check if this is OK.\nThe paragraph of concern: \"We denote the variable set in the system ....\"\n\n3. minor points (typo)\n(3-1) The first equation between line 72-73: $pa(\\theta_{ij,t}^s)\\to pa(\\theta_{j,t}^s)$\n(3-2) Figure 2-(f) in the main paper: the x and y labels need to correct.\n\n[1] Biwei Huang et al., AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning, ICLR 2022\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["Research Integrity Issues (e.g., plagiarism)"]}, "extra_text": {}}, {"review_id": "8MHz7rP4k5O", "reviewer": "Reviewer_rs6w", "summary": "The paper proposes a factored adaptation framework for reinforcement learning in non-stationary environments. The paper first formalizes the notion of Factored Non-stationary MDP (FN-MDP), which augments a factored MDP with time-evolving latent change factors under the Markovian assumption. The generative model in AdaRL is adapted to the time-varying setting by introducing additional equations for the latent change factors update. The causal structure is captured by a DBN, where the edges are represented using binary masks. These masks can be inferred under certain identifiability assumptions.\n\nThe paper then proposes FN-VAE to model the dynamics of the latent change factors, the state transitions, and the reward function. It includes some sparsity and smoothness regularization terms, and can be adapted for continuous or discrete changes. The FANS-RL framework combines model learning with policy optimization. The policy depends on the states and change factors for the dynamics, and the graph structure is used to select only the dimensions which affect the reward.\n\nExperiments are performed on four popular benchmarks, which have been modified to be non-stationary with both continuous and discrete changes. The results show that FANS-RL with FN-VAE performs better than some existing methods. The paper also includes ablation studies and experiments testing various aspects of the proposed method.", "questions": "1. The inference networks infer $\\theta_t$ from $\\tau_{0:t}$ as described in the paper but in Eq. (4) $q_\\phi$ also includes  dependence on $\\theta_{t-1}$. This dependence is not clear, does this represent the time-evolving hidden state of the LSTM, $h_{t-1}$?\n\n1. The one-step prediction encoders in Eq. (6) and (7) should have dependence on $\\theta^s_{t+1}$ and $\\theta^r_{t+1}$, respectively, as specified in the FN-MDP. Though Fig 2(d) shows that state and reward prediction help improve performance, predicting  $(s_{t+1},r_{t+1})$ from $(\\theta^s_{t}, \\theta^r_{t})$ seems to ignore the latent change factor dynamics. A clarification from the authors would be helpful.\n\n1. The multiple loss functions terms each have an associated weight, which seems like a hyper-parameter tuning nightmare. This can hurt the applicability of the proposed model in practice. Could the author comment on this aspect? On a related note, Section D.4.1 specifies the values used for these weights but is missing a description of the hyper-parameter selection method.\n\n1. The experiments verifying the values of learned reward CFs are insightful. It would be interesting to verify the dynamics CF in a similar experiment.\n\n**Minor comments:**\n\n- Eq. (2) seems to have a typing error, $r_t$ should be function of $s_t,a_t$ instead of $s_{t-1},a_{t-1}$.\n- Line 87 typing error, $c$ should not have subscript $i$.\n- Line 93-94 is confusing and seems to be making contradictory claims.\n- Algorithm 1, initial values of $\\theta_{old}$ not specified.\n- Algorithm 1, line 3, what is $k$? What does $\\tau^i_{0:k}$ denote?\n- The smoothing technique applied is slightly unusual. Based on my reading, most works employ moving average-based smoothing.\n- Figure 2(f) has incorrectly labelled axes.\n- Appendix lines 190-191 state that CF inference networks are fully-connected, shouldn’t they be LSTMs?\n- There is a lot of notation used in the paper, which naturally introduces scope of typing errors. I assume such issues will be fixed.", "limitations": "The paper provides a brief discussion on some limitations, namely the scalability and applicability of the approach to complex problems. The appendix provides a more detailed discussion on the effect of different assumptions or inductive biases, and acknowledges the lack of theoretical guarantees for the proposed method. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Learning effective and stable policies in non-stationary environments is an active area of research within the community. This is an important area of research and contributions in this direction can facilitate the application of deep RL to real-world scenarios. This work formalizes a non-stationary analogue of factored MDPs and proposes algorithms to model the environment and combine it with policy optimization to provide a general-purpose framework. The presented approach seems reasonable and theoretically sound. To the best of my knowledge, this represents novel and original work.\n\nThe writing is clear and concise. The factored generative model is quite complex with many moving parts, but the explanation and the diagrams make things fairly straightforward to follow. It is appreciable that both continuous and discrete changes are considered, and the framework is also adapted to image observations. The experiments are suitable and include comparison with recent works. The extensive results both in the main paper and the appendices are impressive. Notably, the authors perform significance tests for comparison with previous work.\n\nThe paper presents good quality and original work, with clear explanation of concepts and extensive results. The quality of submission can be enhanced by addressing some clarification questions and suggestions, provided below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept