Reinforcement Learning in Presence of Discrete Markovian Context Evolution 
Authors: Hang Ren
Aivar Sootla
Taher Jafferjee
Junxiao Shen
Jun Wang
Haitham Bou Ammar
Venue: ICLR 2022 Poster (2022)
Keywords: context-dependent Reinforcement Learning
model-based reinforcement learning
hierarchical Dirichlet process

Abstract: We consider a context-dependent Reinforcement Learning (RL) setting, which is characterized by: a) an unknown finite number of not directly observable contexts; b) abrupt (discontinuous) context changes occurring during an episode; and c) Markovian context evolution. We argue that this challenging case is often met in applications and we tackle it using a Bayesian model-based approach and variational inference. We adapt a sticky Hierarchical Dirichlet Process (HDP) prior for model learning, which is arguably best-suited for infinite Markov chain modeling. We then derive a context distillation procedure, which identifies and removes spurious contexts in an unsupervised fashion. We argue that the combination of these two components allows inferring the number of contexts from data thus dealing with the context cardinality assumption. We then find the representation of the optimal policy enabling efficient policy learning using off-the-shelf RL algorithms. Finally, we demonstrate empirically (using gym environments cart-pole swing-up, drone, intersection) that our approach succeeds where state-of-the-art methods of other frameworks fail and elaborate on the reasons for such failures.

Published as a conference paper at ICLR 2022
REINFORCEMENT LEARNING IN PRESENCE OF DIS- CRETE MARKOVIAN CONTEXT EVOLUTION
Hang Ren' Huawei UK R&D
Aivar Sootla' Huawei UK R&D
Taher Jafferjee Huawei UK R&D
Junxiao Shen
Huawei UK R&D University of Cambridge
Jun Wang
Haitham Bou-Ammar
University College London jun. wangecs. ucl. ac. uk
Huawei UK R&D and Honorary Lecturer at UCL haitham. annar@huawei. com.
ABSTRACT
We consider a context-dependent Reinforcement Learning (RL) setting, which is characterized by: a) unknown finite number of not directly observable contexts; b) abrupt (discontinuous) context changes occurring during an episode; and c) Markovian context evolution. We argue that this challenging case is often met applications and we tackle it using a Bayesian model-base bodellbased approach and varia- tional inference. We adapt sticky Hierarchical Dirichlet Process (HDP) prior for model learning, which is arguably best-suited for infinite Markov chain modeling. We then derive a context distillation procedure, which identifies and removes spuri- ous contexts in an unsupervised fashion. We argue that the combination of these two components allows inferring the number of contexts from data thus dealing with the context cardinality assumption. We then find the representation of the optimal policy enabling efficient policy learning using off-the-shelf RL algorithms. Finally, demonstrate empirically (using gym environments cart-pole swing-up, drone, intersection) that our approach succeeds where state-of-the-art methods other frameworks fail and elaborate on the reasons for such failures.
INTRODUCTION
Our world becomes more automated every day with the development of self-driving cars, robotics, and unmanned factories. Many of these automation processes rely on solutions sequential decision- making problems. Reinforcement Learning (RI.) has recently been shown to be an effective tool for solving such problems achieving notable successes, c.g., solving Atari games (Mnih al.) 2013), defcating the (arguably all-time) best human players in the game of GO (Silver et al. 2016). accelerating robot skill acquisition (Kober al. 2013). Most of the successful RL algorithms rely on abstracting the sequential nature of the deccssionmmkinng as Markov Decision Processes (MDPs). which typically assume both stationary transition dynamics and reward functions.
As classic RL departs well-behaved laboratory setting. stationarity assumptions can quickly be- come prohibitive, sometimes leading to catastrophic consequences. an illustration, imagine an autonomous agent driving vehicle with changing weather conditions impacting visibility and tyre grip. The agent must identify and quickly adapt to these weather conditions changes in order to avoid losing control of the vehicle. Similarly, an unmanned acrial vehicle hovering around fixed set of coordinates needs to deal with sudden atmospheric condition changes (e.g.. wind, humidity etc). Another similar and realistic example is actuator failure, which changes how the action affects the MDP. Following Menke & Maybeck (1995) we distinguish "soft" (a percentage drop in action efficiency) and "hard" (action is not affecting the MDP) failures. The failures can also be dynamic as a "soft" failure in one actuator can overload other actuators introducing a chain of failures. An environment with a fixed weather condition or an actuator failure can be modeled as an MDP, however, with the changing weather or arising failures the environment becomes non-s -on-stationary
"quull contribution

---

Published as a conference paper at ICLR 2022
We can model this type of environments by making MDP state transitions dependant on the context variable, which encapsulates the non-stationary and/or other dependencies. This kind of contextual Markov Decision Processes (C-MDPs) incorporate number of different RL settings and RL frame- works (Khetarpal et al. 2020): non-stationary RI,, where the context changes over time and the agent needs to adapt to the context (e.g., the weather conditions are slowly changing over time); continual and/or meta RL, where the context is sampled from distribution before the start of the episode (c.g. the weather changes abruptly between the instances the vehicle has been deployed).
Although a significant progress in solving specific instances of C-MDPs has been made, the setting with countable number of contexts with Markovian transitions between the contexts has not received sufficient attention in the literature - a gap we are aiming to fill. The closest related works consider only special cases of our setting assuming: no context transitions (Xu et al., 2020); Markovian context transitions with a priori known context transition times (Xie et al. 2020p; finite state-action spaces (Choi et al. 2000). To enable sample efficient context adaptation Xu al. (2020) and Xie) et al. (2020) developed model-based reinforcement learning algorithms. Specifically, Xie et al. (2020) learned latent space variational auto-encoder model with Markovian evolution continuous context-space, while Xu et al. (020)) adopted Gaussian Process model for MDPs and Dirichlet process (DP) prior to detect context changes without modeling the context evolution. Note the DP prior, which is conjugate prior for categorical distribution, is not sufficient to model context transitions. We also propose model-based RI. algorithm. however, we model the context and state transitions using the Hicrarchical Dirichlet Process (HDP) (Teh et al., 2006; Fox et al., 2008a) prior and neural network with outputs parametrizing Guassian distribution (i.c., its mcan and variance), respectively, and refer to the model as HDP-C-MDP, We chose the HDP prior since it only requires the knowledge of an upper bound on context cardinality for tractable inference, and it is better suited for Markov chain modeling than other priors (Teh et al., 2006). Inspired by [Blei et al. (2006) we derive a model learning algorithm using variational inference, which is amenable to RL applications using off-the-shelf algorithms.
Our algorithm relies on two theoretical results: a) we propose context distillation procedure (i.e., removing spurious contexts); b) we show that the optimal policy depends on the context belief (context posterior probability given past observations). We derive another theoretical result, which shows performance improvement bounds for the fully observable context case. Equipped with these results, we experimentally demonstrate that we can infer the true context cardinality from data. Further, the context distillation procedure can be used during training as a regularizer. Interestingly, can also be used to merge similar contexts, where the measure of similarity only implicitly defined through the learning loss. Thus context merging is completely unsupervised. We then show that our model learning algorithm appears to provide an optimization profile with fewer local optima than the maximum likelihood approach, which we attribute to the Bayesian nature of our algorithm. Finally, we illustrate RL applications on an autonomous car left turn and an autonomous drone take-off tasks. We also demonstrate that state-of-the-are algorithms of different frameworks (such as continual RL and Partially-Obsery Observable Markov Decision Processes (POMDPs)) fail to solve C-MDPs in our setting, and we elaborate on potential reasons why this is the case.
2 PROBLEM FORMULATION AND RELATED WORK
We define a contextual Markov Decision Process (C-MDP) as tuple M. (C, R, ra). where S is the continuous state space; A is the action space; 7a [0, 1] is the discount factor; and denotes the context set with cardinality (C). In our setting, the state transition and reward function depend on the context, i.c., Ps context distribution probability Pe: To x C + [0, 1] is conditioned on T, the past states, actions and contexts (an, ag. Co..... C4-1, 3. Our definition a generalization of the C-MDP definition
S x A x S + 10. R
S x A R. Finally, the
by Hallak et al. (2015). where the contexts are stationary, i.c., Pe C + [0, 1). We adapt our definition in order to encompass all the settings presented by Khetarpal et al. 2020), where such
C-MDPs were used but not formally defined.
Throughout the paper, we will restrict the class of C-MDPs by making the following assumptions: (a) Contexts are unknown and not directly observed (b) Context set cardinality is finite and we know its upper bound K; (c) Contexts switches can occur during an episode and they are Markovian. In particular, we consider the contexts eg representing the parameters of the state
2

---

Published as a conference paper at ICLR 2022
transition function 22, and the context set C to be subset of the parameter space O To deal with uncertainty, we consider set C such that: a) |C| K |C|; b) all its elements 02 C are sampled from a distribution H (A), where À is a pointing toward particular parameter vector Oz., which leads to:
Let 2 E [L,..., K) be the index variable
21 Po Cat(
2 2 ( |0| Cat(p
(1)
St 8t-1, 1,-1, 2,, (0m) ~ pl ( St a,- 2n,,,
Ok H(),
where Pa is the initial context distribution, while R P P represents the context transition operator.
As the reader may notice our model is tailored the case, where the model parameters change abruptly due to external factors such weather conditions, cascading actuator failures etc. The change is formalized by a Markov variable 2,, which changes the MDP parameters 0z,- Our approach can also be related to switching systems modeling (cf. [oox et al. (2008a); [Becker-Ehmck (eekk/Beeeer et al. (2019); Dong et al. (2020)) and in this case the context representing the system's 's mode. While we can draw parallels with these works, we improve the model by using nonlinear dynamics (in comparison to (Becker-Ehmck et al. (2019): Fox et al. (2008a)), by using the HDP prior (Becker-Ehme) Ehmnkk et al.) (2019); Dong et al.J (2020) use maximum likelihood estimators), and finally, by proposing the distillation procedure and using deep learning (in comparison to to/Fox et al. (2008a)). Also note that typically a switching system aims represent complex nonlinear (Markov) model using collection of simpler (e.g.. linear) models, which is different from our case. Other restrictions on the space of C-MDPs lead to different problems and solutions (Khetarpal et al. 2020). We briefly mention few notable cases, while relegating detailed discussion to Appendix A Assuming that the context is deterministic with 24 t puts us in the non-stationary RL setting (cf.
Chandak al. (2020)), where
opposed to our case of possibly abrupt changes. Restricting the context to a stationary distribution sampled in specific time points (c.g., at the start of the episode) can be tackled from the continual RL (cf. Nagabandi et al.) (2018)) and meta-RL RL perspectives (cf. Finn et al. (2017)) but both are not designed to handle the Markovian context case. Finally, our C-MDP can be seen as a POMDP. Recall that POMDP is defined by a tuple Mpo = {X. Px, Po, R, Td. p(zo)) where X, A, O are the state, action, observation spaces, respectively; Px X' x A x X + [0, 1] is the state transition probability; Po : X' x A x O - [0, 1] is the conditional observation probability; and p(æo) is the initial state distribution. In our case, x 28) and 8.
it is common to assume a slowly smoothly changing non-stationarity
REINFORCEMENT LEARNING FOR MARKOV PROCESSES WITH MARKOVIAN CONTEXT EVOLUTION
There are three main components in our algorithm: the HDP-C-MDD derivation, the model learning algorithm using probabilistic inference and the control algorithms. We firstly briefly comment on each on these components to give an overview of the results and then explain our main contributions to each. The detailed description of all parts of our approach can be found in Appendix.
In order to learn the model of the context transitions, we choose the Bayesian approach and we employ Hierarchical Dirichlet Processes (HDP) as priors for context transitions, inspired by time- series modeling and analysis tools reported by Fox et al. (2008a b) (see also Appendix C.Ip. We improve the model by proposing a context spuriosity measure allowing for reconstruction of ground truth contexts. We then derive model learning algorithm using probabilistic inference. Having model, we can take fef-the-sheil algorithms such as Model Predictive Control (MPC) approach using Cross-Entren gradient approach Soft-actor critic (SAC) (ef. [Haarnjja et al. (2018) and Appendix (.6),, which are
Minimization (CEM) (cf. Chua et al.) (2018) and Appendix C.5), ), or a policy-
both well-suiter for model-based reinforcement learning. While MPC can be directly applied to our model, for policy-based control we first derive the representation of the optimal policy.
Generative model: HDP-C-MDP. Before presenting our probabilistic model, let us develop some necessary tools. A Dirichlet process (DP). denoted as DP H), is characterized by a concentration parameter y and base distribution (A) defined over the parameter space e A sample G from DP H) is a probability distribution satisfying (G(A)),... G(A,)) ~ Dir(yH (A)), (A,)) for every finite measurable partition A1..., A, of θ, where Dir denotes the Dirichlet distribution.
3

---

Published as a conference paper at ICLR 2022
Sampling G is often performed using the stick-breaking process (Sethuraman, 1994) and constructed by randomly mixing atoms independently and identically distributed samples Os from H:
k-1 1k Beta(1,7). Bx uk II (1-v))
G= 3400
where 50, is the Dirac distribution at 0g. We note that the stick- breaking procedure assigns progres- sively smaller values to 3% for large k, thus encouraging smaller number of meaningful atoms. The Hierarchical Dirichlet Process (HIP) is group of DPs sharing a base distribution. which itself is sample from a DP: G ~ (7 H), Gj ~ (a G) for all j = 0, 1, 2,... TTh et al. 2006). The distribution G guarantees that all G, inherit the same set of atoms, i.c., atoms of G. while keeping the benefits of DPs in the distributions Gj. can be shown that Gj = Σiz Ps:.0 for some Pÿk whose sampling can be performed using another stick-breaking process (Teh al. 2006). We consider its modified version introduced by Fox et al. (2011):
)) Pjk Pik 15x);
Pyk a,K,, Kr Beta aB + Kōgk, a
Σ a3 + Kdjz
(3)
where k 2 1, 2 0, 55z is the Kronecker delta, the parameter K. 2 0, called the sticky factor, modifies the transition matrix priors encouraging self-transitions. The sticky fac- tor serves as another measure of regularization reducing the average number of transitions. In our case, the atoms {0x} forming the context set C are sampled from H (A), while Pyk are the parameters of the Hidden Markov Model: Po is the initial context distribution and Pj are the rows the transition matrix R. Our probabilis- tic model is constructed in Equations 12 3 and illustrated in Figure T as a graphical model. We stress that the HDP in its stick-breaking con-
N
struction assumes that [C] is infinite and count- able. In practice, however, we incorporate the upper cardinality bound K by using a truncated variational distribution, as explained later.
Figure 1: HDP-C-MDP
Context Distillation. HDP-C-MDP promotes small number of meaningful contexts and some contexts will almost surely be spurious, i.c., we will transition to these contexts with a very small probability. While this probability is small we may still need explicitly remove these spurious contexts. Here we propose measure of context spuriosity and derive a distillation procedure removing these spurious contexts. As a spuriosity measure we will use the stationary distribution of the chain pº, which is computed by solving nd pdd R. The distillation is then performed as follows: in stationarity the probability mass of a context is smaller than a threshold Edistil then transitioning this context is unlikely and it can be removed. We develop the corresponding distilled Markov chain in the following result, which we prove in Appendix B.1 while the distillation algorithm can be found in Appendix C.4
Theorem Consider Markov chain p' 101111pp R with stationary distribution pd and distilled
Ii = S 2 Edistil} and spurious In = 1i < Edistit) state indexes, respectively. Then a) the matrix R RI, + RI, I, RI,,I,) RI,,,,)-R1,,1 is valid probability transition matrix; b) the Markov chain p = -p' - such that its stationary distribution & x pf,-
Model Learning using Probabilistic Inference. We aim to find variational distribution g(v, nt. 0) to approximate the true posterior p,,, (D)),, for dataset DD ((s a ) where {s and a² {ai) are the state and action sequences in the i-th trajectory. We minimize KLL (g(v, μ. 0) p(v, μ. D())), or equivalently, maximize the evidence lower bound (ELBO):
N
ELBO- Ey(n,)
Σ
log ((s aa aa n. 0) KCL (c)) ()) 0))
(4)
4

---

Published as a conference paper at ICLR 2022
The variational distribution above involves
random variables v, p. 8. To reach tractable solution, we assume [C] K and exploit the standard mean-field assumption (Blei et al. 2017) and the K-trruncated variational distribution similarly to Blei et al. (200)); Hughes et al. (2015)
Bryant & Sudderth (2012) as follows:
K II
K-1
q(,,,, ,4, 0) ((w))(1))(()) 2(010)
8(0m Ox), a(v1D) - II
d(ve(Pe), dluk - 1) 1,
kml
K-1
((pp)) = II II j=0 1=1
Beta Pyk Psk.
> A5s g05x= 10 1)
(5)
where hatted symbols represent free parameters. For 0 and v, we seek MAP point estimate instead of full posterior (see Appendix C.3 for a discussion on our design choices). Random variables not shown in the truncated variational distribution are conditionally independent of data, and thus can be discarded from the problem. We maximize ELBO using stochastic gradient ascent, while the gradient computations are performed using the following two techniques: (a) we compute the exact context posterior using forward-backward backward message passing algorithm, (b) use implicit reparametrized gradients to differentiate with respect to parameters of variational distributions (Figurnov et al.) 2018). We present the detailed derivations in Appendix C.2 We also can perform context distillation during training as discussed in Appendix C.4 While adding some computational complexity, this procedure acts as regularization for model learning as we show in our experiments.
Representation of the optimal policy. First, we notice that the model in Equation iis POMDP. which we get by setting T; := :2 8t) and o, == st. the POMDP case, we cannot claim that 0 depends only on o, and a, Therefore the Bellman dynamic programming principle does not hold for these variables and solving the problem is more involved. practice, one constructs the belief state b, p(x, () (Astrom 1965). where 1 = {by. Ogr,, act / is called the information state and is used to compute the optimal policy. Since the belief state a distribution, is generally costly to estimate in continuous observation or state spaces. In our case, estimating the belief is tractable, since the belief of the state st is the state itself (as the state 8, is observable) and the belief of 2,, which we denote bj. is vector of a fixed length at every time step (as 2 is discrete and bj is its filtering distribution). We have the following result with the proof in Appendix B.2
bi+1. where (bi+1); X N; Ej P( 18,, Os, at)pp((b bi);, where (bī); ) are the entries of bi; b) the optimal policy can be com- puted as π (ss b²) = argmax, (88, b² ,a), where the value function satisfies the dynamic program- ming principle Q (:,, bi, at) = r(8,, (e, at) + 2 Σ, N₂ maxa; ((8+11, bi,1, a(+1) d8
Theorem a) The belief of can be computed as p(z++1 (
Computational framework. Algorithm T summarizes our approach and is based on the standard model-based RL frameworks (c.g., |Pineaa et al.] (202T). Effectively, we alternate between model updates and policy updates. For the policy updates we relahel (recompute) the beliefs for the historical transition data. As MPC methods compute the sequence of actions based solely on the model such relabeling is not required.
Performance gain for observable contexts. It is not surprising that observing the ground truth of the contexts should improve the maximum expected return. In particular, even knowing the ground truth context model we can correctly estimate the context 2++1 only a posteriori, i.e., after observing the next state S1+1- Therefore at every context switch we can mislabel it with high probability. This leads to performance loss, which the following result quantifies using the value functions. We have the following result with the proof in Appendix B.3
Assume know the true transition model of the contexts and states and consider two settings: we observe the ground truth zt and we estimate it using bj. Assume we computed the optimal model-based policy It <-(|st, bā) with the return R and the optimal ground-truth policy Tgg (-) (t, 2441) with the corresponding optimal value functions Vga (s, and Qge (s, 2, a), then:
Theorem 3
M 2 (Q(st... Ztm +1, aim) (8 2tm +1, atmm))s
Ez Vet 21) - -R Braatt wug
Σn
OUT m
where M is the number of misidentified context switches in a trajectory T.
5

---

Published as a conference paper at ICLR 2022
Algorithm 1: Learning to Control HDP-C-MDP
Input:
Edistill distillation threshold. warm number of trajectories for warm start, Niraj number of newly collected trajectories per epoch, Nepochs number of training epochs, AGENT policy gradient MPC agent
Initialize AGENT with RANDOM AGENT, D 0;
for i == I,..., "..,,NNpppce NNepochs do
Sample Niraj (Nwarm if i = 1) trajectories from the environment with AGENT;
Set Deew {(st, )²)} [Neraa where s {- and a' {a:) are the state and d') action sequences in the i-th trajectory. Set D = D U D new
Update generative model parameters by gradient ascent on ELBO in Equation
Perform context distillation with Edistill:
AGENT POLICY then
Sample trajectories for policy update from D; Recompute the belief's using the model for these trajectories;
Update policy parameters end end
return AGENT
4 EXPERIMENTS
In this section, we demonstrate that the HDP offers an effective prior for model learning, while the distillation procedure refines the model and can regulate the context set complexity. We also explain why sfate-of-the-art methods from continual RL, meta-RL and POMDP literature can fail in our setting. We finally show that our algorithm can be adapted to high dimensional environments. We delegate several experiments to Appendix due to space limitations. We show that we can learn additional unseen contexts without relearning the whole model from scratch. We also illustrate how the context distillation during training can be used to merge contexts in an unsupervised manner thus reducing model complexity. We finally show that our model can generalize non-Markovian and state dependent context transitions.
We choose the switching process to be chain, however, enforce cool-off period, i.c. the chain cannot transition to new state until the cool-off period has ended. This makes the context switching itself non-stationary MDP. This is done avoid switches every time step, but also to show that our method is not limited to the stationary Markov context evolution.
Control Baselines: (1) SAC algorithm with access to the ground truth context information (one-hot- encoded variable zạ) denoted as FI-SAC; (2) SAC algorithm with no context information denoted as NI-SAC (3) continual RL algorithm for contextual MDPs (Xu et al. 2020 where a Gaussian process is used to learn the dynamics while identifying and labeling the data with contexts, which denoted as GPMM; (4) POMDP approach, where the context set cardinality known and the belief is estimated using an RNN, while PPO (Schulman et al. 2017 Kostrikov) 2018) is used to update the policy. We denote this approach as RNN-PPO.
Modeling Prior Baselines: (1) a model with sticky Dirichlet priors pj ~ Dir(a; = a/K + kŌ (2) model which removes all priors and conducts maximum-likelihood (MLE) learning. All the other relevant experimental details (including hyper- sypee-parameters) are provided in Appendix D
A00
A80
(a) HDP
(b) Dirichlet
(c) MLE
(d) HDP distillation
Figure 2: Part-Pole Swing- Up. Transition matrices, initial p(=g) and stationary p(z0) distributions of the learned context models for Result A. Zo Z4 stand for the learned contexts.

---

Published as a conference paper at ICLR 2022
23
23
20-
20
20-
03
03
89
20
(a) HDP
(b) Dirichlet
(c) MLE
Figure 3: Cart-Pole Swing-Up. Time courses the learned context models for Result A, CO and C11 stand for the ground true contexts, while 70- Z4 the learned contexts.
Initial testing on Cart-Pole Swing-up Task (Lovatto, 2019). We attempt to swing up and balance a pole attached cart. This environment has four states and one action. We introduce the contexts by multiplying the action with a constant thus modulating the actuation effect. We will allow negative X modeling catastrophic (or hard) failures, and positive X modeling soft actuation failures. Result A: HDP is an effective prior for learning an accurate and interpretable model. In Figure 2, we plot the expectation of po and R extracted from the variational distribution q(u) for HDP, Dirichlet and MLE priors for the Cart-Pole Swing-up Environment with the context set
C C= (1) and C K 5. The MIE learning appears be trapped in local optimum as the results in Figure 2(c) suggest. A similar phenomenon has been reported by Dong et al. (2020), where an MLE method was used and heuristic entropy regularization and temperature annealing method is adopted to alleviate the issue. All in all, while MLE learning can appear to be competitive with a different random seed, this approach does not give consistent results. The use of Dirichlet priors appears to provide a better model. Furthermore, with an appropriate distillation threshold the distilled transition matrices with HDP and Dirichlet priors are very similar to each other. However, the threshold for Dirichlet prior distillation needs to be much higher calculations of the stationary distributions suggest. This implies that spurious transitions are still quite likely. In contrast, the HDP prior helps to successfully identify two main contexts (20 and Z2) and accurately predict the context evolution (see Figure 3) Furthermore, the model is more interpretable and the meaningful contexts can often be identified with naked eye.
Result B: Distillation acts as a regularizer. We noticed that the context Z2 has a low probability mass in stationarity, but high probability of self-transition (Figure 2(a).. This suggest that spurious transitions can happen, while highly unlikely. We speculate that the learning algorithm tries to fit the uncertainty in the model (e.g.. due to unseen data) to one context. This can lead to over-fitting and unwanted side-effects. Results in Figure 2(d) suggest that distillation during training can act as a regularizer when we used a high enough threshold Edistil = 0.1. We proceed by varying the context set cardinality |C] (taking values 4, 5, 6, 8, 10 and 20) and the distillation threshold daistil (taking values 0, 0.01, and 0.1). Note that distill during training and we refer to the transition matrix for the distilled Markov chain as the distilled transition matrix. As the ground truth context cardinality is equal to two, the probability of the third most likely context would signify the learning error. In T we present the stationary probability of the context with the third largest probability mass. In particular, for Č = 20 the probability mass values for this context are larger than 0.01. This indicates a small but not insignificant possibility of transition to this context, if the distillation does not remove this context. We present some additional details on this experiment in Appendix F. Overall, we can conclude that is safe to overestimate the context cardinality.
Result C: MDP, POMDP and continual RL methods can be ineffective. In Figure 4(a) we plot the learning curves for our algorithms and compare them to each other for x = -1. CEM, which is known to perform well in oow-dimensional environments, learns faster than SAC. Note that there is no significant performance loss of C-SAC in comparison with the full information case exhibiting the power of our modeling approach. We evaluated FI-SAC, C-SAC, C-CEM on three seeds.
We now compare the control algorithms for various values of X. We present the results of our experiments in Table 2 and we also discuss the comparison protocols in Appendix E.6 Here
7

---

Published as a conference paper at ICLR 2022
Table 1: Comparing the probability mass of the third most probable state in the stationary distribution. We vary the cardinality of the estimated context set C and the distillation threshold Edistil- Red indicates underestimation of distillation threshold.
CC|
4
5
8
10
20
Edistil +
8.58e-03
7.06e-03
3.71e-03
6.85e-03
2.20e-03
2.25e-02
0.01
1.06e-03
1.24e-03
1.37e-03
2. 2.19--03
2.56e-03
1.60c-02
0.1
1.21e-03
1.54e-03
1.70e-03
2.80e-03
3.54e-03
9.86e-03
Carteal soludden Vilues
Cartect ecmorae
16
(a) Learning curves
(b) GPMM X
(c) GPMM X = 0.5
(d) RNN "Beliefs"
Figure 4: Cart-Pole Swing-Up Up. Learning curves for x= -1 (a), time courses the learned context
models using GPMM (b)-(c) and the learned model belief by RNN-PPO (d).
focus on the reasons why both RNN-PPO and GPMM can fail some experiments and seem to perform well in others. In GPMM, it is explicitly assumed that the context does not change during the episode, however, the algorithm can adapt to a new context. While a posteriori context estimation has limited success for X 0.5 (see Figure 4). the context adaptation is rather slow for our setting resulting in many context estimation errors, which reduces the performance, Furthermore, it appears that estimating hard failures is a challenge for GPMM. RNN-PPO appears perform very well for > 0 (see Table 2 and fail for X = 11 however, when we plot the output of the RNN, which is meant to predict the beliefs, we see that the average context prediction is quite similar across different experiments (see Figure 4) It is worth noting that the mean of the true belief variable is 0.5 for all X. both contexts are equally probable at every time step. Therefore, the RNN approach does not actually Icarn a belief model, but an "average" adjustment signal for the policy, and hence it will often fail to solve C-MDP. Interestingly, with x = 0.5 our modeling algorithm learns only one meaningful context with high distillation threshold while still solving the task. This is because for both x 0.5 and x = 1 the sign of optimal actions for swing up are the same and both have sufficient power to solve the task. We compare to further baselines in Appendix E.6
Our model is effective for control in twelve dimensional environments (Drone and Intersec- tion). In the drone environment (Panerati et al., (02T). the agent aims at balancing roll and pitch angles of the drone, while accelerating vertically, i.c., the task is to maximize the upward velocity. This environment has twelve states (positions, velocities, Euler angles, angular velocities in three dimensions) and four actions (motor speeds in rotation per minute). In the highway intersection environment (Leurent) 2018), the agent aims at performing the unprotected left turn maneuver with an incoming vehicle turning in the same direction. The goal of the agent is to make the left turn and follow the social vehicle without colliding with it. The agent measures positions, velocities and headings in æ, y axes of the ego and social vehicles (twelve states in total), while controlling the steering angle and acceleration/ deceleration. In both environments we introduce the contexts by multiplying the maximun actuation effect by constant X. specifically, motor speeds in the drone environment and steering angle the highway intersection environment. The results in Table demonstrate that both MPC and policy learning approaches with the model are able to solve the task, while using no information (NT) about the contexts dramatically reduces the performance. Note that in the drone environment C-CEM algorithm exhibits slightly better performance than C-SAC, while in the intersection environment C-SAC controls the car much better. We can only hypothesize that the policy's 'ss feedback architecture (mapping states actions) is better suited for complex tasks such as low-level vehicle control, where MPC approaches require a substantial tuning and computational effort compete with policy based approach.
8

---

Published as a conference paper at ICLR 2022
failure +
hard
soft = 0.1
soft a 0.3
soft = 0.5
algo |
FI-SAC
84.50 = 1.79
76.63 + 8.54
84.75
3.07
86.92 + 1.03
C-SAC
85,38 = 1.64
76.80 8.91
86.76 + 2.88
88.35 = 1.30
C-CEM
87.63 0.14
60.15
83.15 7.72
89.08 + 1.90
GPMM
3.50 + 18.59
3.55 7.83
10.64 16.10
49.61 = 19.13
RNN-PPO
-0.17 18.06
64.10 + 21.37
74.58 + 20.66
67.01 + 8.52
Table 2: Mean + standard deviation of expected return for: our algorithms (C-SAC, C-CEM), continual RL algorithm (GPMM), POMDP algo (RNN-PPO), and SAC with a known context (FI-SAC). For soft failure experiments, we have increased the maximum applicable force by the factor of two. Best performances are highlighted in bold.
FI-SAC
NI-SAC
C-SAC
C-CEM
Drone
36.13 + 0.26
0.80 + 3.08
28.41 + 1.16
32.30 + 2.78
Intersection
572.09 + 20.25 499.62 + 19.98
555.11 + 20.21
529.75 + 78.12
Table 3: Mean + standard deviation of expected return for various algorithms and tasks over three seeds. Our contextual approaches, which marked by the letter C, are competitive with FI-SAC (SAC with full context information) and outperform NI-SAC (SAC with no context information).
CONCLUSION AND DISCUSSION
We studied a hybrid discrete-continuous variable process, where unobserved discrete variable repre- sents the context and observed continuous variables represents the dynamics state. We proposed variational inference algorithm for model learning using a sticky HDP prior. This prior allows for effective learning of an interpretable model and coupled with context distillation procedure offers a powerful tool for learning C-MDPs. In particular, we showed that the combination of the HDP prior and the context distillation method allows learning the context cardinality. We also showed that the model quality is not affected if the upper bound on context cardinality set is overestimated. Furthermore, we illustrated that the distillation threshold can be used regularization rrf parameter and it can also be used to merge similar contexts in an unsupervised manner. Furthermore, we present additional experiments in Appendix suggesting that our model can potentially generalize to non-Marke Man-Markovian and wate-dependen settings. While we presented several experiments in various environments, further experimental evaluation is required, c.g., using Benjamins et al. (202I).
We showed that continual and meta-RL approaches are likely to fail as their underlying assumptions on the environment do not fit our setting. The learned models do not appear to capture the complexity of Markovian context transitions. This, however, should not be surprising as these methods are tailored to a different problem: adapting existing policy model to a new setting. If the context is very different and or the contexts changing too fast then the continual and meta-RI. algorithms would struggle by design. We derived our policy by exploiting the relation of our setting and POMDPs. We demonstrated the necessity of our model by observing that standard POMDP approaches (i.e., modeling the context dynamics using an RNN) fail to learn the model. We attribute this behavior to the lack of effective priors and model structure. While in some cases it can appear that the RNN policy is effective, disregarding the context altogether has a similar effect.
Our model-based algorithm can be further enhanced by using synthetic one-step transitions similarly to Janner et al. (2019), which would improve sample efficiency. We can also use an ensemble of models, which would allow to constantly improve the model using gross-validation over models in the ensemble. However, evaluation of the model quality is more involved since the context is unobservable. In future, we also plan to extend our model to account for a partially observable setting, i... where the state only indirectly measured similarly to POMDPs. This setting would allow for a rigorous treatment of controlling from pictures in the context-dependers setting. While we show that we can learn unseen contexts without re-leaming the entire model, this procedure is not fully automated. Hence it can benefit from adaptation of continual learning methods in order to increase efficiency of the learning procedure.

---

Published as a conference paper at ICLR 2022
REFERENCES
Josh Achiam. Openai spinning up documentation, 2018. URL http::ts///spmmmmp.....opop//
com/ com/en/Dee http//ateestaaaae hatp//aggorttttaae tmms sac .htll
Maruan Al-Shedivat, Trapit Bansal, Yura Burda, Ilya Sutskever, Igor Mordatch, and Pieter Abbeel.
Continuous adaptation via meta-learning in nonstationary and competitive environments. In 6th International Conference on Learning Representations, ICLR 2018. 2018.
Karl Astrom. Optimal control of Markov processes with incomplete state information. Journal of mathematical analysis and applications, 10(1):174-205, 1965.
Taposh Banerjee, Miao Liu, and Jonathan P. How. Quickest change detection approach to optimal control in Markov decision processes with model changes. 2017 American Control Conference, IEEE, 2017.
ACC 2017, pp.
Philip Becker-Ehmek, Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for variational Bayes filtering. arXiv preprint arXiv:1905. 12434, 2019.
Carolin Benjamins, Theresa Eimer, Frederik Schubert, André Biedenkapp. Bodo Rosenhahn, Frank Hutter, and Marius Lindauer. Carl: A benchmark for contextual and adaptive reinforcement learning. arXiv preprint (tXiv:2110.02100 2021.
Abraham Berman and Robert Plemmons. Nonnegative matrices in the mathematical sciences. SIAM, 1994.
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah Goodman. Pyro: Deep Universal Probabilistic Programming. Journal of Machine Learning Research, 2018.
David M Blei, Michael I Jordan, et al. Variational inference for Dirichlet process mixtures. Bayesian analysis, 1(1):121-143, 2006.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the Americun statistical Association, 112(518):859--5) 2017.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference on Machine Learning, pp. 1613-1622. PMLR, 2015.
Haitham Bou-Ammar, Eric Faton, Paul Ruvolo, and Matthew E. Taylor. Online multi-task learning for policy gradient methods. In Proceedings of the 31th International Conference on Machine Learning, ICML, volume 32, pp. 1214-1214, 2014.
Michael Bryant and Erik Sudderth. Truly nonparametric online variational inference for hierarchical Dirichlet processes. Advances in Neural Information Processing Systems, 25:2699 2012.
Yash Chandak, Georgios Theocharous, James Kostas, Scott M. Jordan, and Philip S. Thomas. Learning action representations for reinforcement learning. Proceedings of the 36th International Conference Machine Leurning, ICML 2019, volume 97, pp. 941-950. PMLR, 2019.
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, pp. 1414-1425, 2020.
Samuel PM Choi, Dit- Yan Yeung, and Nevin L Zhang. Hidden-mode Markov decision processes for nonstationary sequential decision making. In Sequence Learning, pp. 264-287. Springer, 2000.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine Deep reinforcement
learning in handful of trials using probabilistic dynamics models. Advances in Neural Information Processing Systems, 31, 2018
Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Model-baseds reinforcement learning via meta-policy optimization. In 2nd Annual Conference on Robot Learning, CoRL 2018, volume 87 of Proceedings of Machine Learning Research, pp. 617-629. PMLR, 2018
10

---

Published as a conference paper at ICLR 2022
Bruno Castro da Silva, Eduardo W. Basso, Ana L. C. Bazzan, and Paulo Martins Engel. Dealing with non-stationary environments using context detection. In William W. Cohen and Andrew W. Moore (eds.), Machine Learning. Proceedings of the wventy-Third International Conference (ICML 2006), volume 148 of ACM International Conference Proceeding Series, pp. 217-224. ACM, 2006.
Matthias Delange, Rahaf Aljundi. Mare Masana, Sarah Parisot, Xu Jia, Ales Leonardis, Greg Slabaugh, and Tinne Tuytelaars. A continual learning survey: Defying forgetting in classification tasks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
Zhe Dong, Bryan Seybold, Kevin Murphy, and Hung Bui. Collapsed amortized variational inference for switching nonlinear dynamical systems. In International Conference Machine Learning, pp. 2638-2644 2647. PMLR, 2020.
Finale Doshi- Velez and George Dimitri Konidaris. Hidden parameter Markov decision processes: A semiparametric regression approach for discovering latent task parametrizations. In Subbarao Kambhampati (ed.), Proceedings of the Twenty-Fifth International Joint Conference Artificial Intelligence, IJCAI, pp. 1432-1440, 2016.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RL3: Fast reinforcement learning via slow reinforcement learning. CoRR, 2016. URL http://arxiv
(rg/abs/1611. 02779
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. In Advances in Neural Information Processing Systems, pp. 441-452, 2018.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML, volume 70, pp. 1126-1135. PMLR, 2017.
Emily Fox, Erik Sudderth, Michael Jordan, and Alan Willsky. Nonparametric bayesian learning of switching linear dynamical systems. Advances in neural information processing systems, 21: 457-464, 2008a.
Emily Fox, Erik Sudderth, Michael Jordan, and Alan Willsky. Bayesian nonparametric inference of switching dynamic linear models. IEEE Transactions Signal Processing, 59(4):1569 1585, 2011.
Emily B. Fox, Erik B. Sudderth, Michael I. Jordan, and Alan S. Willsky. An HDP-HMM for systems with state persistence. In Machine Learning, Proceedings of the Twenty-Fijh Fiith International Conference, volume 307, pp. 312-319. ACM, 2008b.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty deep learning. In international conference on machine learning, pp. 1050-1059. PMLR, 2016.
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 on machine learning, pp. 1861-1870, 2018.
Emmanuel Hadoux, Aurélie Beynier, and Paul Weng. Sequential decision-making under non- stationary environments via sequential change-point detection. In Learning over multiple contexts (LMCE), 2014.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.022: 502.02225, 2015.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data aug- mentation. In 202/ IEEE International Conference on Robotics and Automation (ICRA), pp. 13611-13611 IEFE, 2021.
Matthew Hausknecht and Peter Stone. Deep recurrent Q-learning for partially observable MDPs, In AAAI fall symposium series, 2015.
11

---

Published as a conference paper at ICLR 2022
Milos Hauskrecht. Value-function approximations for partially observable Markov decision processes. Journal of artificial intelligence research, 13:33-92 94, 2000.
Franz S Hover and Michael S Triantafyllou. System design for uncertainty. Mass. Inst. Tech- nol, 2009. URL http:////ooo..ttteee/////ree///a ennggneeinnnngngng
/
2 2-017j desi gn of wwwoofeeeeeettroeeeeeeeeeeeo www.ectooecccccccc.......... robotioobcccc- systems fall 1011-2009/ course-text/
Michael Hughes, Dae Il Kim, and Erik Sudderth. Reliable and scalable variational inference for the hierarchical Dirichlet process. In Artificial Intelligence and Statistics, pp. 370-378, 2015.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for POMDPs. In International Conference on Machine Learning, pp. 2117-2126. PMLR, 2018.
Martin Jankowiak and Fritz Obermeyer. Pathwise derivatives beyond the reparametterization trick. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research, pp. 2240-2249, 2018.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model- based policy optimization. In Advances in Neural Information Processing Systems, volume 32, pp. 12519-12530, 2019.
Khimya Khetarpal, Matthew Riemer, Irina Rish, and Doina Precup. Towards continual reinforcement learning: A review and perspectives. CoRR, abs/2012.13490, 2020. URL ltttps://arxiv
org/abs/2012. 13490
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnlo,, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv:/901.0576); 2019.
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. Advances in neural information processing systems, 28:2575-2588 2015.
Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 32(11))1238-11244 2013.
Ilya Kostrikov, Pytorch implementationn of reinforcement learning algorithms. hattps://githm
com/ wkkostrikoo//yyorccche.... a2c -wtpppppoo----aaaa acwtr-akkkrraaa. 2018.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing 10044.3364 13649, 2020.
deep reinforcement learning from pixels. arXiv preprini J
Gilwoo Lee, Brian Hou, Aditya Mandalika, Jeongseok Lee, Sanjiban Choudhury, and Siddhartha S. Srinivasa. Bayesian policy optimization for model uncertainty. 7th International Conference on Learning Representations, ICLR 2019, 2019.
Edouard Leurent. An environment for autonomous driving decision-making https://github
com/e wwl.lueeeetiiigee hhwghwyggwwhaaoee env 2018
Siyuan Li, Fangda Gu, Guangxiang Zhu, and Chongjic Zhang. Context-awarc policy reuse. In Edith Elkind, Manuela Veloso, Noa Agmon, and Matthew E. Taylor (eds.), Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, AAMAS, pp. 989-997, 2019.
Angelo Lovatto. gym-cartpole-swinggppp
oog.riiiuuuuoooooooooooiiooaaaaa for openai gym. https: /tt::://gitttt http:://h.....ooooooooooaoooaotr ca wwwppoae--ttoinuuuuuuue 2019.
Timothy E Menke and Peter S Maybeck. Sensor/actuator failure detection in the Vista F-16 by multiple model adaptive estimation. IFFE Transactions on aerospace and electronic systems, 31 (4):1218-1229, 1995.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 112.502,,,, 2013.
12

---

Published as a conference paper at ICLR 2022
Anusha Nagabandi, Chelsea Finn, and Sergey Levine. Deep online learning via meta-learning: Con- tinual adaptation for model- based RL. In International Conference on Learning Representations, 2018.
Dimitrios Noutsos. On Seerron-Frobente property of matrices having some negative entries. Linear Algebra and its Applications, (1102-3):132-1115 2006.
Sylvie CW Ong, Shao Wei Png. David Hsu, and Wee Sun Lee. Planning under uncertainty for robotic tasks with mixed observability. The International Journal of Robotics Research, 29(8):0033-1088 2010.
Sindhu Padakandla, Prabuchandran K. J., and Shalabh Bhatnagar. Reinforcement learning in non- stationary environments. CoRR, abs/1905.0393 2019. URL htttp://arxiv.org//aa T905.0397 03970
Jacopo Panerati, Hehui Zheng, SiQi Zhou, James Xu, Amanda Prorok, and Angela P. Schoellig. Learning fly-a a gym environment with pybullet physics for reinforcement learning of multi- agent quadcopter control. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2021.
Luis Pineda, Brandon Amos, Amy Zhang, Nathan O. Lambert, and Roberto Calandra. Mbrl- lib: A modular library for model-based reinforcement learning. Arxiv, 2021. URL https //arxiv. org/ abs/ 2 2104. 10159
Josep M Porta, Nikos Vlassis, Matthijs TJ Spaan, and Pascal Poupart. Point-based value iteration for continuous pomdps. Journal of Machine Learning Research, 7(Nov): 2267, 2006.
Shenghao Qin, Jiacheng Zhu, Jimmy Qin, Wenshuo Wang, and Ding Zhao. Recurrent attentive neural process for sequential data. arXiv preprint arXiv:1910.09323, 2019.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off- off-policy meta-reinforceemnt learning via probabilistic context variables. In Proceedings of the 36th International Conference on Machine Learning, ICML, volume 97, pp. 5331-5340, 2019.
Matthew Riemer, Ignacio Cases, Robert Ajemian, Miao Liu, Irina Rish, Yuhai Tu, and Gerald Tesauro. Learning learn without forgetting by maximizing transfer and minimizing interference. In 7th International Conference Learning Representations, ICLR 2019, 2019.
Hippolyt Ritter, Aleksandar Botev, and David Barber. scalable laplace approximation for neural networks. In 6th International Conference on Leurning Representations, ICLR 2018-CC 1018-Conference Truck Proceedings, volume International Conference on Representation Learning, 2018.
David Rolnick, Arun Ahuja, Jonathan Schwarz, Timothy P. Lillicrap, and Gregory Wayne. Experience replay for continual learning. Advances in Neural Information Processing Systems, pp. 348-358, 2019.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. ProMP: Proximal meta-policy search. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimovv Proximal policy optimization algorithms. arXiv preprint arXivv1707.00347, 2017.
Jayaram Sethuraman. A constructive definition of Dirichlet priors. Statistica sinica, pp. 639-650, 1994.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(75877:4844488) 2016.
Pranjal Tandon. Pytorch implementation soft-actor-criticc httpttrrrrrtt..ttt.ps///giiiiii https:://githuu..com////
pytorch- sottt-soff--oo wwwooftttttr-erreooo
2018.
13

---

Published as a conference paper at ICLR 2022
Yee Whye Teh, Michael I Jordan, Matthew J Beal, and David M Blei. Hierarchical Dirichlet processes. Journal of the American Statistical Association, 101(476):1566-33 1581, 2006.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681-688. Citeseer, 2011.
Annie Xie, James Harrison, and Chelsea Finn. Deep reinforcement learning amidst lifelong non- stationarity. arviv preprint arXiv:2006. 10701, 2020.
Mengdi Xu, Wenhao Ding, Jiacheng Zhu, Zuxin Liu, Baiming Chen, and Ding Zhao. aask-agnostic online reinforcement learning with infinite mixture of Gaussian processes. Advances in Neural Information Processing Systems, 33, 2020.
Denis Yarats and Ilya Kostrikov. Soft actor-critic (sac) implementation in pytorch. https:// github. wwt.c///eeniaaaaattt//orroroaee. sac, 2020.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin
Gal, and Doina Precup. Invariant causal prediction for block MDPs. In Proceedings of the 37th International Conference on Machine Learning, ICMI., volume 1 119, pp. 11214-11224, 2020.
Pengfei Zhu, Xin Li, Pascal Poupart, and Guanghui Miao. On improving deep reinforcement learning for POMDPs. arXiv preprint arXiv:1704.07978. 2017.
Luisa M. Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep RL via meta-learning. In 8th International Conference on Learning Representations, ICLR 2020, 2020.
14

---

Published as a conference paper at ICLR 2022
Appendices
Detailed Literature Review
al
B Proofs
a4
Proof of Theorem
a4
B.2 Derivation of dynamic programming principle and proof of Theorem
a6
B.3
Performance gain for observable contexts
a8
Algorithm Details
a9
C.I
Slightly more details the Hierarchical Dirichlet Processes
a9
Variational Inference for Probabilistic Modeling
a10
Justification the variational distributions
4
Context Distillation
al3
C.5
MPC using Cross EotsopE Method
a15
C.6
Soft-Actor Critic
a15
Experiment Details
D.1
Learning algorithms
a16
D.2 Environments and their Models
a16
D33
Hyper-parameters
a19
Additional Experiments
E.1
HDP is an effective prior for learning an accurate and interpretable model
a19
E.2
Distillation acts as regularizer
a21
E.3
Context cardinality model complexity
a22
Breaking the assumptions in Hidden Markov Models
Learning new contexts
E.6
Comparing to POMDP and C-MDP methods
a24
E.7
Experiments with a larger number of contexts
a25
A DETAILED LITERATURE REVIEW
For convenience, we reproduce the definitions and assumptions from the main text to make the appendix self-contained. We define a contextual Markov Decision Process (C-MDP) as a tuple M (C, S. A, Pe, Ps, Ra), where is the continuous state space; AA is the action space; Ta E [0. 1] is the discount factor; and C denotes the context set with cardinality |C]. In our setting, the state transition and reward function depend on the context, i.e., Ps
x A S [0. 1]. R. C S x A + R. Finally, the context distribution probability Pe T, C - [0. is conditioned on To the past states, actions and contexts (8g: dgg Co...., 01-1, Ct-1, ) Our definition generalization of the C-MDP definition by Hallak al. (2015), where the contexts stationary, i.c.. : [0, 1]. We adapt our definition in order to encompass all the settings presented by Khetarpal
al. 2020). where such C-MDPs were used but not formally defined.
Throughout the paper, we will restrict the class of C-MDPs by making the following assumptions: (a) Contexts unknown and not directly observed (b) Context cardinality is finite and know its upper bound K: (c) Context distribution Markovian. In particular, we consider the contexts e₁. representing the parameters of the state transition function Og, and the context set C be subset of the parameter space @ To deal with uncertainty, we consider set such that: a) |C| K > [C]: b) all its elements 0₂ are sampled from a distribution H (A), where À is hyper peremeerreeter. Let [0...., K)) be the index variable pointing toward partieular parameter vector 0z.- We thus
al

---

Published as a conference paper at ICLR 2022
write the environment model as:
20 Po Cat(po), 21 2t-1 {p;}, Cat(p
(Ala)
lcl {0x k=k ~ 0,-11 Oz.), Os A H
8t 8;-1, a1-1;
(Alb)
For convenience we also write R (p,,... Piyl representing context transition operator, even
[C] is countable and infinite.
Other restrictions the space of C-MDPs lead to different problems and solutions. We present some settings graphically in Figure Al adapted from the review by Khetarpal et al. (2020).
B
c
D
20
20
2
Figure A1: Graphical models for C-MDP modeling. In all panels a stands for action, 8 state, o- observation, context, while the operator' indicates the next time step. A: The context variable z evolves according to Markov process; B: The context variable z is drawn once per episode; C: The context variable z evolves according Markov process and depends on state and/or action variables; D: POMDP.
Setting A. In this setting, a Markov assumption is made on the context evolution. Only a few works make such an assumption similarly our work. For example, Choi al. (2000) assume a discrete state, action and context spaces and fix the exact number of contexts thus significantly limiting applicability of their approach. We can also mention the work by Xu et al. (2020), where the individual points are classified to different context during the episode. However, the agent explicitly assumes that the context will not change during the execution of its plan. The Markovian context evolution can also be related to switching systems (cf. [Becker- et al. (2019); (11:):Foos Fox et al.) (2008a); Dong al. (2020)). where the context representing the system mode. While we take inspiration from these works, we improve the model in comparison to Bccker-Ehmck et al. (2019); Fox al. (008a) by using non-linear dynamics, comparison to Fox et al. (2008a) by proposing the distillation procedure and using deep learning, in comparison to Becker-Ehmek et al. (2019); Dong al. 2020) by using the hierarchical Dirichlet process prior.
Setting B. Here, an episodic non-stationarity constraint on the context probability distribution Pc is introduced, that is, the context may change only (without loss of generality) at the start of the episode. Episodic non-stationarity constraint is widely adopted by meta-RI. (Doshi Velez & Konidaris 2016
Finn et al. 2017 Al-Shedivat et al. [2018] Rothfuss et al.) [dL]]2011// Clavera et al. 22018 Rakelly et al.
2019 Zintgraf et al., 2020 Xie et al. (020)) and continual RL (Chandak et al. 2020; Nagabandi
et al. communities. We distinguish optimization based and context-based Meta-RI families of methods.
2018 2018| Khetarpal et al. 2020 Bou-Ammar et al. 2014 Riemer et al. 2019 Rolnick et al. 2019)
Optimization-based Meta-RL methods (Finn et al.) 2017 Al-Shedivat et al.) 2018 Rothfuss et al.
2019 Clavera al. 2018) sample training contexts from stationary distribution Pe and optimize model or policy parameters, which can be quickly adapted to test context. In contrast, context- based Meta-RL methods (Rakelly et al.) 2019 Zintgraf et al.) 2020 Xie al. [020 Duan et al.) 2016) attempt to infer a deterministic representation or probabilistic beliei on the context from the episode history. We also remark recent work on domain shift while controlling from pixels Hansen & Wang (202T): Kostrikov et al. (2020). In these papers the authors assume that some pixels may change from an episode to an episode (e.g., red color becomes blue), but the underlying dynamics stay the same. Besides the episodic nature of the context change, this setting differs from ours two fronts. First, controlling from pictures constitutes a POMDP problem, where the true states (positions, velocities, accelerations) are observed through proxy (through pixels) and hence need to be inferred. Second, the underlying dynamics stay the same (Hansen & Wang 202T: Kostrikov et al.] 2020) and only observations change. Our case the opposite: the underlying dynamics change, but the observation
a2

---

Published as a conference paper at ICLR 2022
function stays the same. In future work aim to extend our methods to control from pixels. We also note the work by Ball et al. (2021), who proposed to learn a contextual policy by augmenting an offline world model. In the online setting the algorithm had to learn only the current context thus improving efficiency of the learning procedure while delivering good performance. Continual (lifelong) RL. mainly adopts context incremental setting, where the agent is exposed of contexts (Delange et al., (021). While the agent's goal is still to adapt efficiently to the unseen contexts, the emphasis on overcoming catastrophic forgetting, i.e., maintaining a good performance on old contexts while improving performance on the current one (Rolnick et al. 2019 Riemer et al. 2019).
sequence
We further remark that the Markovian context model and episodic non-stationarity assumptions are often incompatible. Indeed, learning the transition model for Pe in C-MDPs with an episodic non-stationarity constraint is somewhat redundant because Pe can be assumed to be a stationary distribution rather than a Markov process. On the other hand, if the context changes according to a Markov process then the episodic non-stationarity assumption may not be enough to capture the rich context dynamics. This can lead not only to suboptimal policies, but also to policies not solving the task at all. Xie et al. (020)) tried to combine the two settings by taking a hierarchical view on non-stationary modeling. In particular, they assume that the context changes in a Markovian fashion, but only between the episodes and during the episode the context does not change. This setting allows to model a two time-scale process: the context transitions on a slow time-scale prescribed by the episodes, while the process transitions on fast time-scale prescribed by the steps. While this approach has its merits, it also has the some limitations, c.g., the hierarchical model is artificially imposed on the learning process.
Settings C and D. Many frameworks can fit into the settings C and D, (c.f., Ong et al. 2010); Chandak et al. (hang et al. (2020)), however, this makes it hard to compare the differences and similarities between them. We will refer the reader to the review by Khetarpal et al. (2020) for a further discussion. We make a few comments, however, on change point defection methods and on
the relation to POMDP formulation (Astrom) 1965 Hauskrecht] 2000) in the case of unobservable contexts. Change-poin detection methods can be readily used for the context estimation in an online
fashion (da Silva et al. 2006 [Hddox et al. [0144] Banerjee et al. 2017: Padakandla et al. 2019 Li
ct 2019), however, these methods either track thange-points in a heuristic way (da Silva et al. (0066
Hadoux al. 2014) assume strong prior knowledge the non- non-stationarity (Banerjee et al. 2017
Padakandla et al. 2019). A POMDP is defined by tuple Mp = {X, A, O, Px, Px,Pe Po. R, Ya. [
where X, A. O are the (unobservable) state, action, observation spaces, respectively; Px : A x
X' - [0, 1] is the state transition probability; Po X' x A x O - [0, 1] is the conditional observation probability; and p(zo) is the initial state distribution. In our case, the state is 2, (st, 21) and the observation is o = St. A key step in POMDP literature is estimating or approximating the belief, i.c.. the probability distribution of the current state using history of observations and actions (Hausknecht) & Stone, 2015 Zhu et al.) 2017 Igl et al. 2018). Since the belief is generally dufinite-dimnnsion is prudent to resort to sampling, point-estimate 2015. Zhu al.) 2017 [gl et al. 2018). In our case, this is not necessary, however, as the context is discrete variable. We further stress, that while general C-MDPs can be represented using POMDPs, this representation may not offer any benefits at all due to complexity of belief estimation.
or other approximations (Hausknecht & Stone)
Finally, we can classify the literature by cardinality of C. Some Meta-RL RL algorithms mainly consider a continuous contextual set by implicitly explicitly parametrizing C with real-valued vectors (Doshi
Velez & Konidaris) 2016 Zintgraf et al.) 2020: Xie et al. 2020; Rakelly et al., (019). On the other hand, exploiting discrete C can be motivated by controlling switching dynamic systems (Fox et al.
2008a) and may achieve better interpretability. Due to complexity, however, some works fix the
2017 Padakandla al. 2019 Lee et al. contextual cardinality C| (Choi al. 2000 Banerjee al.
2019), or infer [C] from data in an online fashion (Xu et al. 2020 da Silva et al. 2006 Hadoux et al.
2014). Besides, most works in continual RL adopt a setting, where there is no explicit specification
on C. but the agent can directly access the discrete contexts of all time steps/episodes (Rolnick et al.
2019 Riemer et al. 2019). In contrast, our approach learns the context evolution from data.
a3

---

Published as a conference paper at ICLR 2022
B PROOFS
B. PROOF OF THEOREM I
For completeness, we restate the theorem below.
Theorem A1 Consider Markov chain p' pp-11 R with stationary distribution pd and distilled I {ilp, 2 Edistil) and spurious I2 = p < Edisitl state indexes, respectively. Then a) the matrix R RI,II, + Rug Tg (I - RIg,I, is valid probability transition matrix; b) the Markov chain p = p - R such that its stationary distribution pr x Pi
Let us first provide an insight into our technical result. order to so consider the Markov chain evolution:
11-1RR, RI, I Pi, RIG.I11 P\z = PT, 1-1 RI, + PT, RI.II.
p = Pz,
Now assume that ph,, s Př, for all large enough 4, which leads to
p P RIA In PIa RIII p ~ pz, 1-1 RI,I, + Pia RIg.I4
If Pł, is small enough, then the Markov chain will rarely end up in the states with indexes Iz Meaning that we can remove these states, but we would need to take them into account while computing the new probability transitions. Furthermore, we need to do so while obtaining new Markov chain the process. The solution this question is straightforward, i.e., we can solve for P7, to obtain:
ph. = - (RI, P + RI -I-RIR,.5 RIA.TA) RTI,,11)
Remarkably, the resulting process is also Markov chain! In order to verify this we need to prove the following:
a) the matrix I - RI,, is invertible
b) the matrix (I - -I,,,, is nonnegative;
c) the right eigenvector of the matrix RI, + Rr, RIZ,I,) -1 can be chosen be the vector of ones 1.
We will need develop some mathematical tools in order prove these statements. For matrix A E R2 the spectral radius r(A) is the maximum absolute value of the eigenvalues X of A and r(A) max; [A)] The matrix A E R= is called nonnegative if all its entries nonnegative and denoted as A > 0, if at least one element is positive we write A > 0 and if all elements are positive we write A > 0. The matrix A =x called reducible there exists
B
C
permutation matrix T such that TAT
for some square matrices B and D. The D
0
matrix is called irreducible if such permutation matrix does not exist. The matrix A E Raxx is called M-matrix, if all its off-diagonal entries are nonpositive and can be represented as = - with 2 r(B) (Berman & Plemmons) 1994). Interestingly the inverse of an M-matrix is always nonnegative (the opposite generally talse) Berman & Plemmons 1994). We will also use the following results:
Proposition A1 (Perron-Frobenius Theorem) Let A B be a nonnegative matrix, then the spectral radius r(A) is eigenvalue of A and the corresponding left and right eigenvectors can be chosen be nonnegative.
If A is additionally irreducible then r(A) is simple eigenvalue of A and the corresponding left and right eigenvectors can be chosen be positive.
Proposition A2 Consider A, B E RN and let > B 2 0, then r(A) 2 r(B). If additionally A is irreducible then the inequality is strict rí A) > r B).
a4

---

Published as a conference paper at ICLR 2022
This result is well--nnww, but for completeness show the proof here (we adapt similar technique to Noutsos (2006)). Let b be the right nonnegative eigenvector of B corresponding to r(B) and let a be the left nonnegative eigenvector of A corresponding to r r(A). Now we have
Proof:
r(A)a' a" Ab
a Bb r(B)a
If a is positive, then the first part is shown. If a²'h O, then by perturbing A and B to A' and B' in such a way that for their corresponding cigenvectors we have (a')"b' Hence the first part of the proof is shown.
make a continuity argument
Now consider the case when A is irreducible. According to Proposition A1 since A is irreducible the spectral radius r r(A) is simple eigenvalue, the corresponding eigenvector a can be chosen to be positive. Let us prove the second part by contradiction and assume that r r(A) = r(B), which implies that a Ab- a Bb and consequently Ab Bb since Ab 2 Bb and a >> Now since we also have A > B this means that the vector b has at least one zero entry. Since we also have Ab = r(A)b, the vector b is an cigenvector of A corresponding to r(A) and has zero entries by assumption above. This contradicts Cerroo-Frobenius theorem since the eigenspace corresponding to r(A) is a ray (since r(.A) is a simple eigenvalue) and b does not lie in this eigenspace. Therefore, (A) (B).
Now we proceed with the proof. We note that the stationary distribution (a positive normalized left eigenvector of the transition matrix) is unique if the Markov chain is irreducible (i.e., the transition matrix is irreducible) due to Proposition A It is also almost surely true if the transition model is learned. This is because the subset of reducible matrices is measure zero in the set of nonnegative matrices.
Recall that have a Markov chain p' p with stationary distribution pe, that is p p R. Consider the index sets: the distilled context indexes I₁ = filp? 2 Edistil J and the spurious context indexes I2 filp; < Edistil)
First, we will show that the matrix N Rr,T, RT, J.g (I RT is nonnegative. Since the matrix RI,I, 2 is nonnegative submatrix of R, due to Proposition A2 we have that r r(RI,IZ)) < r(R) 1. This means that RI,I, is an M-matrix, which implies that it is invertible and its inverse is nonnegative matrix. Hence R is nonnegative (Berman & Plemmons 1994).
Since R describes a Markov chain we have RI 1 (where 1 is the vector of ones) and p R p
11, RI, I, 1T, Rí, 2, 1
(A2a)
17, RI 12 + RI, 12 Pf, Pi, RI + DT, RTI
(A2b)
(A2c)
p p7, II Dt, RI,LII
(A2d)
Now using elementary algebra can establish that Rli, lzy
R1g = (RI,,,, + RI (I I- RIG,Is) -RI, 11, =
RI l7, + Rr Is (I RIa
RI, In 17, due to A2b
RI, 1I, + Rr, I, 1z, due A2a
11,
Similarly for P R p2, we have:
P 20
R
Pt RI,, + RI, RI,,1,) RI,,
p RI,I, +T Rr, 1n RI,,I,)
RI,II
due A2d
p
RTY,I
+ P R
A2c
PT.
Since 7 pT, the normalized vector p is the stationary distribution of the distilled Markov chain and hence p p This completes the proof.
a5

---

Published as a conference paper at ICLR 2022
B.2 DERIVATION OF DYNAMIC PROGRAMMING PRINCIPLE AND PROOF OF THEOREM 2
For completeness we reproduce the theorem formulation below
Theorem A2 a) The belief of z can be computed as p(z++1 = where (bi-1); X N E(p(8:-1 18,, O (z) (bi) where (bī): are the entries of bi; b) the optimal policy can be com- puted as т (s, ba) argmax, (s, ba a) where the value function satisfies the dynamic program- ming principle (s,, b; at) = r(st, by at) + + + N₂ maXa4+1 @(+++1, b1 a(+1) ds
The following definitions and derivations are in line with previous work by Hauskrecht (2000): Porta al. (2006). We introduce the complete information state at the time i follows:
I (bo; Oge, dcct);
where bn )00 (8o). In order to tackle the intractability of the complete information state, one can use any information state that sufficient some sense:
Definition A1 Consider partially observable Markov decision process (X,A,O.Px:Po. R. bu)- Let I be an information state space and I x A x O + I be an update function defining an information process = (( a1 0.). We say that I sufficient information process with regard to the optimal control if it is an information process and for any time step t, it satisfies
p( p( ) p(o: a.-1) = p(o: I a1-).
As sufficient information state of the process I we will use the belief b, defined as follows:
b. p( ) = p(a:ooge act)) p( oge, act)p(== o act).
Effectively, we introduce the belief by = p(a, Ose, act) of the state st and the belief bj p( ace)) of the state given observation or. However, since o 8t consider only the belief of 24:
b p(s; = p( ace) (0))
b; = ) = p( act).
One of the features of our model that the updates of the beliefs can analytically computed.
Lemma A1 The belief bj )() is sufficient state information with the updates:
(bi+1)s N, = Σ p(st+1le ze+1 = =i ir. pp:
Proof:
First, let us check that the belief b, satisfies both conditions of Definition A] By definition
of our belief we have p( Ie,) = b = p( I( which satisfies the first condition, as for the
second condition we have:
[ 1-x
p(o.|IE[1,
p(o,10 a - D.
= p(or) az-1).
2--1EX
Now straightforward derivations yield:
b = = 8441, x 2(2441 14.1 a.)=
p(s141l2441, If an) Σ p(2cm122)pp2 p(s 2411, a.) Σ ((
completing the proof.
a6

---

Published as a conference paper at ICLR 2022
As discussed by Porta al. 2006), POMDP in continuous state, action and observation spaces also satisfy Bellman dynamic programming principle, albeit a different space. Recall that the control problem is typically formulated
T
Tb
t=0
where 7 {Zo, do..... .0, ar -1: ΦT) and the horizon T can be infinite. As we do not have access to the state transitions we need to rewrite the problem in the observation or the belief spaces. We have
a.) 0++1 ~ p(1++1 De. (( 12))
Eæo
Σ (
r'
T
E
y rr (be, az) = 0(,
1=0
r (br, ar) _(bb-1,, (1)
E,
where =ibu, Ou ao.
:.. 01-1,00111110000 0r) and:
/ re(2,, (x) dæ.
P. Pb a:)
Given this reparameterization we can introduce the value functions and derive the Bellman equation similarly to Porta et al. (2006), which can be written in terms of the Q-function as follows:
Q(b bj. a.) / Pr (r|s ar ds,t
of pioctibb, bj a,) max Q(bf+1 bi+1 ar41 dd0+11
We, however, have additional structure that allows for simplified value functions. First note that (b, a,) = (st bi a,), i.c., our reward depends directly on the observation and the belief of
bt
b
Now we need to estimate IE' maxa44 Q(b bi+1 a++1), where the expectation taken over new observation, the probability dis- tribution of which can be computed as:
St
St++1
p(0 b. bi, a.) p(s St, bi a.)
Σ p(s 2111 (b)))))(())) No.
T
Figure A2: Graphical model for policy optimization
This allows rewriting the second part of the Bellman equation as follows:
/ plo0++11b bi, a.) max Q(bi+1 bi41 act ddo
ac+1
/
Σ N max ((5 bis1, ax+1) d8 a++1
where N₁, bi+1 depend on s441. Finally, we have:
Q bj, a1) r(s,, a.) + 7 / /
N max Q(s bi ds asll
which completes the proof.
a7

---

Published as a conference paper at ICLR 2022
B.3 PERFORMANCE GAIN FOR OBSERVABLE CONTEXTS
It is not surprising that observing the ground truth of the contexts should improve the maximum expected return. In partieular, even knowing the ground truth context model we can correctly estimate the context 2t+1 only a posteriori, i.e., after observing the next state 81+1- This means that with every context switch we will mislabel a context with high probability. This will lead to a sub-optimal action and performance loss, which the following result quantifies using the value functions.
Theorem A3 Assume we know the true transition model of the contexts and states and consider two settings: we observe the ground truth and we estimate it using by Assume we computed the optimal boseel-based policy T 81, bj) with the return R and the optimal ground-truth policy ggg(-188 (-181, Ze+1) with the corresponding optimal value functions (8 2) and Qut (a,(s,,2 2, a), then:
gg
M
Es, s0 Ve (Su, 21) R R2E, E, B,,att ot Or
Σ 7m'
(Q(stm
Sem +1, arn) Q (m 2m +1, atm )),
m==
where M is the number of misidentified context switches in a trajectory
Proof: Let us consider the best case scenario. As we assume that we have access to the ground truth model for computing the policy T-|1 bj), has the highest probability mass in the vector bj. That is, we can assume that we can identify the correct context a posteriori. In effect, we can use a priori estimate of ž++1 by transitioning the next time step. i.e., using the vector b; R. We can also assume that the action distributions of the policy и and the ground truth policy Tgt are identical provided our a priori estimate of the context and the ground truth context are the same. This, however, is almost surely not true when the context switch occurs, as we need at least one sample from the transition model in the new context. Now we can estimate the effect of this mismatch on the performance, this will provide us with lower bound on the performance gain.
assume that the ground truth value of 2
Let (s). (st a) be the optimal value functions for the ground truth policy नet satisfying the Bellman equation:
(st 2) = Ex Ex (aapp (r(s, a, + (
Ve (8, = Ba Ozn a).
Consider particular realization of the stochastic context variable = (which is independent of st, a.) and assume the context switched only once at 41 Then we have
T
R(s0, 21) = Ba xx)))
Σ t=0
(rr(St, a22 81+1) = Eajay Nat
Σ 1
rrr( e ai, S1+1)+
+ Orsy 11) titt t- ( ar, ae, suni)
(gu 21) - a "EENTTTT 1 (r(s), : a, 86,11 + 7 VVe (8,,41, + + * -12 (r(St,, a, 81.+1) + V (8t +1, 2t, #t))
(80, 21) 7"1 " Bn o (Qez (8t, Et +1 +1, a Qut (8t Etg +1,
In effect, we are using the Qge function estimate the performance loss of one mistake. The same procedure can be repeated for context realizations with M misidentified switches, where the number M depends on the realization of the context variable
R( 21) =
M Σ
V(80, 21)
7
E ai
as (Q( 8t... 21, ag Q(8tmm 26 at ))
1==1
Now averaging over the context realizations proves the result.
a8

---

Published as a conference paper at ICLR 2022
C ALGORITHM DETAILS
There are three main components in our algorithm: the generative model derivation (HDP-C-MDP), the model learning algorithm with probabilistic inference and the control algorithms. We firstly briefly comment on each on these components to give an overview of the results and then explain our main contributions to each.
In order to learn the model of the context transitions, we choose the Bayesian approach and we employ Hierarchical Dirichlet Processes (HDP) as priors for context transitions inspired by time-seriss modeling and analysis tools reported by Fox et al. (2) We improve the model by proposing a context spuriosity measure allowing for reconstruction of ground truth contexts. We then derive model learning algorithm using probabilistic inference. Having a model, we can take off-the-shelf frameworks such (Pineda et al.) (02I), which can include a Model Predictive Control (MPC) approach using Cross-Entropy Minimization (CEM) (cf. CChua et al. (208)) and Appendix C.5 or
a policy-gee -rradient approach Soft-actor critic (cf. Haarnoja al. (2018) and Appendix (.6)) which are both well-suited for model-based reinforcement learning. While MPC can be directly applied to
our model, for policy-based control we needed to derive the representation of the optimal policy and prove the dynamic programming principle for our C-MDP (see Theorem 2 the main text and its proof in Appendix B.I. We summarize our model- based approach Algorithm Al
Algorithm A1: Learning to Control HDP-C-MDP
Input: Edistill distillation threshold, warm number of trajectories for warm start, Nurj number of newly collected trajectories per epoch, Nepochs number of training epochs, AGENT policy gradient MPC agent
Initialize AGENT with RANDOM AGENT, D = 0:
for i 1,..., N Nepochs do
Sample Niraj (N (NWarm if = 1) trajectories from the environment with AGENT;
Nusii Set Duew {(s', a²)) Fi= where sf {s and a'
{aa) (a:)²= are the state and
action sequences in the i-th trajectory. Set D = D U D new;
Update generative model parameters by gradient ascent on ELBO in Equation 4
Perform context distillation with Edistill;
if AGENT is POLICY then
Sample trajectories for policy update from D; Recompute the beliefs using the model for these trajectories;
Update policy parameters end end
return AGENT
C.1 SLIGHTLY MORE DETAILS ON THE HIERARCHICAL DIRICHLET PROCESSES
A Dirichlet process (DP), denoted as DP H). is characterized by a concentration parameter and base distribution H (A) defined over the parameter space O A sample G' from DP (1,,1 H)) is probability distribution satisfying (G) G(A,)) ~ Dir(" (AI), HH (A,)) for every finite measurable partition A₁, ..., A, of θ where Dir denotes the Dirichlet distribution. Sampling G is often performed using the stick-breaking process (Sethuraman, 1994) and constructed by randomly mixing atoms independent and identically distributed samples O₂ from H:
2k Bcta(1,7), Bx 2k
[1- v)). G-Σ 3500
(A3)
1=1
where 89, is the Dirac distribution at Og. and the resulting distribution of B3 (B....Bb<)) is called GEM(7) for Griffiths- Engen-
(Teh et al., 2006). The discrete nature of G motivates the application of DP as a non-parametric prior for mixture models with an infinite number of atoms Oş- We note that the stick- brecking procedure assigns progressively smaller values to 3% for large k thus encouraging smaller number of meaningful atoms.
The Hierarchical Dirichlet Process (HDP) is a group of DPs sharing a base distribution, which itself is sample from DP: G DP H), Gj DP(a, G) for all j = 0, 2,,.,. (Teh et al.
a9

---

Published as a conference paper at ICLR 2022
2006). The distribution G guarantees that all G3 inherit the same set of atoms, i.e., atoms of G, while keeping the benefits of DPs in the distributions Gj. HDPs have received a significant attention the literature (Teh et al. 2006 Fox al. 2008b a with various applications including Markov chain modeling.
In our case, the atoms {0x} forming the context set C are sampled from H It can be shown that random draw Gj from (a, G)) can be done using Pj GEM(() and 02 G. However, since 0g sampled from C, Gj also a distribution over and
Gj
=
Σ
Psh00
K=
A=
for some Pyk, which can be sampled using another consider its modified version introduced by Fox et al. (2011):
construction (Teh et al. 2006). We
Σ 0. + kdjz )) Pyk Pijk (1
Mik K, Beta aBBe kōje,
p32)s
(A4)
where 2 1, 2 O, ōjz is the Kronecker delta, the parameter x 2 0, called the sticky factor, modifies the transition matrix priors encouraging self-transitions. The sticky factor serves as another measure of regularization reducing the average number of transitions. Thus DP G) can serve as the prior for the initial context distribution Pu and each row Pj in the transition matrix R.
21
22
2T
Vk
80
1
32
ST
iT-1
N
Figure A3: A probabilistic model for C-MDP with Markovian context
In summary, our probabilistic model constructed Equations AT] A4 and illustrated in Fig-
ure A3 as graphical model. We stress that the HDP in its stick-breaking construction assumes that
Cl is infinite and countable. In practice, however, make approximation and set C K with large enough K.
C.2 VARIATIONAL INFERENCE FOR PROBABILISTIC MODELING
Recall that our context MDP is represented as follows
2 žt, {p Mul(pz.); 2p Pa ~ Mul(po),
0
H()):
81-1, a:-1; =t; (0 p(..)11111 1,))
and depict our generative model graphical one in Figure A3 Also recall that the distributions Pj have the following priors:
Σ aB; + kōσ² ))
Pjk = Pik
pg));
Mip a, K, ~ Beta aBe + kō
Ph Beta(L. 7),
By = 1
(1 (1-11)) vz))
(A5)
a10

---

Published as a conference paper at ICLR 2022
where & N21, N20, t E N20 and özz is the Kronecker delta function.
We aim to find a variational distribution g(v, μ.. 0) to approximate the true posterior p(v, μ. ( where D = {(s', ] is a data set, s' [s]]]=- and a' {a)) the state and action sequences in the i-th trajectory. This is achieved by minimizing K.C (g(,, p μ. (())), or equivalently, maximizing the evidence lower bound (ELBO):
N
ELBO Bi(
Log log pogpp(s'aa" μ. 0)
KCC (c(v. ne, 0))
The variational distribution above involves
random variables v, O. To reach
tractable solution, we set (C] K and exploit a mean-field truncated variational distribution (Blei
et al. 2006 [Hughes al. 2015 Bryant & Sudderth 2012).
We construct the following variational
distributions:
K II
K
(v,,,,
M4 0) ((V)(())((), a(0)0)
8(0n Ok), a(viD) = II o(vx(De), dluk 1) = 1, k=1
k=1
K-1 II II j=0 k=1
((p())
Beta
Pik Mykk
Pskk Prj
Σ Ass
005K= = 1)
(A6)
where the hatted symbols represent free parameters. Random variables shown in (A6) are conditionally independent of the data, and thus can be discarded from the problem.
We maximize ELBO using stochastic gradient ascent. In particular, given sub-sampled batch B = {(s', aa) }>> the gradient of ELBO is estimated as:
N D. ELBO
ù. Bq(u) pog la', m. (a. Vun (
B
+ V l() + V log
where we apply implicit reparameterization method (Figurnov al. 2018 Jankowiak & Obermeyer 2018) for gradients with respect to the expectations over Beta distributions. For computing gradient with respect to the likelihood term log p(s"la,, μ, 0), we exploit a message passing algorithm to integrate out the context indexes zi.T- We present the details of the gradient computations in what follows.
Gradient of log (ogp(s" "aaaa,,,, have:
μr, (1,0), 0). We drop the dependency on μ. θ in the following derivations. We
p(st v log s(zi's',aa)
V log Logp(s''aa = (a) V Voggp(s s a²) = Ep(za(aiaa))
( V Vlog((s', z²|a')] Ep(z")s)
[V (oggp2>'1'' [s
where z² [=f/=1 the context index sequence. Since the second term equals zero, have:
Vlog p(s" (a') Ep at)) [V Vogg(ss 2""
Ep(efle V logg(s((s) ag. =f)((()))
a log Wooe pl
Context index posteriors p( agglss a') and p( 2f lss,, a') required to compute the above expectation can be obtained by the message passing algorithm. The forward pass can written as:
m=f=1) = p(z), si sb. a = p(s) 186, ag =1)r(z()
(2(2 = p(2), si. 1)) 1) Σ p(2,, 2|_1 si si s/ ab.-
st
= =p(s/|s/ ai-1 2/) Σ p(=))=/_1))) (z(
all

---

Published as a conference paper at ICLR 2022
The backward pass can be written mo(21) = 1
ai- m = Ep(s) p(s) s
ms ms(zi-1)= = p(. s
#i
Combining the forward and backward messages, we have:
p( a') x p(zi. sirr b a') = my mz(2))mo
p( = a²) x p(z 24. s s 8 s/) a') =my(=/_1)) )p (p(ss(s(__1)1) - m)
The forward pass estimates the posterior context distribution at time give the past observations and actions (i.e., for l S t), which is similar to filtering process. The backward pass estimates the context distribution at time give the future observations and actions (i.e., for k 2 t). Combining both passes allows compute the context distribution at time given the whole trajectory.
Gradient of ELBO
VE ELBO
=
Ve [IL (1()))((()))))) - Vo (C(()))
i ELBO
N
V Bo(()
[ogg(ss(a', 11. 1.,0] (
B
(A7)
V ELBO N
B
Vo loloppss'aa' la, n. 0. + V logr(0).
where the terms in blue involve differentiating expectation over Beta distributions, which we compute by adopting the implicit reparameterization (Figurnov et al. 2018 Jankowiak & Obermeyer 2018). Considering a general case where Po(x) and the cumulative distribution function (CDF) of (T) is (2)), it has been shown that:
VoFo(z) (x
V (w) (o) = Epsls)
(Vof(()) (Vstat)
+ Valo(z) Va
C.3 JUSTIFICATION FOR THE VARIATIONAL DISTRIBUTIONS
Here, we provide both intuitive and empirical justifications for the choice of variational distributions in (A6).
The mean-field approximation is mainly based on the tractability consideration where a reparame- terizable variational distribution is required for gradient estimation (Blei et al.. 2017). Besides, the truncation level is set to K, which reduces an infinite dimensional problem to finite one.
The intuition of choosing the point estimation for q(v) is the following: The ((4)) in (A6) induces
variational distribution ((p), following Pyk 1.ik IIk_1 (1 p.s) in A5.). When observing reasonable amount of trajectories, the optimal 9- (p) should center around the ground ground-truth truth initial context distribution and the context transition. The HDP prior in our generative model specifies that Pjle Pyla, 3 ~ DP (a, (3) (Teh et al., 2006), which means B3 serves as the expectation of the initial context distribution and each
in the context transition. Intuitively, the optimal q (.)), which is induced from 9º (v), should center around the stationary distribution of the context chain. Therefore, each factor q" (Px) ) in the optimal q" (v/ supposed to be uni-modal, and thus point estimation could be reasonable simplification. This intuition was supported by our computing 3 and stationary distributions during training resulting in similar distilled Markov chains.
We then conduct simple numerical study to verify the intuition above, which reveals another problem when using full Beta variational distribution for u. Consider the following context Markov chain:
(0.6
p P
0.7 0.3)
Po
0.4
R
0.2
0.8
We set hyper-parameters in A55)
1, K
0, K = 2 and assume that, with sampled
trajectories D. the learned optimal q* (4) is given by:
(pm) )a (mr)q" (μa)) = Beta (oo1 13.8, 9.2) Beta (etaammm144000660 14.00, Beta (3.0, 12.0)-
al2

---

Published as a conference paper at ICLR 2022
14
0.08
0.05
10
Irobability
0.05
0.04
0.03
0.02
0.01
0.0
0.2
0.4
0.6
0.8
1.0
10
15
20
25
6
12-norm
the gradient
(a) Point estimation loss
(b) Point estimation gradient
0.30
0.25
Probabiitty
0.05
0.00
20
P1
12-norm
gradient
(c) Beta variational distribution loss
(d) Beta variational distribution gradient
Figure A4: Loss function of (v)
A8
and the empirical distribution of gradient under two configurations
Recall that Ps1 Myi in AS The assumed 4" (H4r) has the mean at its ground-truth value and the variance 0.01. According to A7), maximizing ELBO w.r.t g(v) is equivalent to:
min [KC (#) [|(()))) + Ve ICC (c(()))((()) ((2)
(A8)
where q(4) is fixed the assumed optima q" (m).
We investigate two configurations of q(v):: (1) a point estimation where g(vl)) 8(11 [1): (2) full Beta variational distribution where q( () Beta( /1 f11, (12). The results shown in Figure A4 The optimal point estimation is q" (v) 8( (0.493) (labeled by the red star in Figure A4(a), while the optimal Beta distribution in Figure A4(c) is q (v) Beta(V) |4.86, 4.78). The Beta optimum is uni-modal and its mode 0.505 very close to the point estimation 0.493, which is consistent with intuition. Comparing Figure A4(b) with A4(d) observe that the point estimation can provide gradients better suited for optimization while using Beta variational distribution potentially suffers from vanishing gradients. We observed this phenomenon our model learning experiments well.
Since the transition function p( a:-1, Oz,) is modeled by neural networks, it is generally hard predict any property of the true posterior of 0 and choose an appropriate variational distribution. Bayesian neural network literature attempt to tackle this problem (Welling & Teh, 2011 Blundell al., 2015 Kingma et al. 20155) Gal & Ghahramani] 2016 Ritter et al. 2018). However, most methods have considerably high computational complexity and it is not trivial to evaluate the quality of generated probabilistic prediction.
this work, we explicitly assume the distribution of 8, whose parameters are fitted by neural networks. The transition model is still capable of generating probabilistic prediction with point estimation of 0. This many model-based RL works.
is followed by
C.4 CONTEXT DISTILLATION
At every iteration of model learning, we can extract MAP parameter estimates {0 and approx- imated posteriors of Po and R (PI...., pK), which are induced from q(Hz). Let also define
al3

---

Published as a conference paper at ICLR 2022
Algorithm A2: Context distillation
Inputs:
Edistil distillation threshold; R expected context transition matrix; 3 weights of HDP's base distribution
Determine the distillation vector e using one of the following choices:
(a) stationary distribution of the chain, v such that v' rR;
(b) weights of HDP"s base distribution, v 3.
Determine distilled context indexes In and spurious context indexes In as follows:
I { Z Edistil J. I2 {ile; Edistil
Compute R as follows:
AGENT is MPC then
R R +
RI,,IA) R.I.I..
end
else if AGENT is POLICY then
R R-
+ RI -I RI 0) R-
(I T-RI,Ia)-RRR RI
end
return R distillee probability transition matrix.
the expected context initial distribution po - Eg(gn) [Po] and the expected context transition matrix R (x(n) RR. These MAP estimates are used during training as well as testing for sampling the values Hence distilling R has an effect on training as well as testing.
Our distillation criterion is based on the values of the stationary distribution of the context Markov chain. Recall that one can compute the stationary distribution 8 by solving pa = p R. Now the meaningful context indexes I (ilp, 2 Edistit and spurious context indexes In {ilp < Edistil) can be chosen using a distillation threshold Edistil- Then, we distill the learned contexts by simply discarding Ôrg- Meanwhile, the context Markov chain also needs to be reduced. For Po. we gather those dimensions indexed by I₁ into a new vector in and re-normalize po. In addition, R can be reduced R following the Theorem in the main text.
Perhaps, a less rigorous, but definitely simpler approach is choosing the index sets I₁ and I2 using @ MAP estimation of 3 computed using D. Since optimizing the KLL (q(μ) [(p(([] term in ELBO is essentially driving the posterior of Pu:k toward B. Therefore, & can be seen 'summary' distribution contexts and
consider the k-th context redundancy if _ is small. not clear has a direct relation the stationary distribution of the Markov chain with the transition probability R. However, we have observed that the magnitudes of the entries of and p are correlated. Hence, in order to avoid computing an eigenvalue decomposition at every context estimation one employ distillation using B.
Both approaches are summarized in Algorithm A2 For policy optimization, we actually need to keep the number of contexts constant dealing with changing state-belief space can be challenging during training. Therefore, the transition matrix R has the same dimensions as R, where the transition probabilities between spurious contexts and from meaningful to spurious context are set to zero. We can still remove the spurious contexts after training both from the model and the policy.
al4

---

Published as a conference paper at ICLR 2022
C.5 MPC USING CROSS-ENTROPY METHOD
This procedure is gradient-free, which benefits lower- dimensional settings, but could be less efficient in higher dimensional environments. It is also known to be more ef- ficient than random shooting methods (Chua al. 2018). The idea of this approach is quite simple and sketched in Algorithm A3 At the state 8, with an action plan {axtk=r we sample plan updates {6% {6k) We then roll-out trajectories and compute the average returns for f--I
E++1
21+h
S1+
51+1
each plan {ax + / ki= We pick Nelite best performing
at+hh
1 compute the empiri- action plans {fax 6 Tk=t fj=1"
cal mean f4 and variance Eg of the elite plan updates 8%
Figure A5: Graphical model for the MPC problem
and then compute the updates on the action plan distribu- tion follows:
F4 Σ == (1 - (1-1,)D 1, + 1, L
1, + 1.
(A9)
where 1, is the learning rate.
Algorithm A3: MPC based on CEM
Inputs: {ax ) { 1++ [E&)eer 1++H Nepoches Nelite. Niraces Npops lr 8, H
for j = 0, ......... N epoochs do
Sample action plan updates {(6) 14+ Npors where of ~ NN Ex); Ji=1
Roll-out Ntraces for each update plan;
Compute the returns 1/Neraces Eperrco Dket r(s ax + 8%)) with sp for all p:
Pick Nelite best performing action plans;
Update the sampling distributions {M*
1+ {Ex 1 in A9
end
C.6 SOFT-ACTOR CRITIC
We reproduce the summary of the soft-actor critic algorithm by Achiam (2018). which we found very accessible. The soft-actor critic algorithm aims at solving modified RL problem with an entropy-regularized objective:
π argmax ET
ai) + + (H((1())))
where H(P) Er [ogg( P(x) and is called the temperature parameter. The entropy regular- ization modifies the Bellman equation for this problem as follows:
Q" (s, a) (' [r(s, a) + (1"" a') - (a'ls')))
The algorithm largely follows the standard actor-critic framework for updating value functions and policy, with a few notable changes. First, two Q functions are used in order to avoid overestimation of the value functions. In particular, the loss for value learning as follows:
Lvalue, ( D) (ss.arr,'' .dy-D ,d)-
(Qs, (s, a) v(r, v(r.s',d)) d))²
(A10)
min Q (s' a') a log (ppa'ls')) /=1,2
y(r d) + r(1 - d)
(A11)
For policy updates the reparameterization trick is used allowing for differentiation of the policy. Namely, the policy loss function is as follows:
4palicy (1b, D) = - - E...D.-Niio..
min Qo, (s',ār) a log ग (āplss),
(A12)
j=1.2
a = tanh (ng + o e E), E &~ S~000
(A13)
a15

---

Published as a conference paper at ICLR 2022
Algorithm A4: Soft-actor critic (hasic version)
Inputs: Nepocha number epochs. Nupd number of gradient updates per epochs, Ntarget-foi freq target value function update frequency, Nsamples number of steps per epoch, 4, w learning rates
Notal- = 0.
Initialize parameters 4/, Φi,, Prarget,i = Φ;
for j ...... Nepochs do
Sample Naamples steps from the environment with a ~ (-()))) resulting in the buffer update
Dnew {(s,, ag 8/ Ti, d.) (d))limmp
Set = Dnew U JD;
Sample batch B from the buffer D; for k = O..... Nupa do
Npo f Neotal-upd + 1;
Update parameters of the value functions Φ; t- Φi - 4, 0; Lvalue, (fp B); Update parameters of the policy vb +b 1, V Lpolicy (1/, B); mod (Neosal-upd Niarget-fred = then
Update parameters of the target value function Otarget, + w uargeet, (1 w)o;;
end
end
end
return Te
D EXPERIMENT DETAILS
D.I LEARNING ALGORITIMMS
Model learning We implemented the model learning using the package Pyro (Bingham et. 2018 which is designed for efficient probabilistic programming. Pyro allows for automatic differentiation, iee. we do not need to explicitly implement message passing and reparametrized gradients for the ELBO gradient computation. We still need, however, forward message pass to compute the belief' estimate, e.g.. to perform filtering the variable when needed.
PPO with an RNN model We modified an implementation PPO by Kostrikov (2018) to account for our belief model. In our implementation, the RNN with hidden state h at time & taking the inputs h,-1, $t-1, a;-1, while producing the output hT- What left is to project the hidden state onto the belief space using a decoder, which we have chosen as b, = softmax(W7 Wh,), where the length of the vector b, is equal the number of contexts. The architecture depicted in Figure A6(a) Note that one can see the RNN and the decoder architecture as a model for the sufficient information state for the POMDP. We have experimented with different architectures. e.g., projecting larger belief' space to account for spurious contexts, removing the decoder altogether etc. These architectures, however, did not yield reasonable results.
GPMM. We took the implementation by Xu et al. (2020), which is able to swing-up the pole attached the and adapt to environments with different parameters.
SAC. We based our implementation largely on (Tandon 2018 with some inspiration from Yaratt Kostrikov| (020).. We use two architectures: full information policy (see Figure A6(c) and model-based policy (see Figure A6(b)) The full information policy is using one hot encoded true context and is, therefore, used as a reference for the best case performance only.
CEM-MPC. We implemented the algorithm from scratch in PoTch.
D.2 ENVIRONMENTS AND THEIR MODELS
Cart-Pole Swing Up We largely followed the setting introduced by Xu et al. (2020), that we set the maximum force magnitude to 20, time interval to 0.04, and time horizon to 100. We took
a16

---

Published as a conference paper at ICLR 2022
ht
critic
Vt
St
bi
RNN
Decoder
at
actor
at
(a) RNN policy
b?
critic
V8
critic
Ve
St
bi
St.
b
Model
at-1
actor
at
actor
at
(b) Model-based policy
(c) Fl policy
Figure A6: Policy architectures
p
dt
P
dt
1
8/
+
dt
MLP
(a) Cart-Pole
E'
dt
R
E
1
dt
P
sin
dt
V
COS
W
MLP
E7
①
dt
MLP
E
()² Id
a
(b) Vehicle
(c) Drone
Figure A7: Structure of the neural networks predicting the mean of the transition probability. The blocks R and dt stand for multiplication with the rotation matrix R and discretization time dt. MLP denotes multi-layer perceptron. E stands for Eurler angles (pitch, roll, yaw), p. v, and w stands for position, velocity and angular velocity in the world frame. Operator stands for the next time step
the implementation by Lovatto 2019) and modified it to fit our context MDP setting. The states of the environment are the position of the mass (p), velocity of the mass (v), deviation angle of the pole from the top position (0) and angular velocity (Ò). For GPMM we replaced 0 with sin and (0) as was done by Xu et al. (2020). We set the reward function to cos(o), where is the deviation from the top position. Our transition model predicts the mean change between the next and the current states and its variance as it is common in model-based RI. Therefore, the structure of our neural network model for mean prediction state and MLP is multi-layer perceptron predicting s' The variance in the transition model is trained parameter. The structure the neural network predicting the mean of the transition model depicted in Figure (7(a)
= s ML, a), where s' is the successor
Drone Take-off The drone environment (Panerati et al. 2021 has 12 states: position in the world frame (pz, Py and pz), yaw, pitch, roll angles (t), and Φ), velocities in the world frame (Uz, Uy and v.) and angular velocities in the world frame (wz, Wy and w.). The prediction of the transition model is similar to the cart-pole model with one notable exception: the neural network for the mean prediction has additional structure. Note that estimate spacial positions, roll, pitch and yaw
al7

---

Published as a conference paper at ICLR 2022
angles given position and angular velocities using crude but effective formulae:
(
sin
t@)
cos tan(f) sin (p)
Apz Apy
Ux Vy
A ② AA A vb
cos (b)
Wx
dt
dt
sin up
cos R
wy
Ap2/
U2
0
cos(O)
cos (0)
where the formula for angular velocities can be found, for example, in Hover & Triantafyllou (200)) We will refer to the matrix R as the rotation matrix with slight abuse of notation. We also choose special features for the MLP: angular velocities, velocities, sines and cosines of the Euler angles (F), actions and actions squared. Using these expression we impose the structure on the neural network depicted in Figure A7(c)
Table Al: Haper-parameters for model learning.
Cart-Pole Swing-Up
Intersection
Drone Take-Off
K
[4, 5, 6,8, 10,
10
10 1
2
2
Model Prior
1 10a
1- 10
5 10
R
3 K/5
6
stdg
0.1
transition cool-off
Network dimensions
128, 4) ReLU
{6, 64,4] ReLU
(20, 128, 6) ReLU
Activations
Optimizer
Clipped Adam
Clipped Adam
Clipped Adam 10 10
Learning rates (0,
(5 10-10-3 10-2, 20 (5 10
20 10
(5 10
Table A2: Hyper-parameter for SAC experiments.
Cart-Pole Swing-Up
Intersection
Drone Take-Off
roll-outs at warm-start
100
200
200
roll-outs per iteration
1
# model iterations at startstart
500
500
500
Runner
# model iterations per epoch
500
200
200
# agent updates at warm start
1000
1000
100
# agent updates per epoch
200
100
150
# epochs
500
500
500
Model frequency update
100
100
80
Model batch size
20
50
50
Agent batch size
256
256
256
Prior
Training distillation threshold
0.1
0.05
0.02
Testing distillation threshold
0.05
0.02
Policy network dimensions
{4, 25, 2} ReLU
{12, 256, 256, 4) ReLU
{12, 256, 4}
Policy networks activations
ReLU
SAC
Value network layer dims
{4, 256 1} ReLU
{12, 256, 256, 1}
{12, 256, 1}
Value networks activations
ReLU -0.01
ReLU
Target entropy
-0.05
-0.1
Initial temperature
0.8
0.2
0.6
Discount factor
0.99
0.999
0.999
Target value fn update freq
4
4
4
Optimizer
Adam
Adam
Adam
Optimization
Policy learning rate
7 10
10
3-10-4
Value function learning rate
3- 10-00-4
10-4
10-4
3--0
Temperature learning rate
5- 10 or 7 10-5
1 10-4
1 1-
Linear Learning Decay
True
True 10-6
True
Weight Decay
10
10
Left turn on the Intersection in Highway Environment We take the environment by Leurent 22018), but use the modifications made by Xu et al. (2020) including the overall reward function structure. We, however, do not penalize the collisions. and increase the episode time from 40 to
a18

---

Published as a conference paper at ICLR 2022
Table A3:
for MPC experiments.
Cart-Pole Swing-Up
Intersection
Drone Take-Off
# roll-outs at warm-start
100
200
200
# roll-outs per iteration
20
20
20
Runner
# model iterations warm-start
500
500
500
model iterations per epoch epochs
500
50
60
10
3
Model batch size
100
100
50
Prior
Training distillation threshold
0.1
0
0
Testing distillation threshold
0.02
0.02
(a) HDP
(b) Dirichlet
(c) MLE
(d) HDP w distillation
Figure A8: Transition matrices, initial p(zo) and stationary p(= distributions the learned context
models in the Cart-Pole Sw Experiment for Result A. Zo Z4 stand for the learned contexts. Reproduction of Figure from the main text.
100 time steps. We again predicted the difference between current and next steps for the mean, and used the simplified model for the position, i.e., Apr Rs div., Apy Rs diiy for both social and ego vehicles.
D.3 HYPER-PARAMETER
For model learning experiments used 500 trajectory roll-outs and 500 epochs for optiiiattion In the cart-pole environment we used the higher learning rate for hard failure experiments when X < and used the lower learning rate for the soft failure experiments >> O. We use the weight decay to avoid gradient explosion in the value functions and the policies. Similarly, Clipped Adam optimizer (available in Pyro) was used avoid gradient explosion in model learning.
All the hyper-parameters are presented in Tables AT A2 and A3
E
ADDITIONAL EXPERIMENTS
E.
HDP IS AN EFFECTIVE PRIOR FOR LEARNING AN ACCURATE AND INTERPRETABLE MODEL.
We plot the time courses of the context evolution and the ground truth context evolution in Figure A9 As the results in Figure A8 (reproduction of Figure 2 in the main text) suggested the MLE method did not provide an accurate context model, while both DP and HDP priors provided models for reconstructing the true context cardinality after distillation. The difference was the choice of the distillation threshold, which had to be significantly higher for the DP prior. This experiment indicates that DP prior can be good tool for modeling context transitions, but HDP provides sharper model fit and a more interpretable model.
For completeness, we performed the same experiment, but with different seed and setting Edisril 0.1 during training. We plot the results in Figures A10 and AII In this case, all models (including the MLE method) coupled with distillation provided accurate estimate of the context evolution. This suggests that the optimization profile for the MLE method has many local minima (peaks and troughs in ELBO), which we can be trapped in given unlucky seed.
a19

---

Published as a conference paper at ICLR 2022
22
2
23
20
20
00
Thme
Time
(a) HDP
(b) Dirichlet
(c) MLE
Figure A9: Time courses the learned context models in Cart-Pole Swing-Up Experiment. "Unlucky' random seed for MLE was used. CO and C1 stand for the ground true contexts, while 20 - Z4 are the learned contexts. Reproduction of Figure from the main text.
30
22
30
22
NO
P24
40.0
(a) HDP
(b) Dirichlet
(c) MLE
Figure A10: Transition matrices, initial p(0)) and stationary p(2x) distributions of the learned context models in Palt-Pole Swing-Up Experiment. Distillation during training and a "lucky" seed were used. Zo Z4 are the learned contexts.
21
24
21
22
2
C1
CL
20
60
20
40
20
40
50
Time
Time 1b
Time 30
(a) HDP
(b) Dirichlet
(c) MLE
Figure A11: Time courses the learned context models in Cart-Pole Swing- Up Experiment. Distillation during training and a "lucky" seed were used. CO and C'1 stand for the ground true contexts, while Z0 24 are the learned contexts.
a20

---

Published as a conference paper at ICLR 2022
51
80
87
86
12
0
(a) No distillation
(b) Distillation threshold 0.01
(c) Distillation threshold 0.1
Figure A12: Influence of distillation during training in Cart-Pole Swing- Up Experiment with [ 8.
Table A4: Comparing the probability mass of the third most probable state in the stationary distribu- tion. We vary the cardinality of the estimated context set C and the distillation threshold ddistil- Red indicates underestimation of the distillation threshold. Reproduction of Table from the main text.
[
4
6
8
10
20
100
Edistil
0
8.58e-03
7.06e-03
3.71e-03
6.85e-03
2.20e-03
2.25e-02
1.37e-01
0.01
1.06e-03 03
1.24e-03
1.37c-03
2.19c-03
2.56e-03
1.60e-02
8.84e-03
0.1
1.21e-03
1.54c-03
1.70c-03
2.80c-03
3.54e-03
9.86e-03
5.03e-02
E.2 DISTILLATION ACTS AS REGULARIZER
After the first experiment, noticed that the context 72 has low probability mass in stationarity, but a high probability of self-transition (see Figure A8 reproduction of Figure in the main text). This suggest that spurious transitions can happen, while highly unlikely. We speculate that the learning algorithm tries to fit the uncertainty in the model (e.g., due to unseen data) to one context. This can lead to over- fitting and unwanted side-effects. Results in Figure A8 (reproduction of Figure in the main text) suggest that distillation during training can act as a regularizer when we used high enough threshold Edistil = 0.1. We further validate our findings by changing the context number upper bound (C] between 4 and 20. In particular, we proceed by presenting the results for varying [C] (taking values 4, 5, 6, 8, 10, 20 and 100) and the distillation threshold Ediatll (taking values 0, 0.01, and 0.1). Note that we distill during training and we refer to the transition matrix for the distilled Markov chain the distilled transition matrix. First, consider the results in Figure A12
where plot the learned MAP estimates of the transition matrices with [č] 8. Note that using distillation with both thresholds prevents overfitting to one context. One could argue that instead of context distillation during training one could simply use distillation after training. This approach, however, can lead to emergence of a spurious context with a large probability raising the possibility of inaccurate model predictions. Furthermore, the large number of spurious contexts can lead to large probability mass concentrated in one of them in stationarity. Indeed, consider Table A4 where we plot the context with third largest probability mass. particular, for |C] 20 the probability mass values for this context are larger than 0.01. This indicates a small but not insignificant possibility of transition to this context, if the distillation does not remove it.
a self-transition
We verify that the distilled transition matrices for various cardinalities |č] are close each other. Let Rx denote the estimated transition matrix for C We use the following metric for comparison
K with only two contexts chosen a posteriori.
Rg Aslli
8x
|Rsl1
That is, we compare all the transition matrices to the case of ( = 5. In Table A5 we present the results, which indicate that the estimated transition matrices are quite close to each other. Furthermore, the distillation during training helps recover the true context regardless of the upper bound (C].
a21

---

Published as a conference paper at ICLR 2022
[ -Rlll1
Table A5: Comparing the estimated transition matrices using the metric ر| the cardinality C and the distillation threshold Edistil-
We vary
||
CC|
4
6
8
10
20
100
Edistil
0
8.26e-03
7.69e-03
2.79e-03
1.29e-02
2.89e-02
1.42e-01
0.01
0.08e-03
6,41e-03
8,88e-03
3.61e-03
2.14e-02
1.77e-02
0.1
2.09e-03
1.66e-03
4.87e-03
1.38e-02
2.19e-02
2.60e-02
To summarize, our experiments suggest that the context set cardinality [ can be confidently overestimated and the context set can be reconstructed using our distillation procedure.
E.3
CONTEXT CARDINALITY MODEL COMPLEXITY
78
23
79
2
21
75
21
21
25
Thme
(a) Edissill = 0
(b) Edistill 0.05
(c) Ediaci 0.1
Figure A13: Time courses of the learned context models for various distillation thresholds (-0.5), {0.05), (0.1), {0.5). {1} We deduce that the learned context sets: Co ([-1), (0.05, 0.1), (0.5), (1)}, Co.os 1f-1, -0.5), (0.05, 0.1, 0.5), (1)), (1))
where C
Co.1 = (1-1,-005) 0.105,0.1,0333
Next we demonstrate that our algorithm can merge similar context automatically using an ap- propriate distillation threshold during training. We increased the force magnitude by two-fold in order to create large number of distinct contexts and chosen the context set C {[-1), [-0.5), (0.5), {0.05), (0.1), (0.5), {1}} and vary the distillation threshold édistill setting it to O, 0.05 and 0.1 during training. In order to get the estimated context sets distilled one more time after training with edietill = 0.02. We set sufficiently high context cardinality estimate K = 12. Results in Figure A13 suggest that the learned contexts sets for various distillation threshold are as follows:
E (1 ((-1) 1], (-0.5), (0.05, (0.5), (1), Co.o5 = (( ((-1, - (0.05,0.11 (1)),
E..1 (1-1, -0.5), (0.05,..1.0
This indicates the ability to change the model structure using the distillation threshold. Furthermore. the trade-off between continuous and discrete dynamics decided automatically using the distillation threshold.
E.4 BREAKING THE ASSUMPTIONS IN HIDDEN MARKOV MODELS
One of the major assumptions in our framework is the Markovian nature of the switches, which can limit its applicability. We conduct two simple experiments to assess if these assumptions can be broken without detrimental effects. First, we model the context transitions as process, where the next context depends on the previous context rather than the current context. This process is clearly
a22

---

Published as a conference paper at ICLR 2022
non-Markovian. However, results in Figure A14 correct contexts in the context realization. Second, we make the context state dependent. In particular, we have:
indicate that model successfully predicts the
c 0, if pos E [0, 0.1),
0-0.1 0.2), [0.1, 0.2),
[0.2, 0.3),
[-0.3, 0.4),
c= if pos E [0, - 0.1),
[-0.2, - 0.3),
[0.3, 0.4),
where pos is the position of the cart. Again our model handles this case as the results in Figure AIS suggest.
So how our Markovian model handles predictions for non-Markovian models? In both cases, the key comparing context extimations/preeiitionss rather than context models, which are almost surely incorrect. However, the predictive power of our model is still preserved since it relies heavily on the observed state history. This allowed us to correctly estimate the contexts in these experiments.
Z4
100
Z0
85.3
0.1
0.2
14.3
0.2
Z3
Z1
55.1
0.8
1.0
40.8
2.3
Z2
51.7
0.2
(9) Prabbbbiit
60
21
Z3
17.9
0.0
0.2
81.5
0.3
20
40
Z4
53.6
0.9
0.8
39.7
5.0
C1
p(z))
0.2
CO
p(z.)
55.5
0.1
0.2
43.9
0.3
20
40
60
Time (t)
(a) Transition matrix
(b) A realization of the process
Figure A14: Learning non-Markovian transitions.
ZO
Z2
Z3
Z4
100
ZO
79.7
0.1
0.3
0.4
19.5
Z3
2.1
1.6
80
Z2
Z2
44.0
0.8
3.7
2.4
49.0
-60
(9%) Pillяbдодk
21
Z3
44.9
0.8
1.7
5.1
47.4
20
17.6
0.1
0.3
81.7
C1
p(2))
41.5
0.8
1.1
1.6
55.1
p(z.)
47.0
0.1
0.3
0.4
52.2
20
40
60
Time (t)
(a) Transition matrix
(b) A realization of the process
Figure A15: Learning state-dependent transitions.
E.5 LEARNING NEW CONTEXTS
As we mentioned above we can remove the new contexts following the proposed distillation procedure. However, model is flexible enough to add new (unseen) contexts without learning the state
a23

---

Published as a conference paper at ICLR 2022
transition model from scratch. We design the following experimental procedure to demonstrate this ability: (1) train the model on dataset D1 which contains two contexts (CO and C1); (2) reset free parameters in the variational distribution ((Vli) and q (i) while preserving ( (3) re-train the model on dataset D2 which contains the original contexts and two new contexts (C2 and
C3). We present the training results on the sets Dj and D2 in Figure AA66 and A17 respectively.
These results suggest that our model able to learn new contexts (C2 and C3 in Figure A17 while preserving the original contexts (CO and C1 in both Figures A16 and A17
74
23
(96) Aaloeeaa
-00
48.0
-40
20
-20
20
40
60
90
Time It]
(a) Learned context model on 2)
(b) Learned context transitions on Dr
Figure A16: Learning the context on the set D1. CO and C1 stand for the ground true contexts, while ZO-Z are the learned contexts.
+
24
100
24
20
72.3
7.8
8.9
1.0
9.9
23
22
(%) Aaoooooi
21
9.6
8.0
71.8
1.0
8.6
23
7.9
71.2
C1
28.3
20.7
24.4
2.7
23.9
150
175
Time (t)
(a) Context model
(b) Context transition
Expanding the model by learning new contexts on the set D2. CO-C3 stand for ground true contexts, while ZO-Z4 are the learned contexts.
Figure A17:
E.6 COMPARING TO POMDP AND C-MDP METHODS
However, for > O, we again increased the foree magnitude by two-fold. For SAC-based algorithms we present the statistics for 50 episodes and 3 seeds. To avoid unfair comparison to POMDP and continual RL methods, we pick their best performance and explain why these methods are not well-suited for our problem. For GPMM we ten separate experiments with 15 episode each and picked three best runs and the best episode performance. For RNN-PPO we run three experiments, but picked the best learned policy over time and over the runs. Furthermore, that the number contexts is known and can be hard coded the RNN architecture. We present our experimental results in Table A6 While it seems that RNN-PPO learns to swing-up correctly for
We chosen the context set as == = [-1, for x E
the RNN in more favorable setting
assume
a24

---

Published as a conference paper at ICLR 2022
failure
hard
soft =
soft a 0.3
soft 0.5
algo |
FI-SAC
84.50 1.79
76.63 + 8.54
84.75 3.07
86.92 + 1.03
C-SAC
85.38 = 1.64
=
86.76 2.88
88.35 = 1.30
C-CEM
87.63 0.14
25.95
83.15
89.08 1.90
DNNMM
-7.73 17.62
2..11
2
66.35 = 19.67
ANPMM
-3.35 15.64
8.07 + 16.48
32.08 19.54
57.22 + 25.55
GPMM
3.50 + 18.59
+
10.64 + 16.10
49.61 = 19.13
RNN-PPO
-0.17 + 18.06
64.10
74.58 + 20.66
67.01 + 8.52
Table A6: Mean + standard deviation for: our algorithms (C-SAC, C-CEM), continual learning algos (GPMM, DNNMM, ANPMM), POMDP algo (RNN-PPO), and SAC with a known context (FI-SAC). For soft failure experiments, we have increased the maximum applicable force by the factor of two. Reproduction of Table from the main text.
soft failures and GPMM almost learns swing-up for X 0.5, the learned belief models indicate that this not so. In fact, plotting the evolution of the belief for RNN-PPO and the context evolution for GPMM illustrates that the algorithms do not learn the context model (see Figure 4 in the main text). In order to illustrate that it is not Gaussian Processes that cause failure in GPMM, Gaussian Process mixture with Deep Neural Network and Attentive Neural Process (Qin et al. 2019; Kim et al.) (019) mixtures (DNNMM and ANPMM, respectively) using the code from Xu ct al. (020). The results this case are similar to the GPMM case, i... we manage to get reasonable rewards for 0.5, but fail for other cases.
replace
E.7 EXPERIMENTS WITH LARGER NUMBER OF CONTEXTS
We further test our approach by introducing larger number of contexts. We have the following parameter sets
[0, 0.1], Cgravity [9.82, 50], Ceart mass
[0.5, 5], Cmax force [20, 40],
Ciriction
with 16 parameters in total. The contexts for our experiment are as follows:
Co Cfriction 0.1,
Cgravity 9.82,
Ceart mass 0.5,
Cmax force 40,
C Ciriction = 0.1,
Cgravity = 9.82,
Ceart = 0.5,
Cmax force
20,
C2 Cfriction 0.1,
Cgravity 9.82,
Ceart mass 5.
Cmax force 40,
C3
Ciriction = 0.1,
Cgravity = 9.82,
Ceart 5,
Cmax force 20,
C₁
Cfriction 0.1,
Cgravity 50,
Ceart mass 0.5,
Cmax force 40,
C Ciriction = 0.1,
Cgravity = 50,
Ceart 0.5,
Cmax force
20,
C. Ofriction 0.1,
Cgravity 50,
Ceart mass
Cmax force 40
C Ofriction = 0.1,
Cgravity = 50,
Geart 5,
Cmax force
20,
Cs
Ciriction O,
Cgravity 9.82,
Ceart mass 0.5,
Cmax force 40
Cy Ciriction = O,
Cgravity = 9.82,
Ceart 0.5,
Cmax force
20,
C10 Ofriction 0,
Cgravity = 9.82,
Ceart mass
Cmax force
40
Cn Ofriction = O,
Cgravity = 9.82,
Ceart 5,
Cmax force
20,
C12 Cfriction 0.
Cgravity 50.
Ceart mass 0.5.
Cmax force 40
C1a Ciriction = 0,
Cgravity = 50,
Ceart 0.5,
Cmax force
20,
C14: Ctriction 0,
Cgravity 50,
Ceart mass
Cmax force 40,
Cis Ciriction = 0.
Cgravity = 50,
Ceart mass 5,
Cmax 20.
We leam model with K = 20. We further add a structure the contexts transition matrix: form every contexts we can switch only to and from two contexts, i.c, from the context we can switch to/from the context and the context where operations should be understood modulo
a25

---

Published as a conference paper at ICLR 2022
K (i.e., - K - 1, - 0), ie::
0.6
p(2i)2)) 0.2 j mo + 1, K) or = mod(i 1, K);
otherwise
Similarly to our previous settings we have a transition cool-off period of 5 time steps.
As the results in Figures (T8(a) and A18(b)
suggest suggest, identify the following meaningful contexts: Zi 1C1, Ca J. Z4 = (Co. C7, C14, C15J. Z10 {C4, Cs, C12, C13). Z14 {C2, C3, Cio. Cub Z18 = (C,, Ca Our algorithm does not distinguish the difference in the values of the friction parameter. The difference in maximum force is not identified for larger masses and stronger gravity, but it can be identified for cart mass 0.5 and gravity 9.82. The other contexts are characterized by different gravity values and cart masses. The only exception is the ground truth context C4, which is present both Z10 and Z18. This, however, bound happen the ground truth contexts are hard to separate.
60
80
20
40
60
(0)
(a) Context realization
(b) Context realization
Figure A18: Learning a context model with 16 contexts.
a26


Reviews:
[{"review_id": "sFm0fBj5IKA", "reviewer": "Reviewer_w49w", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper considers the important and timely problem of context shift in reinforcement learning. This is a popular area of research at the moment. Recent works such as SODA and DRQ have invigorated the communities interest in domain shift and related problems. The treatment given to the problem here is Bayesian, using HDPs and probabilistic inference to handle the test time context switching. It's nice to see such a treatment, as this was a gap in the current literature. \n\nThe domain shift problem is in this paper cast as a contextual MDP (C-MDP). As best as I can tell, a contextual MDP is just an MDP where the state transitions and reward distribution shift around, possibly during training and possibly just at test time. (small note: why does the reward function lack the C subscript if it depends on context?) This machinery seems to play nicely into the eventual HDP model the author's develop. Although, I'm not entirely convinced this additional competing standard is needed. There is already a lot of interesting work on domain shift in robotics. I didn't feel like the new framework was justified, beyond making it easier to recast the problem as probabilistic inference. This is a minor complaint and just an observation. There are always issues with too many competing standards and all that. \n\nIt seems that the number of different contexts must be finite and a hand designed heuristic is required to prune connections and avoid spurious correlations. Perhaps this is a standard part of the package whenever you're using HDP. In any case, often the core problem in domain shift is that the learned distributions are not wide enough. And I worry that this sort of limit on the number of contexts might be a hinderance for wider adoption of this model (which I would like to see!) Do we know where the upper bound on problem dimensionality is? There are new domains (distractor control suit) where you can set distractor intensity, in essence producing arbitrarily many distinct contexts for the environment. Is this limit on the number of atoms going to be the first thing that breaks down as the number of contexts scales? Or something else? \n\nThe trick used to derive equation (4) is cute. Thank you for calling this to my attention. \n\nThere are more clear descriptions in the appendix, but it was not clear to me from reading section 3 what the step by step algorithm is. I could not reproduce these results from the main body of the paper alone. \n\nCan you discuss a little how RNN-PPO relates to RL2, which in essence just fuses the belief state estimate and PPO training steps? There is a general lack of comparison with and concern for both meta learning and domain shift baselines in this paper. Perhaps the authors feel that contextual MDPs is a sufficiently distance problem and prior art in these areas is therefore irrelevant. If that is the case, I would like to see a little further justification. Domain shift in general is a rich area of research (SODA, DRQ). Is this paper justified in ignoring that prior art? Have I fundamentally misunderstood your problem setting?\n\nThe interpretability experiments in Result A are strong. I really like seeing results like this. \n\nThe choices of Drone and Intesection are a little strange. Are there no Deep Mind Control Suit or similar environments that could work? I haven't seen these environments before in RL, so it is hard for me to evaluate how impressive these results are. A lot of the results are done on cartpole, which is fine for trying to pull apart the method. Although it is a bit concerning because cartpole famously works even when the underlying RL algorithm is broken. I would like to see results on further hard RL environments. In Table 3, the environment is called Highway. At the bottom of page 8 (bold text) it's called Intersection. \n\nI am more than willing to bump my score if some of the questions I asked are answered. It's mostly a few points of clarification. ", "summary_of_the_paper": "This paper considers the problem of MDPs with underlying context variables. As these context variables shift, so too does the environmental MDP an RL agent encounters. In particular, this paper considers domain shift in state transitions and reward functions. A new method which relies on hierarchical dirchlet clustering is introduced. This model, as far as I can tell, uses bayesian probabilistic inference to see what parts of the underlying environment probability distribution are useful across context shifts. \n\nA theorem is provided which draws a relationship between the number of mis-identified context switches and the error in the estimated value function. This quantifies the robustness of the algorithm under sudden context shift. ", "summary_of_the_review": "This paper is an interesting look at domain shift from a bayesian probabilistic modeling perspective. The experiments do a good job of taking apart the design choices, in particular results A and B. I have a few lingering questions about this paper's relationship to prior work in Meta Learning and Domain Shift. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "o6rffcDQzot", "reviewer": "Reviewer_q8Tn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Strengths:\n- This paper studies a pretty difficult problem, and proposes to leverage the Hierarchical Dirichlet Process, which can accurately model changes in the agent’s environment and adapt its behavior accordingly.\n\n- It is overall well written and very thorough. Similarly, the experimental section studies several alternative priors and RL algorithms in non-stationary settings, and probes their learned model in various ways.\n\n- The experiment in Appendix A5.4 shows (in a simple experiment) that the model is flexible enough to accommodate new contexts, which is an important capability for continual/lifelong settings.\n\nWeaknesses:\n- At times, the paper becomes convoluted and overloaded with details. For example, the presentation of the HDP could be condensed. I would also suggest a table of contents for the appendix.\n\n- In the experiments, there seems to be a tension between the selected K (cardinality of context set) and distillation threshold. Specifically, if one overestimates K significantly, then the distillation threshold needs to be larger to compensate for an increased number of spurious contexts. However, the threshold can’t be so large that it prunes out relevant contexts. So, careful selection of these hyperparameters is important.\n\n- The experimental settings are relatively simple as they operate in low-dimensional state-action spaces (the largest being 12-dim states and 4-dim actions). The settings also have a relatively small number of contexts that they switch between. I’m curious how well the HDP prior can model large numbers of modes (for example, 100) which meta-RL algorithms often operate in.\n\n- The contexts in all three experimental domains correspond to multiplying the actions by some constant. To see how scalable the model is, it’d be useful to study contexts with multiple variables that parameterize the transition dynamics in more complex ways.\n\n- Assumptions about the environment have to be made for their model to work well, which could be unrealistic in some setups. How effective is the HDP when the assumptions do not hold? For example, what if the mode transitions are not Markovian and depend on earlier modes, or depend on the states and actions? What if we underestimated the cardinality of the context set?\n\nFinally, there are some details that are unclear to me:\n- In the comparison to RNN-PPO, is there a reconstruction loss on prediction of the next state s_t+1? This supervision should be useful for the RNN model, but it’s not clear if it’s added based on the description in Appendix A4.1. \n\n- How do the model learning and policy learning components interact, e.g., are they trained simultaneously? This detail seems to be missing in Section 3.\n\nOther comments:\n- Some of the titles and axes of Figures 5(b-d) are cut off.\n\n- In Tables 2 and 3, it’d be clearer if the highest number of each column were bolded.\n", "summary_of_the_paper": "This paper proposes to model mode transitions (where each mode corresponds to a different MDP) in reinforcement learning with the Hierarchical Dirichlet Process (HDP) prior in a contextual MDP. It makes some assumptions about the evolving environment, such as how it can only take on a finite number of modes and the mode transitions are Markovian. During model learning, it also prunes out modes that have low transition probabilities into them to avoid modeling spurious contexts.", "summary_of_the_review": "The paper studies a challenging problem setting, and demonstrates the effectiveness of the HDP to model this particular setting. The experimental section presents a lot of varied analysis of their model already, but is missing results that demonstrate its scalability to more complex environments. Since the main argument of this work is that the HDP is best suited here, I think there are aspects (described in the main review) that still need to be evaluated.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "fGikr0nkyY4", "reviewer": "Reviewer_HPaX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Weakness:\n- It would be informative to see some experiments with problems that have already been tackled in the literature  (in the context-dependent RL or PO-MDPs framework), to see how the developed method compares to existing ones on problems they were optimized for.\n- Results for the running time are not presented. It would be interesting to see though; it is not clear how efficiently the proposed algorithm can be implemented.\n- (4) seems to be quite a simplification. Can you provide some intuition on how good/bad it is in some realistic examples?\n- The paper is lacking a concise presentation of the algorithm\n\nStrength:\n- The method is based on an elegant approach with Dirichlet Processes\n- Solves a previously untackled but realistic problem\n- The algorithm outperforms the baselines\n- Clear intuition is provided for most steps\n\n\nFurther questions:\n-----------------\nCould the proposed method work in the more general transition learning setup?", "summary_of_the_paper": "The paper proposes a novel approach for policy learning in a context-dependent reinforcement learning setting (a subset of POMDPs). In particular, the approach is a Bayesian method using a generative model based on a hierarchical Dirichlet process.\n\nThe experiments demonstrate that the proposed method performs well even on tasks where the benchmark algorithms chosen by the authors struggle or even fail.\n", "summary_of_the_review": "The proposed method is based on an elegant approach with Dirichlet Processes, and performs well on tasks where the benchmark algorithms chosen by the authors struggle or even fail. Additional experiments, metrics and explanations would be informative though.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XOM48CiYgTo", "reviewer": "Reviewer_w8qw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Pros:\n- This paper is well-written and the proposed methods are discussed in detail. The literature review is pretty comprehensive and the authors did a good job at classifying the related works. \n- As far as I can tell, the idea is novel although it is built on a line of previous works.\n- The theoretical results are principled and clearly defined. I didn’t check all the details but it seems the proofs are good.\n- The experiments are relatively well done, although I would suggest including more baselines.\n\nCons:\n- More details need to be added to the algorithm part. For example, you could point out the update rule or algorithm rather than just say “Update generative model/policy parameters”. Also, I suggest putting one overall algorithm in the main text. In addition, is the distillation step supposed to be included in algorithm A1? \n- In general, the model is relatively complex. Ideally, a principled Bayesian inference procedure estimated the posterior distribution for all the parameters. In this work, the parameter (e.g. the transition functions) are obtained by point estimation. What will be the error caused by this approximation? \n- Experiments are not sufficient. In terms of the learning curve, only comparing with CEM could not be sufficient. It seems that CEM is even outperforming the proposed method in this case. Adding more natural baselines and the corresponding training curve will be a strong plus. \n\nMinor concerns:\n- It would be better to introduce the full name of the methods (e.g. Soft Actor-Critic) before using their abbreviations.\n- Figure 5 is wrongly cropped.\n\nQuestion:\nIn the related work [Xu et al., 2020], their method can model an infinite number of environments based on the Dirichlet process. In this work, the number of contexts is always truncated by a fixed number, is there any possibility that this work can deal with the potential infinite contexts with HDP?\n", "summary_of_the_paper": "This paper proposed a contextual Markov Decision Process with a Hierarchical Dirichlet Process transition prior (HDP-C-MDP) that aims at solving the problem of non-stationary changing environments. Specifically, the context variable is obtained via a truncated variational inference and a context distillation method is applied to remove the spurious ones. The authors also provide theoretical studies concerning the distillation method and the intuition that performance improvement could be brought by knowing the underlying contexts. The authors compared their proposed method with a few reasonable baselines and demonstrated the effectiveness of their method. ", "summary_of_the_review": "This paper addresses a very interesting problem where there exists unknown nonstationarity in RL problems. The proposed method takes the Hierarchical Dirichlet Process as the transition prior. Overall the idea of using HDP is original. However, this paper still lacks sufficient experimental results that make the proposed method substantially stand out. I will give a 6 but and I am still on the boardline. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "G5F_-y2pGw9", "reviewer": "Reviewer_D986", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "CMDP is a powerful framework capable of modeling a broad range of applications. There is a growing literature on RL in CMPs, but most such works, to my knowledge, assume adversarially chosen contexts or contexts sampled in an i.i.d. fashion. The paper is amongst very few works considering the so-called Markovian contexts, an evolution model which, I believe, is relevant in many applications. Learning in CMDPs is a challenging task and so is the problem studied in this paper.  \n\nThe studied problem, while interesting and important, is also highly relevant for the ICLR community as learning in this setting involves some representation learning. \n\nThe paper is well-written and mostly well-organized, and presents an adequate literature review. I have limited knowledge about learning in the CMDP setting here and the Bayesian approach employed. It appears to me that the authors rely on some elements from existing works while some elements sound novel. Nonetheless, due to challenging nature of the learning problem considered, I consider a fruitful orchestration of all these as a contribution. It is also evident that they solidly compare their model and chosen algorithmic elements to existing ones. \n\nIt is so pity that almost all pseudo-codes are presented in the appendix, and a large part of the main text is occupied by experimental results. Perhaps bringing a high-level pseudo-code of the main algorithm to the main text could improve the current organization. \n\nI have one technical comment: The authors use a context distillation procedure to remove spurious contexts. Doesn’t removal of such contexts affect the Markovian property of the contexts?\n \nMinor: \n\n- p. 2: this is case => this is the case\n\n- p. 3: independent and identically distributed => independently … \n\n", "summary_of_the_paper": "This paper studies reinforcement learning where the environment is mathematically modeled as a contextual Markov decision process (CMDP) with finite action-space and context-space, and continuous state-space. The contexts, which are assumed unobservable, may abruptly evolve according to a Markov chain -- hence the name, discrete Markovian context evolution. \n\nThe main contribution of the paper is to present an algorithm for this setting, whose design follows a Bayesian approach. It also relies on Hierarchical Dirichlet Processes (HDP) as priors for context transitions, and uses a context distillation procedure to remove spurious contexts. The performance of the proposed method is demonstrated through numerical experiments. \n", "summary_of_the_review": "This well-written paper presents a Bayesian learning algorithm for RL in CMDPs, which is an interesting and challenging task, and of high relevance to ICLR.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a Bayesian approach to learning in contextual MDPs where the contexts can dynamically vary during the episode.
The authors did well in their rebuttal and alleviated most of the reviewers' concerns. During the discussion there was an agreement that the paper should be accepted.
Please take all reviewer comments into account when preparing the final version.
Decision: Accept (Poster)