Original: R EINFORCEMENT L EARNING IN P RESENCE OF D ISCRETE M ARKOVIAN C ONTEXT E VOLUTION
Revised: R EINFORCEMENT L EARNING IN P RESENCE OF D ISCRETE M ARKOVIAN C ONTEXT E VOLUTION

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We consider a context-dependent Reinforcement Learning (RL) setting, which is characterized by: a) an unknown ﬁnite number of not directly observable contexts; b) abrupt (discontinuous) context changes occurring during an episode; and c)
Revised: We consider a context-dependent Reinforcement Learning (RL) setting, which is characterized by: a) an unknown ﬁnite number of not directly observable contexts; b) abrupt (discontinuous) context changes occurring during an episode; and c)

Original: Markovian context evolution.
Revised: Markovian context evolution.

Original: We argue that this challenging case is often met in applications and we tackle it using a Bayesian  approach and variational inference.
Revised: We argue that this challenging case is often met in applications and we tackle it using a Bayesian model-based approach and variational inference.

Original: We adapt a sticky Hierarchical Dirichlet Process (HDP) prior for model learning, which is arguably best-suited for  Markov process modeling.
Revised: We adapt a sticky Hierarchical Dirichlet Process (HDP) prior for model learning, which is arguably best-suited for inﬁnite Markov chain modeling.

Original: We then derive a context distillation procedure, which identiﬁes and removes spurious contexts in an unsupervised fashion.
Revised: We then derive a context distillation procedure, which identiﬁes and removes spurious contexts in an unsupervised fashion.

Original: We argue that the combination of these two components allows to infer the number of contexts from data thus dealing with the context cardinality assumption.
Revised: We argue that the combination of these two components allows inferring the number of contexts from data thus dealing with the context cardinality assumption.

Original: We then ﬁnd the representation of the optimal policy enabling efﬁcient policy learning using off-the-shelf RL algorithms.
Revised: We then ﬁnd the representation of the optimal policy enabling efﬁcient policy learning using off-the-shelf RL algorithms.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Our world becomes more automated every day with the development of self-driving cars, robotics, and unmanned factories.
Revised: Our world becomes more automated every day with the development of self-driving cars, robotics, and unmanned factories.

Original: Many of these automation processes rely on solutions to sequential decisionmaking problems.
Revised: Many of these automation processes rely on solutions to sequential decisionmaking problems.

Original: Reinforcement Learning (RL) has recently been shown to be an effective tool for solving such problems achieving notable successes, e.g., solving Atari games (Mnih et al., 2013), defeating the (arguably all-time) best human players in the game of GO (Silver et al., 2016), accelerating robot skill acquisition (Kober et al., 2013).
Revised: Reinforcement Learning (RL) has recently been shown to be an effective tool for solving such problems achieving notable successes, e.g., solving Atari games (Mnih et al., 2013), defeating the (arguably all-time) best human players in the game of GO (Silver et al., 2016), accelerating robot skill acquisition (Kober et al., 2013).

Original: Most of the successful RL algorithms rely on abstracting the sequential nature of the decision-making as Markov Decision Processes (MDPs), which typically assume both stationary transition dynamics and reward functions.
Revised: Most of the successful RL algorithms rely on abstracting the sequential nature of the decision-making as Markov Decision Processes (MDPs), which typically assume both stationary transition dynamics and reward functions.

Original: As classic RL departs a well-behaved laboratory setting, stationarity assumptions can quickly become prohibitive, sometimes leading to catastrophic consequences.
Revised: As classic RL departs a well-behaved laboratory setting, stationarity assumptions can quickly become prohibitive, sometimes leading to catastrophic consequences.

Original: As an illustration, imagine an autonomous agent driving a vehicle with changing weather conditions impacting visibility and tyre grip.
Revised: As an illustration, imagine an autonomous agent driving a vehicle with changing weather conditions impacting visibility and tyre grip.

Original: The agent must identify and quickly adapt to these weather conditions changes in order to avoid losing control of the vehicle.
Revised: The agent must identify and quickly adapt to these weather conditions changes in order to avoid losing control of the vehicle.

Original: Similarly, an unmanned aerial vehicle hovering around a ﬁxed set of coordinates needs to deal with sudden atmospheric condition changes (e.g., wind, humidity etc).
Revised: Similarly, an unmanned aerial vehicle hovering around a ﬁxed set of coordinates needs to deal with sudden atmospheric condition changes (e.g., wind, humidity etc).

Original: Another similar and realistic example is an actuator failure, which changes how the action affects the MDP.
Revised: Another similar and realistic example is an actuator failure, which changes how the action affects the MDP.

Original: Following Menke & Maybeck (1995) we distinguish “soft” (a percentage drop in action efﬁciency) and “hard” (action is not affecting the MDP) failures.
Revised: Following Menke & Maybeck (1995) we distinguish “soft” (a percentage drop in action efﬁciency) and “hard” (action is not affecting the MDP) failures.

Original: The failures can also be dynamic as a “soft” failure in one actuator can overload other actuators introducing a chain of failures.
Revised: The failures can also be dynamic as a “soft” failure in one actuator can overload other actuators introducing a chain of failures.

Original: With every ﬁxedspeciﬁc weather condition or  actuator failuresuch environments can be modeled as MDPs, however, with the changing weather or arising failures the environment becomes non-stationary.
Revised: An environment with a ﬁxed weather condition or an actuator failure can be modeled as an MDP, however, with the changing weather or arising failures the environment becomes non-stationary.

Original: 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.
Revised: 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.

Original: This kind of contextual
Revised: This kind of contextual

Original: Markov Decision Processes (C-MDPs) incorporate a number of different RL settings and RL frameworks (Khetarpal et al., 2020): non-stationary RL, 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 a distribution before the start of the episode (e.g, the weather changes abruptly between the instances the vehicle has been deployed).
Revised: Markov Decision Processes (C-MDPs) incorporate a number of different RL settings and RL frameworks (Khetarpal et al., 2020): non-stationary RL, 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 a distribution before the start of the episode (e.g, the weather changes abruptly between the instances the vehicle has been deployed).

Original: Although a signiﬁcant progress in solving speciﬁc instances of C-MDPs has been made, the setting with a countable number of contexts with Markovian transitions between the contexts has not received sufﬁcient attention in the literature — a gap we are aiming to ﬁll.
Revised: Although a signiﬁcant progress in solving speciﬁc instances of C-MDPs has been made, the setting with a countable number of contexts with Markovian transitions between the contexts has not received sufﬁcient attention in the literature — a gap we are aiming to ﬁll.

Original: 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., 2020); ﬁnite state-action spaces (Choi et al., 2000).
Revised: 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., 2020); ﬁnite state-action spaces (Choi et al., 2000).

Original: To enable sample efﬁcient context adaptation Xu et al.
Revised: To enable sample efﬁcient context adaptation Xu et al.

Original: (2020) and Xie et al.
Revised: (2020) and Xie et al.

Original: (2020) developed model-based reinforcement learning algorithms.
Revised: (2020) developed model-based reinforcement learning algorithms.

Original: Speciﬁcally, Xie et al.
Revised: Speciﬁcally, Xie et al.

Original: (2020) learned a latent space variational auto-encoder model with Markovian evolution in continuous context-space, while Xu et al.
Revised: (2020) learned a latent space variational auto-encoder model with Markovian evolution in continuous context-space, while Xu et al.

Original: (2020) adopted a Gaussian Process model for MDPs and a Dirichlet process (DP) prior to model static non-evolving contexts.
Revised: (2020) adopted a Gaussian Process model for MDPs and a Dirichlet process (DP) prior to detect context changes without modeling the context evolution.

Original: Note the use of DP, which is a conjugate prior for a categorical distribution, and ﬁts perfectly with the static context case.
Revised: Note the DP prior, which is a conjugate prior for a categorical distribution, is not sufﬁcient to model context transitions.

Original: We also propose a model-based RL algorithm, however, we model the context and state transitions using the Hierarchical 
Revised: We also propose a model-based RL algorithm, however, we model the context and state transitions using the Hierarchical Dirichlet Process (HDP)

Original: Dirichlet Process (HDP) (Teh et al., 2006; Fox et al., 2008a) prior and a neural network with outputs parametrizing a Guassian distribution (i.e., its mean and variance), respectively, and refer to the model as HDP-C-MDP.
Revised: (Teh et al., 2006; Fox et al., 2008a) prior and a neural network with outputs parametrizing a Guassian distribution (i.e., its mean and variance), respectively, and refer to the model as HDP-C-MDP.

Original: We chose the HDP prior since it only requires the knowledge of an upper bound on context cardinality, and it is better suited for Markov chain modeling than other priorssuch as DP (Teh et al., 2006).
Revised: 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).

Original: Inspired by Blei et al.
Revised: Inspired by Blei et al.

Original: we derive a model learning algorithm using variational inference, which is amenable to RL applications using off-the-shelf algorithms.
Revised: we derive a model learning algorithm using variational inference, which is amenable to RL applications using off-the-shelf algorithms.

Original: Our algorithm relies on two theoretical results guiding the representation learning: a) we propose a 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).
Revised: Our algorithm relies on two theoretical results: a) we propose a 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).

Original: We derive another theoretical result, which shows performance improvement bounds for the fully observable context case.
Revised: We derive another theoretical result, which shows performance improvement bounds for the fully observable context case.

Original: Equipped with these results, we experimentally demonstrate that we can infer the true context cardinality from data.
Revised: Equipped with these results, we experimentally demonstrate that we can infer the true context cardinality from data.

Original: Further, the context distillation procedure can be used during training as a regularizer.
Revised: Further, the context distillation procedure can be used during training as a regularizer.

Original: Interestingly, it can also be used to merge similar contexts, where the measure of similarity is only implicitly deﬁned through the learning loss.
Revised: Interestingly, it can also be used to merge similar contexts, where the measure of similarity is only implicitly deﬁned through the learning loss.

Original: Thus context merging is completely unsupervised.
Revised: Thus context merging is completely unsupervised.

Original: We then show that our model learning algorithm appears to provide an optimization proﬁle with fewer local maxima and minima than the maximum likelihood approach, which we attribute to the Bayesian nature of our algorithm.
Revised: We then show that our model learning algorithm appears to provide an optimization proﬁle with fewer local optima than the maximum likelihood approach, which we attribute to the Bayesian nature of our algorithm.

Original: Finally, we illustrate RL applications on an autonomous car left turn and an autonomous drone take-off tasks.
Revised: Finally, we illustrate RL applications on an autonomous car left turn and an autonomous drone take-off tasks.

Original: We also demonstrate that state-of-the-art algorithms of different frameworks (such as continual RL and Partially-Observable Markov Decision Processes (POMDPs)) fail to solve C-MDPs in our setting, and we elaborate on potential reasons why this is the case.
Revised: We also demonstrate that state-of-the-art algorithms of different frameworks (such as continual RL and Partially-Observable Markov Decision Processes (POMDPs)) fail to solve C-MDPs in our setting, and we elaborate on potential reasons why this is the case.

Original:  P ROBLEM F ORMULATION AND R ELATED W ORK
Revised:  P ROBLEM F ORMULATION AND R ELATED W ORK

Original: We deﬁne a contextual Markov Decision Process (C-MDP) as a tuple M c = (cid:104)C , S , A , P C , P S , R , γ d (cid:105) , where S is the continuous state space; A is the action space; γ d ∈ [0 , 1] is the discount factor; and C denotes the context set with cardinality |C| .
Revised: We deﬁne a contextual Markov Decision Process (C-MDP) as a tuple M c = (cid:104)C , S , A , P C , P S , R , γ d (cid:105) , where S is the continuous state space; A is the action space; γ d ∈ [0 , 1] is the discount factor; and C denotes the context set with cardinality |C| .

Original: In our setting, the state transition and reward function depend on the context, i.e., P S : C × S ×
Revised: In our setting, the state transition and reward function depend on the context, i.e., P S : C × S ×

Original: A × S → [0 , 1] , R : C × S ×
Revised: A × S → [0 , 1] , R : C × S ×

Original: A → R .
Revised: A → R .

Original: Finally, the context distribution probability P C : T t × C → [0 , 1] is conditioned on T t - the past states, actions and contexts { s 0 , a 0 , c 0 , . .
Revised: Finally, the context distribution probability P C : T t × C → [0 , 1] is conditioned on T t - the past states, actions and contexts { s 0 , a 0 , c 0 , . .

Original: . , a t − 1 , c t − 1 , s t } .
Revised: . , a t − 1 , c t − 1 , s t } .

Original: Our deﬁnition is a generalization of the C-MDP deﬁnition by Hallak et al.
Revised: Our deﬁnition is a generalization of the C-MDP deﬁnition by Hallak et al.

Original: (2015), where the contexts are stationary
Revised: (2015), where the contexts are stationary

Original: , i.e., P C : C → [0 , 1] .
Revised: , i.e., P C : C → [0 , 1] .

Original: We adapt our deﬁnition in order to encompass all the settings presented by Khetarpal et al.
Revised: We adapt our deﬁnition in order to encompass all the settings presented by Khetarpal et al.

Original: (2020), where such
Revised: (2020), where such

Original: C-MDPs were used but not formally deﬁned.
Revised: C-MDPs were used but not formally deﬁned.

Original: Throughout the paper, we will restrict the class of C-MDPs by making the following assumptions: (a)
Revised: Throughout the paper, we will restrict the class of C-MDPs by making the following assumptions: a

Original:  Contexts are unknown and not directly observed (b) Context  cardinality is ﬁnite and we know its upper bound K ; (c) Contexts switches can occur during an episode and they are Markovian .
Revised: ) Contexts are unknown and not directly observed (b) Context set cardinality is ﬁnite and we know its upper bound K ; (c) Contexts switches can occur during an episode and they are Markovian .

Original: In particular, we consider the contexts c k representing the parameters of the state transition function θ k , and the context set C to be a subset of the parameter space Θ .
Revised: In particular, we consider the contexts c k representing the parameters of the state transition function θ k , and the context set C to be a subset of the parameter space Θ .

Original: To deal with uncertainty, we consider a set (cid:101) C such that: a) | (cid:101) C|
Revised: To deal with uncertainty, we consider a set (cid:101) C such that: a) | (cid:101) C|

Original: = K > |C| ; b) all its elements θ k ∈ (cid:101) C are sampled from a distribution H ( λ ) , where λ is a hyper-parameter.
Revised: = K > |C| ; b) all its elements θ k ∈ (cid:101) C are sampled from a distribution H ( λ ) , where λ is a hyper-parameter.

Original: Let z t ∈ [0 , . .
Revised: Let z t ∈ [1 , . .

Original: , K ) be the index variable pointing toward a particular parameter vector θ z t , which leads to:
Revised: , K ) be the index variable pointing toward a particular parameter vector θ z t , which leads to:

Original: [Equation]
Revised: [Equation]

Original: where ρ 0 is the initial context distribution, while R =
Revised: where ρ 0 is the initial context distribution, while R =

Original: [ ρ 1 , ..., ρ | (cid:101) C| ] represents the context transition operator.
Revised: [ ρ 1 , ..., ρ | (cid:101) C| ] represents the context transition operator.

Original: As the reader may notice our model is tailored to the case, where the model parameters change abruptly due to external factors such as weather conditions, cascading actuator failures etc.
Revised: As the reader may notice our model is tailored to the case, where the model parameters change abruptly due to external factors such as weather conditions, cascading actuator failures etc.

Original: The change is formalized by a Markov variable z t , which changes the MDP parameters.
Revised: The change is formalized by a Markov variable z t , which changes the MDP parameters θ z t .

Original: Our approach can also be related to switching systems modeling (cf.
Revised: Our approach can also be related to switching systems modeling (cf.

Original: Fox et al.
Revised: Fox et al.

Original: (2008a);
Revised: (2008a);

Original: Becker-Ehmck et al.
Revised: Becker-Ehmck et al.

Original: (2019); Dong et al.
Revised: Dong et al.

Original: and in this case the context is representing the system’s mode.
Revised: and in this case the context is representing the system’s mode.

Original: While we can draw parallels with these works, we improve the model by using nonlinear dynamics (in comparison to Becker-Ehmck et al.
Revised: While we can draw parallels with these works, we improve the model by using nonlinear dynamics (in comparison to Becker-Ehmck et al.

Original: ; Fox et al.
Revised: ; Fox et al.

Original: (2008a)), by using the HDP prior (BeckerEhmck et al.
Revised: (2008a)), by using the HDP prior (Becker-Ehmck et al.

Original:  Dong et al.
Revised: (2019); Dong et al.

Original: (2020) use maximum likelihood estimators), and ﬁnally, by proposing the distillation procedure and using deep learning (in comparison to Fox et al.
Revised: (2020) use maximum likelihood estimators), and ﬁnally, by proposing the distillation procedure and using deep learning (in comparison to Fox et al.

Original: (2008a)).
Revised: (2008a)).

Original: Also note that typically a switching system aims to represent a complex nonlinear (Markov) model using a collection of simpler (e.g., linear) models, which is different from our case.
Revised: Also note that typically a switching system aims to represent a complex nonlinear (Markov) model using a collection of simpler (e.g., linear) models, which is different from our case.

Original: Other restrictions on the space of C-MDPs lead to different problems and solutions (Khetarpal et al., 2020).
Revised: Other restrictions on the space of C-MDPs lead to different problems and solutions (Khetarpal et al., 2020).

Original: We brieﬂy mention a few notable cases, while relegating a detailed discussion to Appendix A. Assuming P C is deterministic with z t = t puts us in the non-stationary RL setting (cf.
Revised: We brieﬂy mention a few notable cases, while relegating a detailed discussion to Appendix A. Assuming that the context is deterministic with z t = t puts us in the non-stationary RL setting (cf.

Original: Chandak et al.
Revised: Chandak et al.

Original: (2020)), where it is common to assume a slowly or smoothly changing non-stationarity as opposed to our case of possibly abrupt changes.
Revised: (2020)), where it is common to assume a slowly or smoothly changing non-stationarity as opposed to our case of possibly abrupt changes.

Original: Restricting the context to a stationary distribution sampled in speciﬁc time points (e.g., at the start of the episode) can be tackled from the continual RL (cf.
Revised: Restricting the context to a stationary distribution sampled in speciﬁc time points (e.g., at the start of the episode) can be tackled from the continual RL (cf.

Original: Nagabandi et al.
Revised: Nagabandi et al.

Original: and meta-RL perspectives (cf.
Revised: and meta-RL perspectives (cf.

Original: Finn et al.
Revised: Finn et al.

Original: , but both are not designed to handle the
Revised: , but both are not designed to handle the

Original: Markovian context case.
Revised: Markovian context case.

Original: Finally, our C-MDP can be seen as a POMDP.
Revised: Finally, our C-MDP can be seen as a POMDP.

Original: Recall that a POMDP is deﬁned by a tuple M po =
Revised: Recall that a POMDP is deﬁned by a tuple M po =

Original: {X , A , O , P X , P O , R , γ d , p ( x 0 ) } , where X , A , O are the state, action, observation spaces, respectively; P X :
Revised: {X , A , O , P X , P O , R , γ d , p ( x 0 ) } , where X , A , O are the state, action, observation spaces, respectively; P X :

Original: X ×
Revised: X ×

Original: A × X → [0 , 1] is the state transition probability;
Revised: A × X → [0 , 1] is the state transition probability;

Original: P O :
Revised: P O :

Original: X ×
Revised: X ×

Original: A × O → [0 , 1] is the conditional observation probability; and p ( x 0 ) is the initial state distribution.
Revised: A × O → [0 , 1] is the conditional observation probability; and p ( x 0 ) is the initial state distribution.

Original: In our case, x = ( s , z ) and o = s .
Revised: In our case, x = ( s , z ) and o = s .

Original:  R EINFORCEMENT L EARNING FOR M ARKOV P ROCESSES WITH M ARKOVIAN C ONTEXT E VOLUTION
Revised:  R EINFORCEMENT L EARNING FOR M ARKOV P ROCESSES WITH M ARKOVIAN C ONTEXT E VOLUTION

Original: There are three main components in our algorithm: the HDP-C-MDP derivation, the model learning algorithm using probabilistic inference and the control algorithms.
Revised: There are three main components in our algorithm: the HDP-C-MDP derivation, the model learning algorithm using probabilistic inference and the control algorithms.

Original: We ﬁrstly brieﬂy comment on each on these components to give an overview of the results and then explain our main contributions to each.
Revised: We ﬁrstly brieﬂy comment on each on these components to give an overview of the results and then explain our main contributions to each.

Original: The detailed description of all parts of our approach can be found in Appendix.
Revised: The detailed description of all parts of our approach can be found in Appendix.

Original: 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 timeseries modeling and analysis tools reported by Fox et al.
Revised: 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 timeseries modeling and analysis tools reported by Fox et al.

Original: (2008a;b) (see also Appendix C.1).
Revised: (2008a;b) (see also Appendix C.1).

Original: We improve the model by proposing a context spuriosity measure allowing for reconstruction of ground truth contexts.
Revised: We improve the model by proposing a context spuriosity measure allowing for reconstruction of ground truth contexts.

Original: We then derive a model learning algorithm using probabilistic inference.
Revised: We then derive a model learning algorithm using probabilistic inference.

Original: Having a model, we can take off-the-shelf algorithms such as a Model Predictive Control (MPC) approach using Cross-Entropy Minimization (CEM) (cf.
Revised: Having a model, we can take off-the-shelf algorithms such as a Model Predictive Control (MPC) approach using Cross-Entropy Minimization (CEM) (cf.

Original: Chua et al.
Revised: Chua et al.

Original: (2018) and Appendix C.5), or a policygradient approach Soft-actor critic (SAC) (cf.
Revised: (2018) and Appendix C.5), or a policygradient approach Soft-actor critic (SAC) (cf.

Original: Haarnoja et al.
Revised: Haarnoja et al.

Original: (2018) and Appendix C.6), which are both well-suited for model-based reinforcement learning.
Revised: (2018) and Appendix C.6), which are both well-suited for model-based reinforcement learning.

Original: While MPC can be directly applied to our model, for policy-based control we ﬁrst derive the representation of the optimal policy.
Revised: While MPC can be directly applied to our model, for policy-based control we ﬁrst derive the representation of the optimal policy.

Original: Generative model: HDP-C-MDP.
Revised: Generative model: HDP-C-MDP.

Original: Before presenting our probabilistic model, let us develop some necessary tools.
Revised: Before presenting our probabilistic model, let us develop some necessary tools.

Original: A Dirichlet process (DP), denoted as DP ( γ, H ) , is characterized by a concentration parameter γ and a base distribution H ( λ ) deﬁned over the parameter space Θ .
Revised: A Dirichlet process (DP), denoted as DP ( γ, H ) , is characterized by a concentration parameter γ and a base distribution H ( λ ) deﬁned over the parameter space Θ .

Original: A sample G from
Revised: A sample G from

Original: DP ( γ, H ) is a probability distribution satisfying ( G ( A 1 ) , ..., G
Revised: DP ( γ, H ) is a probability distribution satisfying ( G ( A 1 ) , ..., G

Original: ( A r )) ∼ Dir ( γH ( A 1 ) , ..., γH ( A r )) for every ﬁnite measurable partition
Revised: ( A r )) ∼ Dir ( γH ( A 1 ) , ..., γH ( A r )) for every ﬁnite measurable partition

Original: A 1 , ..., A r of Θ , where Dir denotes the Dirichlet distribution.
Revised: A 1 , ..., A r of Θ , where Dir denotes the Dirichlet distribution.

Original: Sampling G is often performed using the stick-breaking process (Sethuraman, 1994) and constructed by randomly mixing atoms independently and identically distributed samples θ k from H :
Revised: Sampling G is often performed using the stick-breaking process (Sethuraman, 1994) and constructed by randomly mixing atoms independently and identically distributed samples θ k from H :

Original: [Equation]
Revised: [Equation]

Original: where δ θ k is the Dirac distribution at θ k .
Revised: where δ θ k is the Dirac distribution at θ k .

Original: We note that the stick-breaking procedure assigns progressively smaller values to β k for large k , thus encouraging a smaller number of meaningful atoms.
Revised: We note that the stick-breaking procedure assigns progressively smaller values to β k for large k , thus encouraging a smaller number of meaningful atoms.

Original: The
Revised: The

Original: Hierarchical Dirichlet Process (HDP) is a group of DPs sharing a base distribution, which itself is a sample from a DP: G ∼ DP ( γ, H ) , G j ∼ DP ( α, G ) for all j = 0 , 1 , 2 , . . .
Revised: Hierarchical Dirichlet Process (HDP) is a group of DPs sharing a base distribution, which itself is a sample from a DP: G ∼ DP ( γ, H ) , G j ∼ DP ( α, G ) for all j = 0 , 1 , 2 , . . .

Original: (Teh et al., 2006).
Revised: (Teh et al., 2006).

Original: The distribution G guarantees that all G j inherit the same set of atoms, i.e., atoms of G , while keeping the beneﬁts of DPs in the distributions G j .
Revised: The distribution G guarantees that all G j inherit the same set of atoms, i.e., atoms of G , while keeping the beneﬁts of DPs in the distributions G j .

Original: It can be shown that G j = (cid:80) ∞ k
Revised: It can be shown that G j = (cid:80) ∞ k

Original: =
Revised: =

Original: ρ jk δ θ k for some ρ jk their sampling can be performed using another stick-breaking process (Teh et al., 2006).
Revised: ρ jk δ θ k for some ρ jk whose sampling can be performed using another stick-breaking process (Teh et al., 2006).

Original: We consider its modiﬁed version introduced by Fox et al.
Revised: We consider its modiﬁed version introduced by Fox et al.

Original: [Equation]
Revised: [Equation]

Original: where k ≥ 1 , j ≥ 0 , ˜ δ jk is the Kronecker delta, the parameter κ ≥ 0 , called the sticky factor, modiﬁes the transition matrix priors encouraging self-transitions.
Revised: where k ≥ 1 , j ≥ 0 , ˜ δ jk is the Kronecker delta, the parameter κ ≥ 0 , called the sticky factor, modiﬁes the transition matrix priors encouraging self-transitions.

Original: The sticky factor serves as another measure of regularization reducing the average number of transitions.
Revised: The sticky factor serves as another measure of regularization reducing the average number of transitions.

Original: In our case, the atoms { θ k } forming the context set (cid:101) C are sampled from H ( λ ) , while ρ jk are the parameters of the Hidden Markov Model: ρis the initial context distribution and ρ j are the rows in the transition matrix R .
Revised: In our case, the atoms { θ k } forming the context set (cid:101) C are sampled from H ( λ ) , while ρ jk are the parameters of the Hidden Markov Model: ρis the initial context distribution and ρ j are the rows in the transition matrix R .

Original: Our probabilistic model is constructed in Equations 1,2,3 and illustrated in Figure 1 as a graphical model.
Revised: Our probabilistic model is constructed in Equations 1,2,3 and illustrated in Figure 1 as a graphical model.

Original: We stress that the HDP in its stick-breaking construction assumes that | (cid:101) C| is inﬁnite and countable.
Revised: We stress that the HDP in its stick-breaking construction assumes that | (cid:101) C| is inﬁnite and countable.

Original: In practice, however, we make an approximation and set | (cid:101) C| = K with a large enough K .
Revised: In practice, however, we incorporate the upper cardinality bound K by using a truncated variational distribution, as explained later.

Original: Context Distillation.
Revised: Context Distillation.

Original: HDP-C-MDP promotes a small number of meaningful contexts and some contexts will almost surely be spurious, i.e., we will transition to these contexts with a very small probability.
Revised: HDP-C-MDP promotes a small number of meaningful contexts and some contexts will almost surely be spurious, i.e., we will transition to these contexts with a very small probability.

Original: While this probability is small we may still need to explicitly remove these spurious contexts.
Revised: While this probability is small we may still need to explicitly remove these spurious contexts.

Original: Here we propose a measure of context spuriosity and derive a distillation procedure removing these spurious contexts.
Revised: Here we propose a measure of context spuriosity and derive a distillation procedure removing these spurious contexts.

Original: As a spuriosity measure we will use the stationary distribution of the chain p ∞ , which is computed by solving p ∞ = p ∞ R .
Revised: As a spuriosity measure we will use the stationary distribution of the chain p ∞ , which is computed by solving p ∞ = p ∞ R .

Original: The distillation is then performed as follows: if in stationarity the probability mass of a context is smaller than a threshold ε distil then transitioning to this context is unlikely and it can be removed.
Revised: The distillation is then performed as follows: if in stationarity the probability mass of a context is smaller than a threshold ε distil then transitioning to this context is unlikely and it can be removed.

Original: 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.
Revised: 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.

Original: Theorem 1 Consider a Markov chain p t = p t − 1 R with a stationary distribution p ∞ and distilled
Revised: Theorem 1 Consider a Markov chain p t = p t − 1 R with a stationary distribution p ∞ and distilled

Original: I 

Revised: I

Original: = { i | p ∞ i ≥ ε distil } and spurious I 

Revised: = { i | p ∞ i ≥ ε distil } and spurious I

Original: = { i | p ∞ i < ε distil } state indexes, respectively.
Revised: = { i | p ∞ i < ε distil } state indexes, respectively.

Original: Then a) the matrix (cid:98) R =
Revised: Then a) the matrix (cid:98) R =

Original: R
Revised: R

Original: I 1 , I 1 + R
Revised: I 1 , I 1 + R

Original: I 1 , I 2 ( I − R
Revised: I 1 , I 2 ( I − R

Original: I 2 , I 2 ) − 1 R
Revised: I 2 , I 2 ) − 1 R

Original: I 2 , I 1 is a valid probability transition matrix; b) the
Revised: I 2 , I 1 is a valid probability transition matrix; b) the

Original: Markov chain (cid:98) p t = (cid:98) p t − 1 (cid:98) R is such that its stationary distribution (cid:98) p ∞ ∝ p ∞I 1 .
Revised: Markov chain (cid:98) p t = (cid:98) p t − 1 (cid:98) R is such that its stationary distribution (cid:98) p ∞ ∝ p ∞I 1 .

Original: Model Learning using Probabilistic Inference.
Revised: Model Learning using Probabilistic Inference.

Original: We aim to ﬁnd a variational distribution q ( ν , µ , θ ) to approximate the true posterior p ( ν , µ , θ |D ) , for a dataset D
Revised: We aim to ﬁnd a variational distribution q ( ν , µ , θ ) to approximate the true posterior p ( ν , µ , θ |D ) , for a dataset D

Original: = { ( s i , a i ) } Ni =1 , where s i
Revised: = { ( s i , a i ) } N i =1 , where s i

Original: = { s it } Tt = − 1 and a i = { a it } Tt = − 1 are the state and action sequences in the i -th trajectory.
Revised: = { s it } Tt = − 1 and a i = { a it } Tt = − 1 are the state and action sequences in the i -th trajectory.

Original: We minimize
Revised: We minimize

Original: KL ( q ( ν , µ , θ ) || p ( ν , µ , θ |D )) , or equivalently, maximize the evidence lower bound (ELBO):
Revised: KL ( q ( ν , µ , θ ) || p ( ν , µ , θ |D )) , or equivalently, maximize the evidence lower bound (ELBO):

Original: [Equation]
Revised: [Equation]

Original: The variational distribution above involves inﬁnite-dimensional random variables ν , µ , θ .
Revised: The variational distribution above involves inﬁnite-dimensional random variables ν , µ , θ .

Original: To reach a tractable solution, we assume | (cid:101) C|
Revised: To reach a tractable solution, we assume | (cid:101) C|

Original: =
Revised: =

Original: K and exploit the standard mean-ﬁeld assumption (Blei et al., 2017) and the truncated variational distribution similarly to Blei et al.
Revised: K and exploit the standard mean-ﬁeld assumption (Blei et al., 2017) and the K -truncated variational distribution similarly to Blei et al.

Original: ; Hughes et al.
Revised: ; Hughes et al.

Original: Bryant & Sudderth (2012) as follows:
Revised: Bryant & Sudderth (2012) as follows:

Original: [Equation]

Original: Beta

Original: [Equation]
Revised: [Equation]

Original: where hatted symbols represent free parameters.
Revised: where hatted symbols represent free parameters.

Original: For θ and ν , we seek a MAP point estimate instead of a full posterior (see Appendix C.3 for a discussion on our design choices).
Revised: For θ and ν , we seek a MAP point estimate instead of a full posterior (see Appendix C.3 for a discussion on our design choices).

Original: Random variables not shown in the truncated variational distribution are conditionally independent of data, and thus can be discarded from the problem.
Revised: Random variables not shown in the truncated variational distribution are conditionally independent of data, and thus can be discarded from the problem.

Original: 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 a forward-backward message passing algorithm, (b) we use implicit reparametrized gradients to differentiate with respect to parameters of variational distributions (Figurnov et al., 2018).
Revised: 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 a forward-backward message passing algorithm, (b) we use implicit reparametrized gradients to differentiate with respect to parameters of variational distributions (Figurnov et al., 2018).

Original: We present the detailed derivations in Appendix C.2.
Revised: We present the detailed derivations in Appendix C.2.

Original: We also can perform context distillation during training as discussed in Appendix C.4.
Revised: We also can perform context distillation during training as discussed in Appendix C.4.

Original: While adding some computational complexity, this procedure acts as a regularization for model learning as we show in our experiments.
Revised: While adding some computational complexity, this procedure acts as a regularization for model learning as we show in our experiments.

Original: Representation of the optimal policy.
Revised: Representation of the optimal policy.

Original: First, we notice that the model in Equation 1 is a POMDP, which we get by setting x t :=
Revised: First, we notice that the model in Equation 1 is a POMDP, which we get by setting x t :=

Original: ( z t , s t ) and o t := s t .
Revised: ( z t , s t ) and o t := s t .

Original: In the POMDP case, we cannot claim that o t + depends only on o t and a t .
Revised: In the POMDP case, we cannot claim that o t + depends only on o t and a t .

Original: Therefore the Bellman dynamic programming principle does not hold for these variables and solving the problem is more involved.
Revised: Therefore the Bellman dynamic programming principle does not hold for these variables and solving the problem is more involved.

Original: In practice, one constructs the belief state b t = p ( x t | I Ct )
Revised: In practice, one constructs the belief state b t = p ( x t | I Ct )

Original: (Astrom, 1965), where I Ct = { b 0 , o ≤ t , a <t } is called the information state and is used to compute the optimal policy.
Revised: (Astrom, 1965), where I Ct = { b 0 , o ≤ t , a <t } is called the information state and is used to compute the optimal policy.

Original: Since the belief state is a distribution, it is generally costly to estimate in continuous observation or state spaces.
Revised: Since the belief state is a distribution, it is generally costly to estimate in continuous observation or state spaces.

Original: In our case, estimating the belief is tractable, since the belief of the state s t is the state itself (as the state s t is observable) and the belief of z t , which we denote as b zt , is a vector of a ﬁxed length at every time step (as z t is discrete).
Revised: In our case, estimating the belief is tractable, since the belief of the state s t is the state itself (as the state s t is observable) and the belief of z t , which we denote as b zt , is a vector of a ﬁxed length at every time step (as z t is discrete and b zt is its ﬁltering distribution).

Original: We have the following result with the proof in Appendix B.2.
Revised: We have the following result with the proof in Appendix B.2.

Original: Theorem 2 a)
Revised: Theorem 2 a)

Original: The belief of z can be computed as p ( z t +1 | I Ct )
Revised: The belief of z can be computed as p ( z t +1 | I Ct )

Original: =
Revised: =

Original: b zt +1 , where ( b zt +1 )
Revised: b zt +1 , where ( b zt +1 )

Original: i ∝ N i = (cid:80) j p ( s t +1 | s t , θ i , a t ) ρ ji ( b zt ) j , where ( b zt )
Revised: i ∝ N i = (cid:80) j p ( s t +1 | s t , θ i , a t ) ρ ji ( b zt ) j , where ( b zt )

Original: i are the entries of b zt ; b) the optimal policy can be computed as π ( s , b z ) =
Revised: i are the entries of b zt ; b) the optimal policy can be computed as π ( s , b z ) =

Original: argmax a Q ( s , b z , a ) , where the value function satisﬁes the dynamic programming principle Q ( s t , b zt , a t )
Revised: argmax a Q ( s , b z , a ) , where the value function satisﬁes the dynamic programming principle Q ( s t , b zt , a t )

Original: =
Revised: =

Original: r ( s t , b zt , a t ) + γ (cid:82) (cid:80) i
Revised: r ( s t , b zt , a t ) + γ (cid:82) (cid:80) i

Original: N i max a t +1 Q ( s t +1 , b zt +1 , a t +1 )
Revised: N i max a t +1 Q ( s t +1 , b zt +1 , a t +1 )

Original: d s t +1 .
Revised: d s t +1 .

Original: Computational framework.
Revised: Computational framework.

Original: Algorithm 1 summarizes our approach and is based on the standard model-based RL frameworks (e.g., Pineda et al.
Revised: Algorithm 1 summarizes our approach and is based on the standard model-based RL frameworks (e.g., Pineda et al.

Original: Effectively, we alternate between model updates and policy updates.
Revised: Effectively, we alternate between model updates and policy updates.

Original: For the policy updates we relabel (recompute) the beliefs for the historical transition data.
Revised: For the policy updates we relabel (recompute) the beliefs for the historical transition data.

Original: As MPC methods compute the sequence of actions based solely on the model such relabling is not required.
Revised: As MPC methods compute the sequence of actions based solely on the model such relabeling is not required.

Original: Input: distill distillation threshold

Original: N warm number of trajectories warm start, N trajnumber newly collected trajectories

Original: Initialize AGENT with RANDOM AGENT , D = ∅ ; for i = 1 , . .

Original: , N epochs do

Original: Sample N traj (

Original: N warm

Original: if

Original: i =

Original: ) trajectories from

Original: the environment

Original: with

Original: AGENT ;

Original: Set D new =

Original: s

Original: i , a i )

Original: } N traj i

Original: =1 , where s i

Original: = { s it

Original: Tt = − 1 and a i = { a it } Tt = − 1 are the

Original: state

Original: and

Original: action

Original: sequences in the i -th trajectory

Original: Set

Original: D

Original: =

Original: D ∪ D new ;

Original: Update generative model parameters by gradient ascent on ELBO in Equation

Original: Perform

Original: context distillation with ε distill

Original: ; if AGENT is POLICY then

Original: Sample trajectories for policy update from D ; Recompute the beliefs using the model for these trajectories; Update policy parameters end endreturn AGENT

Original: Performance gain for observable contexts.
Revised: Performance gain for observable contexts.

Original: It is not surprising that observing the ground truth of the contexts should improve the maximum expected return.
Revised: It is not surprising that observing the ground truth of the contexts should improve the maximum expected return.

Original: In particular, even knowing the ground truth context model we can correctly estimate the context z t +1 only a posteriori , i.e., after observing the next state s t +1 .
Revised: In particular, even knowing the ground truth context model we can correctly estimate the context z t +1 only a posteriori , i.e., after observing the next state s t +1 .

Original: Therefore at every context switch we can mislabel it with a high probability.
Revised: Therefore at every context switch we can mislabel it with a high probability.

Original: This leads to a performance loss, which the following result quantiﬁes using the value functions.
Revised: This leads to a performance loss, which the following result quantiﬁes using the value functions.

Original: We have the following result with the proof in Appendix B.3.
Revised: We have the following result with the proof in Appendix B.3.

Original: Theorem 3 Assume we know the true transition model of the contexts and states and consider two settings: we observe the ground truth z t
Revised: Theorem 3 Assume we know the true transition model of the contexts and states and consider two settings: we observe the ground truth z t

Original: and we estimate it using b zt .
Revised: and we estimate it using b zt .

Original: Assume we computed the optimal model-based policy π ( ·| s t , b zt ) with the return R and the optimal ground-truth policy π gt ( ·| s t , z t +1 ) with the corresponding optimal value functions V gt ( s , z ) and Q gt ( s , z, a ) , then:E z 1 , s 0 V gt ( s 0 , z 1 ) −R ≥ E τ, a gt tm ∼ π gt , a tm ∼ π M (cid:88) m = [Equation]
Revised: Assume we computed the optimal model-based policy π ( ·| s t , b z t ) with the return R and the optimal ground-truth policy π gt ( ·| s t , z t +1 ) with the corresponding optimal value functions V gt ( s , z ) and Q gt ( s , z, a ) , then:

Original: m ( Q

Original: [Equation]

Original: s

Original: [Equation]

Original: m

Original: [Equation]

Original: m

Original: [Equation]

Original: m

Original: [Equation]

Original: −

Original: [Equation]

Original: s

Original: [Equation]

Original: m

Original: [Equation]

Original: a

Original: [Equation]
Revised: [Equation]

Original: m )) , where M is the number of misidentiﬁed context switches in a trajectory τ .
Revised: where M is the number of misidentiﬁed context switches in a trajectory τ .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we demonstrate that the HDP offers an effective prior for model learning, while the distillation procedure reﬁnes the model and can regulate the context set complexity.
Revised: In this section, we demonstrate that the HDP offers an effective prior for model learning, while the distillation procedure reﬁnes the model and can regulate the context set complexity.

Original: We also explain why state-of-the-art methods from continual RL, meta-RL and POMDP literature can fail in our setting.
Revised: We also explain why state-of-the-art methods from continual RL, meta-RL and POMDP literature can fail in our setting.

Original: We ﬁnally show that our algorithm can be adapted to high dimensional environments.
Revised: We ﬁnally show that our algorithm can be adapted to high dimensional environments.

Original: We delegate several experiments to Appendix due to space limitations.
Revised: We delegate several experiments to Appendix due to space limitations.

Original: We show that we can learn additional unseen contexts without relearning the whole model from scratch.
Revised: We show that we can learn additional unseen contexts without relearning the whole model from scratch.

Original: We also illustrate how the context distillation during training can be used to merge contexts in an unsupervised manner thus reducing model complexity.
Revised: We also illustrate how the context distillation during training can be used to merge contexts in an unsupervised manner thus reducing model complexity.

Original: We ﬁnally show that our model can generalize to non-Markovian and state dependent context transitions.
Revised: We ﬁnally show that our model can generalize to non-Markovian and state dependent context transitions.

Original: We choose the switching process to be a chain, however, we enforce a cool-off period, i.e, the chain cannot transition to a new state until the cool-off period has ended.
Revised: We choose the switching process to be a chain, however, we enforce a cool-off period, i.e, the chain cannot transition to a new state until the cool-off period has ended.

Original: This makes the context switching itself a non-stationary MDP.
Revised: This makes the context switching itself a non-stationary MDP.

Original: This is done to avoid switches at every time step, but also to show that our method is not limited to the stationary markov context evolution.
Revised: This is done to avoid switches at every time step, but also to show that our method is not limited to the stationary Markov context evolution.

Original: Control Baselines: (1) SAC algorithm with access to the ground truth context information (one-hotencoded variable z t ) denoted as FI-SAC ; (2) SAC algorithm with no context information denoted as NI-SAC (3) a 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 is denoted as GPMM ; (4) A POMDP approach, where the context set cardinality is known and the belief is estimated using an RNN, while PPO (Schulman et al., 2017; Kostrikov, 2018) is used to update the policy.
Revised: Control Baselines: (1) SAC algorithm with access to the ground truth context information (one-hotencoded variable z t ) denoted as FI-SAC ; (2) SAC algorithm with no context information denoted as NI-SAC (3) a 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 is denoted as GPMM ; (4) A POMDP approach, where the context set cardinality is known and the belief is estimated using an RNN, while PPO (Schulman et al., 2017; Kostrikov, 2018) is used to update the policy.

Original: We denote this approach as RNN-PPO .
Revised: We denote this approach as RNN-PPO .

Original: Modeling Prior Baselines: (1) a model with sticky Dirichlet priors ρ j ∼ Dir ( α k
Revised: Modeling Prior Baselines: (1) a model with sticky Dirichlet priors ρ j ∼ Dir ( α k

Original: = α/K + κ ˜ δ jk ) ; (2) a model which removes all priors and conducts a maximum-likelihood (MLE) learning.
Revised: = α/K + κ ˜ δ jk ) ; (2) a model which removes all priors and conducts a maximum-likelihood (MLE) learning.

Original: All the other relevant experimental details (including hyper-parameters) are provided in Appendix D.
Revised: All the other relevant experimental details (including hyper-parameters) are provided in Appendix D.

Original: Z 0 Z 1 Z 2 Z 3 Z 4 Z
Revised: Z 0 Z 1 Z 2 Z 3 Z 4 Z

Original: Z 1 Z
Revised: Z 1 Z

Original: Z 3 Zp ( z 0 ) p ( z ) 87.7 0.0 0.6 0.0 11.7 48.4 0.1 3.4 0.0 48.
Revised: Z 3 Zp ( z 0 ) p ( z ) 87.7 0.0 0.6 0.0 11.7 48.4 0.1 3.4 0.0 48.

Original: P r o b a b ili t y ( % ) (a) HDP
Revised: P r o b a b ili t y ( % ) (a) HDP

Original: Z 0 Z 1 Z 2 Z 3 Z 4 Z
Revised: Z 0 Z 1 Z 2 Z 3 Z 4 Z

Original: Z 1 Z
Revised: Z 1 Z

Original: Z 3 Zp ( z 0 ) p ( z ) [Table] [Table] [Table]{+86.
+}
Revised: Z 3 Zp ( z 0 ) p ( z ) {+86.
+}

Original: P r o b a b ili t y ( % ) (b) Dirichlet
Revised: P r o b a b ili t y ( % ) (b) Dirichlet

Original: Z 0 Z 1 Z 2 Z 3 Z 4 Z
Revised: Z 0 Z 1 Z 2 Z 3 Z 4 Z

Original: Z 1 Z
Revised: Z 1 Z

Original: Z 3 Zp ( z 0 ) p ( z ) 81.0 0.0 9.1 9.9 0.1 5.9 1.5 58.9 32.4 1.
Revised: Z 3 Zp ( z 0 ) p ( z ) 81.0 0.0 9.1 9.9 0.1 5.9 1.5 58.9 32.4 1.

Original: P r o b a b ili t y ( % ) (c) MLE
Revised: P r o b a b ili t y ( % ) (c) MLE

Original: Z 0 Z 1 Z 2 Z 3 Z 4 Z
Revised: Z 0 Z 1 Z 2 Z 3 Z 4 Z

Original: Z 1 Z
Revised: Z 1 Z

Original: Z 3 Zp ( z 0 ) p ( z )
Revised: Z 3 Zp ( z 0 ) p ( z )

Original: P r o b a b ili t y ( % ) (
Revised: P r o b a b ili t y ( % ) (

Original: [Figure]
Revised: [Figure]

Original: HDP 
Revised: HDP w distillation

Original: [Figure]
Revised: [Figure]

Original: distillation
Revised: C Z Z Z

Revised: Z

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: C

Revised: [Figure]

Revised: C

Revised: [Figure]

Revised: Z

Revised: Z

Original:  Initial testing on Cart-Pole Swing-up Task (Lovatto, 2019).
Revised: Z Z [Figure] Initial testing on Cart-Pole Swing-up Task (Lovatto, 2019).

Original: We attempt to swing up and balance a pole attached to a cart.
Revised: We attempt to swing up and balance a pole attached to a cart.

Original: This environment has four states and one action.
Revised: This environment has four states and one action.

Original: We introduce the contexts by multiplying the action with a constant χ thus modulating the actuation effect.
Revised: We introduce the contexts by multiplying the action with a constant χ thus modulating the actuation effect.

Original: We will allow negative χ modeling catastrophic (or hard) failures, and positive χ modeling soft actuation failures.
Revised: We will allow negative χ modeling catastrophic (or hard) failures, and positive χ modeling soft actuation failures.

Original: Result A: HDP is an effective prior for learning an accurate and interpretable model.
Revised: Result A: HDP is an effective prior for learning an accurate and interpretable model.

Original: In
Revised: In

Original: Figure 2, we plot the expectation of ρ 0 and R extracted from the variational distribution q ( µ ) for HDP, Dirichlet and MLE priors for the Cart-Pole Swing-up Environment with the context set
Revised: Figure 2, we plot the expectation of ρ 0 and R extracted from the variational distribution q ( µ ) for HDP, Dirichlet and MLE priors for the Cart-Pole Swing-up Environment with the context set

Original: C = { 1 , − 1 } and | (cid:101) C|
Revised: C = { 1 , − 1 } and | (cid:101) C|

Original: =
Revised: =

Original: K = 5 .
Revised: K = 5 .

Original: The MLE learning appears to be trapped in a local optimum as the results in Figure 2(c) suggest.
Revised: The MLE learning appears to be trapped in a local optimum as the results in Figure 2(c) suggest.

Original: A similar phenomenon has been reported by Dong et al.[Figure] C C
Revised: A similar phenomenon has been reported by Dong et al.

Original: Z

Original: Z

Original: Z

Original: Z

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: C

Original: C

Original: Z

Original: Z

Original: Z

Original: Z

Original: [Figure] Z [Figure] where an MLE method was used and a heuristic entropy regularization and temperature annealing method is adopted to alleviate the issue.
Revised: (2020), where an MLE method was used and a heuristic entropy regularization and temperature annealing method is adopted to alleviate the issue.

Original: All in all, while MLE learning can appear to be competitive with a different random seed, this approach does not give consistent results.
Revised: All in all, while MLE learning can appear to be competitive with a different random seed, this approach does not give consistent results.

Original: The use of Dirichlet priors appears to provide a better model.
Revised: The use of Dirichlet priors appears to provide a better model.

Original: Furthermore, with an appropriate distillation threshold the distilled transition matrices with HDP and Dirichlet priors are very similar to each other.
Revised: Furthermore, with an appropriate distillation threshold the distilled transition matrices with HDP and Dirichlet priors are very similar to each other.

Original: However, the threshold for Dirichlet prior distillation needs to be much higher as calculations of the stationary distributions suggest.
Revised: However, the threshold for Dirichlet prior distillation needs to be much higher as calculations of the stationary distributions suggest.

Original: This implies that spurious transitions are still quite likely.
Revised: This implies that spurious transitions are still quite likely.

Original: In contrast, the HDP prior helps to successfully identify two main contexts ( Z 0 and Z 2 ) and accurately predict the context evolution (see Figure 3).
Revised: In contrast, the HDP prior helps to successfully identify two main contexts ( Z 0 and Z 2 ) and accurately predict the context evolution (see Figure 3).

Original: Furthermore, the model is more interpretable and the meaningful contexts can often be identiﬁed with a naked eye.
Revised: Furthermore, the model is more interpretable and the meaningful contexts can often be identiﬁed with a naked eye.

Original: Result B: Distillation acts as a regularizer.
Revised: Result B: Distillation acts as a regularizer.

Original: We noticed that the context Z 2 has a low probability mass in stationarity, but a high probability of self-transition (Figure 2(a)).
Revised: We noticed that the context Z 2 has a low probability mass in stationarity, but a high probability of self-transition (Figure 2(a)).

Original: This suggest that spurious transitions can happen, while highly unlikely.
Revised: This suggest that spurious transitions can happen, while highly unlikely.

Original: We speculate that the learning algorithm tries to ﬁt the uncertainty in the model (e.g., due to unseen data) to one context.
Revised: We speculate that the learning algorithm tries to ﬁt the uncertainty in the model (e.g., due to unseen data) to one context.

Original: This can lead to over-ﬁtting and unwanted side-effects.
Revised: This can lead to over-ﬁtting and unwanted side-effects.

Original: Results in Figure 2(d) suggest that distillation during training can act as a regularizer when we used a high enough threshold ε distil = 0 .
Revised: Results in Figure 2(d) suggest that distillation during training can act as a regularizer when we used a high enough threshold ε distil = 0 .

Original: We proceed by varying the context set cardinality | (cid:101) C| (taking values 4 , 5 , 6 , 8 , 10 and 20 ) and the distillation threshold ε distil (taking values 0 , 0 . 01 , and 0 . 1 ).
Revised: We proceed by varying the context set cardinality | (cid:101) C| (taking values 4 , 5 , 6 , 8 , 10 and 20 ) and the distillation threshold ε distil (taking values 0 , 0 . 01 , and 0 . 1 ).

Original: Note that we distill during training and we refer to the transition matrix for the distilled Markov chain as the distilled transition matrix.
Revised: Note that we distill during training and we refer to the transition matrix for the distilled Markov chain as the distilled transition matrix.

Original: As the ground truth context cardinality is equal to two, the probability of the third most likely context would signify the learning error.
Revised: As the ground truth context cardinality is equal to two, the probability of the third most likely context would signify the learning error.

Original: In Table 1, we present the stationary probability of the context with the third largest probability mass.
Revised: In Table 1, we present the stationary probability of the context with the third largest probability mass.

Original: In particular, for | (cid:101) C|
Revised: In particular, for | (cid:101) C|

Original: = 20 the probability mass values for this context are larger than 0 . 01 .
Revised: = 20 the probability mass values for this context are larger than 0 . 01 .

Original: This indicates a small but not insigniﬁcant possibility of a transition to this context, if the distillation does not remove this context.
Revised: This indicates a small but not insigniﬁcant possibility of a transition to this context, if the distillation does not remove this context.

Original: We present some additional details on this experiment in Appendix E.1.
Revised: We present some additional details on this experiment in Appendix E.1.

Original: Overall, we can conclude that it is safe to overestimate the context cardinality.
Revised: Overall, we can conclude that it is safe to overestimate the context cardinality.

Original: [Table]

Original: Result C: MDP, POMDP and continual RL methods can be ineffective.
Revised: Result C: MDP, POMDP and continual RL methods can be ineffective.

Original: In Figure 4(a), we plot the learning curves for our algorithms and compare them to each other for χ = − 1 .
Revised: In Figure 4(a), we plot the learning curves for our algorithms and compare them to each other for χ = − 1 .

Original: CEM, which is known to perform well in low-dimensional environments, learns faster than SAC.
Revised: CEM, which is known to perform well in low-dimensional environments, learns faster than SAC.

Original: Note that there is no signiﬁcant performance loss of C-SAC in comparison with the full information case exhibiting the power of our modeling approach.
Revised: Note that there is no signiﬁcant performance loss of C-SAC in comparison with the full information case exhibiting the power of our modeling approach.

Original: We evaluated FI-SAC, C-SAC, C-CEM on three seeds.
Revised: We evaluated FI-SAC, C-SAC, C-CEM on three seeds.

Original: [Figure]

Original: [Table]

Original: We now compare the control algorithms for various values of χ .
Revised: We now compare the control algorithms for various values of χ .

Original: We present the results of our experiments in Table 2 and we also discuss the comparison protocols in Appendix E.6. 
Revised: We present the results of our experiments in Table 2 and we also discuss the comparison protocols in Appendix E.6. Here we [Table]

Original: Here we focus on the reasons why both RNN-PPO and GPMM can fail in some experiments and seem to perform well in others.
Revised: [Figure] focus on the reasons why both RNN-PPO and GPMM can fail in some experiments and seem to perform well in others.

Original: In GPMM, it is explicitly assumed that the context does not change during the episode, however, the algorithm can adapt to a new context.
Revised: In GPMM [Table] it is explicitly assumed that the context does not change during the episode, however, the algorithm can adapt to a new context.

Original: While a posteriori context estimation has a limited success for χ = 0 .
Revised: While a posteriori context estimation has a limited success for χ = 0 .

Original: 5 (see Figure 4), the context adaptation is rather slow for our setting resulting in many context estimation errors, which reduces the performance.
Revised: 5 (see Figure 4), the context adaptation is rather slow for our setting resulting in many context estimation errors, which reduces the performance.

Original: Furthermore, it appears that estimating hard failures is a challenge for GPMM.
Revised: Furthermore, it appears that estimating hard failures is a challenge for GPMM.

Original: RNN-PPO appears to perform very well for χ > 0 (see Table 2) and fail for χ = − 1 , 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).
Revised: RNN-PPO appears to perform very well for χ > 0 (see Table 2) and fail for χ = − 1 , 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).

Original: It is worth noting that the mean of the true belief variable is 0 . 5 for all χ , as both contexts are equally probable at every time step.
Revised: It is worth noting that the mean of the true belief variable is 0 . 5 for all χ , as both contexts are equally probable at every time step.

Original: Therefore, the RNN approach does not actually learn a belief model, but an “average” adjustment signal for the policy, and hence it will often fail to solve a C-MDP.
Revised: Therefore, the RNN approach does not actually learn a belief model, but an “average” adjustment signal for the policy, and hence it will often fail to solve a C-MDP.

Original: Interestingly, with χ = 0 .
Revised: Interestingly, with χ = 0 .

Original: 5 our modeling algorithm learns only one meaningful context with high distillation threshold while still solving the task.
Revised: 5 our modeling algorithm learns only one meaningful context with high distillation threshold while still solving the task.

Original: This is because for both χ = 0 .
Revised: This is because for both χ = 0 .

Original: 5 and χ = 1 the sign of optimal actions for swing up are the same and both have sufﬁcient power to solve the task.
Revised: 5 and χ = 1 the sign of optimal actions for swing up are the same and both have sufﬁcient power to solve the task.

Original: We compare to further baselines in Appendix E.6.
Revised: We compare to further baselines in Appendix E.6.

Original: Our model is effective for control in twelve dimensional environments (Drone and Intersection) .
Revised: Our model is effective for control in twelve dimensional environments (Drone and Intersection) .

Original: In the drone environment (Panerati et al., 2021), the agent aims at balancing roll and pitch angles of the drone, while accelerating vertically, i.e., the task is to maximize the upward velocity.
Revised: In the drone environment (Panerati et al., 2021), the agent aims at balancing roll and pitch angles of the drone, while accelerating vertically, i.e., the task is to maximize the upward velocity.

Original: This environment has twelve states (positions, velocities, Euler angles, angular velocities in three dimensions) and four actions (motor speeds in rotation per minute).
Revised: This environment has twelve states (positions, velocities, Euler angles, angular velocities in three dimensions) and four actions (motor speeds in rotation per minute).

Original: 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.
Revised: 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.

Original: The goal of the agent is to make the left turn and follow the social vehicle without colliding with it.
Revised: The goal of the agent is to make the left turn and follow the social vehicle without colliding with it.

Original: The agent measures positions, velocities and headings in x , y axes of the ego and social vehicles (twelve states in total), while controlling the steering angle and acceleration / deceleration.
Revised: The agent measures positions, velocities and headings in x , y axes of the ego and social vehicles (twelve states in total), while controlling the steering angle and acceleration / deceleration.

Original: In both environments we introduce the contexts by multiplying the maximum actuation effect by a constant χ , speciﬁcally, motor speeds in the drone[Table] environment and steering angle in the highway intersection environment.
Revised: In both environments we introduce the contexts by multiplying the maximum actuation effect by a constant χ , speciﬁcally, motor speeds in the drone environment and steering angle in the highway intersection environment.

Original: The results in Tabledemonstrate that both MPC and policy learning approaches with the model are able to solve the task, while using no information (NI) about the contexts dramatically reduces the performance.
Revised: The results in Tabledemonstrate that both MPC and policy learning approaches with the model are able to solve the task, while using no information (NI) about the contexts dramatically reduces the performance.

Original: 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.
Revised: 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.

Original: We can only hypothesize that the policy’s feedback architecture (mapping states to actions) is better suited for complex tasks such as low-level vehicle control, where MPC approaches require a substantial tuning and computational effort to compete with a policy based approach. 
Revised: We can only hypothesize that the policy’s feedback architecture (mapping states to actions) is better suited for complex tasks such as low-level vehicle control, where MPC approaches require a substantial tuning and computational effort to compete with a policy based approach. [Table] [Table]

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: We studied a hybrid discrete-continuous variable process, where unobserved discrete variable represents the context and observed continuous variables represents the dynamics state.
Revised: We studied a hybrid discrete-continuous variable process, where unobserved discrete variable represents the context and observed continuous variables represents the dynamics state.

Original: We proposed a variational inference algorithm for model learning using a sticky HDP prior.
Revised: We proposed a variational inference algorithm for model learning using a sticky HDP prior.

Original: This prior allows for effective learning of an interpretable model and coupled with our context distillation procedure offers a powerful tool for learning C-MDPs.
Revised: This prior allows for effective learning of an interpretable model and coupled with our context distillation procedure offers a powerful tool for learning C-MDPs.

Original: In particular, we showed that the combination of the HDP prior and the context distillation method allows learning the true context cardinality.
Revised: In particular, we showed that the combination of the HDP prior and the context distillation method allows learning the true context cardinality.

Original: We also showed that the model quality is not affected if the upper bound on context cardinality set is overestimated.
Revised: We also showed that the model quality is not affected if the upper bound on context cardinality set is overestimated.

Original: Furthermore, we illustrated that the distillation threshold can be used as a regularization trade-off parameter and it can also be used to merge similar contexts in an unsupervised manner.
Revised: Furthermore, we illustrated that the distillation threshold can be used as a regularization trade-off parameter and it can also be used to merge similar contexts in an unsupervised manner.

Original: Furthermore, we present additional experiments in Appendix suggesting that our model can potentially generalize to non-Markovian and state-dependent settings.
Revised: Furthermore, we present additional experiments in Appendix suggesting that our model can potentially generalize to non-Markovian and state-dependent settings.

Original: While we presented several experiments in various environments, further experimental evaluation are required, e.g., usingbenchmarks in Benjamins et al.
Revised: While we presented several experiments in various environments, further experimental evaluation is required, e.g., using Benjamins et al.

Original: We showed that continual and meta-RL approaches are likely to fail as their underlying assumptions on the environment do not ﬁt our setting.
Revised: We showed that continual and meta-RL approaches are likely to fail as their underlying assumptions on the environment do not ﬁt our setting.

Original: The learned models do not appear to capture the complexity of Markovian context transitions.
Revised: The learned models do not appear to capture the complexity of Markovian context transitions.

Original: This, however, should not be surprising as these methods are tailored to a different problem: adapting existing policy / model to a new setting.
Revised: This, however, should not be surprising as these methods are tailored to a different problem: adapting existing policy / model to a new setting.

Original: If the context is very different and / or the contexts changing too fast then the continual and meta-RL algorithms would struggle by design.
Revised: If the context is very different and / or the contexts changing too fast then the continual and meta-RL algorithms would struggle by design.

Original: We derived our policy by exploiting the relation of our setting and POMDPs.
Revised: We derived our policy by exploiting the relation of our setting and POMDPs.

Original: 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.
Revised: 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.

Original: We attribute this behavior to the lack of effective priors and model structure.
Revised: We attribute this behavior to the lack of effective priors and model structure.

Original: While in some cases it can appear that the RNN policy is effective, disregarding the context altogether has a similar effect.
Revised: While in some cases it can appear that the RNN policy is effective, disregarding the context altogether has a similar effect.

Original: Our model-based algorithm can be further enhanced by using synthetic one-step transitions similarly to Janner et al.
Revised: Our model-based algorithm can be further enhanced by using synthetic one-step transitions similarly to Janner et al.

Original: (2019), which would improve sample efﬁciency.
Revised: (2019), which would improve sample efﬁciency.

Original: We can also use an ensemble of models, which would allow to constantly improve the model using cross-validation over models in the ensemble.
Revised: We can also use an ensemble of models, which would allow to constantly improve the model using cross-validation over models in the ensemble.

Original: However, evaluation of the model quality is more involved since the context is unobservable.
Revised: However, evaluation of the model quality is more involved since the context is unobservable.

Original: In future, we also plan to extend our model to account for a partially observable setting, i.e., where the state only indirectly measured similarly to POMDPs.
Revised: In future, we also plan to extend our model to account for a partially observable setting, i.e., where the state only indirectly measured similarly to POMDPs.

Original: This setting would allow for a rigorous treatment of controlling from pictures in the context-dependent setting.
Revised: This setting would allow for a rigorous treatment of controlling from pictures in the context-dependent setting.

Original: While we show that we can learn unseen contexts without re-learning the entire model, this procedure is not fully automated.
Revised: While we show that we can learn unseen contexts without re-learning the entire model, this procedure is not fully automated.

Original: Hence it can beneﬁt from adaptation of continual learning methods in order to increase efﬁciency of the learning procedure.
Revised: Hence it can beneﬁt from adaptation of continual learning methods in order to increase efﬁciency of the learning procedure.

Original: Appendices A Detailed Literature Review
Revised: Appendices A Detailed Literature Review

Original: a
Revised: a

Original: B Proofs
Revised: B Proofs

Original: a
Revised: a

Original: B.
Revised: B

Revised: [Table]

Original: 1 Proof of
Revised: Proof of

Original: Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: [Table]

Original: a
Revised: a

Original: B.2 Derivation of dynamic programming principle and proof of Theorem2 . . . . .
Revised: [Table] 2 Derivation of dynamic programming principle and proof of Theorem

Original: a
Revised: a

Original: B.3 Performance gain for observable contexts
Revised: B [Table] 3 Performance gain for observable contexts

Original: a
Revised: a

Original: C Algorithm Details
Revised: C Algorithm Details

Original: a
Revised: a

Original: C.
Revised: [Table]

Original: 1 Slightly more details on the Hierarchical Dirichlet Processes. . . . . . . . . . .
Revised: 1 Slightly more details on the Hierarchical Dirichlet Processes

Original: a
Revised: a

Original: C.2 Variational Inference for Probabilistic Modeling
Revised: C.2 Variational Inference for Probabilistic Modeling

Original: a
Revised: a

Original: C.3 Justiﬁcation for the variational distributions
Revised: C.3 Justiﬁcation for the variational distributions

Original: a
Revised: a

Original: C.4 Context Distillation
Revised: C.4 Context Distillation

Original: a
Revised: a

Original: C.5 MPC using Cross-Entropy Method
Revised: C.5 MPC using Cross-Entropy Method

Original: a
Revised: a

Original: [Table]
Revised: [Table]

Original: 6 Soft-Actor Critic
Revised: 6 Soft-Actor Critic

Original: a
Revised: a

Original: D Experiment Details
Revised: D Experiment Details

Original: a
Revised: a

Original: D.1 Learning algorithms
Revised: D.1 Learning algorithms

Original: a
Revised: a

Original: D.2 Environments and their Models
Revised: D.2 Environments and their Models

Original: a
Revised: a

Original: D.3 Hyper-parameters 
Revised: D.3 Hyper-parameters .

Original: a
Revised: a

Original: E Additional Experiments
Revised: E Additional Experiments

Original: a
Revised: a

Original: E.1 HDP is an effective prior for learning an accurate and interpretable model. . . .
Revised: E.1 HDP is an effective prior for learning an accurate and interpretable model

Original: a
Revised: a

Original: [Table]
Revised: [Table]

Original: A D ETAILED L ITERATURE R EVIEW
Revised: A D ETAILED L ITERATURE R EVIEW

Original: For convenience, we reproduce the deﬁnitions and assumptions from the main text to make the appendix self-contained.
Revised: For convenience, we reproduce the deﬁnitions and assumptions from the main text to make the appendix self-contained.

Original: We deﬁne a contextual Markov Decision Process (C-MDP) as a tuple
Revised: We deﬁne a contextual Markov Decision Process (C-MDP) as a tuple

Original: M c = (cid:104)C , S , A , P C , P S , R , γ d (cid:105) , where S is the continuous state space; A is the action space; γ d ∈ [0 , 1] is the discount factor; and C denotes the context set with cardinality |C| .
Revised: M c = (cid:104)C , S , A , P C , P S , R , γ d (cid:105) , where S is the continuous state space; A is the action space; γ d ∈ [0 , 1] is the discount factor; and C denotes the context set with cardinality |C| .

Original: In our setting, the state transition and reward function depend on the context, i.e., P S : C × S ×
Revised: In our setting, the state transition and reward function depend on the context, i.e., P S : C × S ×

Original: A × S → [0 , 1] ,
Revised: A × S → [0 , 1] ,

Original: R : C × S ×
Revised: R : C × S ×

Original: A → R .
Revised: A → R .

Original: Finally, the context distribution probability P C : T t × C → [0 , 1] is conditioned on T t - the past states, actions and contexts { s 0 , a 0 , c 0 , . .
Revised: Finally, the context distribution probability P C : T t × C → [0 , 1] is conditioned on T t - the past states, actions and contexts { s 0 , a 0 , c 0 , . .

Original: . , a t − 1 , c t − 1 , s t } .
Revised: . , a t − 1 , c t − 1 , s t } .

Original: Our deﬁnition is a generalization of the C-MDP deﬁnition by Hallak et al.
Revised: Our deﬁnition is a generalization of the C-MDP deﬁnition by Hallak et al.

Original: (2015), where the contexts are stationary, i.e.,
Revised: (2015), where the contexts are stationary, i.e.,

Original: P C : C → [0 , 1] .
Revised: P C : C → [0 , 1] .

Original: We adapt our deﬁnition in order to encompass all the settings presented by Khetarpal et al.
Revised: We adapt our deﬁnition in order to encompass all the settings presented by Khetarpal et al.

Original: (2020), where such C-MDPs were used but not formally deﬁned.
Revised: (2020), where such C-MDPs were used but not formally deﬁned.

Original: Throughout the paper, we will restrict the class of C-MDPs by making the following assumptions: (a)
Revised: Throughout the paper, we will restrict the class of C-MDPs by making the following assumptions: (a)

Original: Contexts are unknown and not directly observed (b) Context cardinality is ﬁnite and we know its upper bound K ; (c) Context distribution is Markovian .
Revised: Contexts are unknown and not directly observed (b) Context cardinality is ﬁnite and we know its upper bound K ; (c) Context distribution is Markovian .

Original: In particular, we consider the contexts c k representing the parameters of the state transition function θ k , and the context set C to be a subset of the parameter space Θ .
Revised: In particular, we consider the contexts c k representing the parameters of the state transition function θ k , and the context set C to be a subset of the parameter space Θ .

Original: To deal with uncertainty, we consider a set (cid:101) C such that: a) | (cid:101) C|
Revised: To deal with uncertainty, we consider a set (cid:101) C such that: a) | (cid:101) C|

Original: = K > |C| ; b) all its elements θ k ∈ (cid:101) C are sampled from a distribution H ( λ ) , where λ is a hyper-parameter.
Revised: = K > |C| ; b) all its elements θ k ∈ (cid:101) C are sampled from a distribution H ( λ ) , where λ is a hyper-parameter.

Original: Let z t ∈ [0 , . .
Revised: Let z t ∈ [0 , . .

Original: , K ) be the index variable pointing toward a particular parameter vector θ z t .
Revised: , K ) be the index variable pointing toward a particular parameter vector θ z t .

Original: We thus a write the environment model as:
Revised: We thus a write the environment model as:

Original: [Equation]
Revised: [Equation]

Original: For convenience we also write R =
Revised: For convenience we also write R =

Original: [ ρ 1 , ..., ρ | (cid:101) C| ] representing the context transition operator, even if | (cid:101) C| is countable and inﬁnite.
Revised: [ ρ 1 , ..., ρ | (cid:101) C| ] representing the context transition operator, even if | (cid:101) C| is countable and inﬁnite.

Original: Other restrictions on the space of C-MDPs lead to different problems and solutions.
Revised: Other restrictions on the space of C-MDPs lead to different problems and solutions.

Original: We present some settings graphically in Figure A1 adapted from the review by Khetarpal et al.
Revised: We present some settings graphically in Figure A1 adapted from the review by Khetarpal et al.

Original: Setting A.
Revised: Setting A.

Original: In this setting, a Markov assumption is made on the context evolution.
Revised: In this setting, a Markov assumption is made on the context evolution.

Original: Only a few works make such an assumption similarly to our work.
Revised: Only a few works make such an assumption similarly to our work.

Original: For example, Choi et al.
Revised: For example, Choi et al.

Original: (2000) assume a discrete state, action and context spaces and ﬁx the exact number of contexts thus signiﬁcantly limiting applicability of their approach.
Revised: (2000) assume a discrete state, action and context spaces and ﬁx the exact number of contexts thus signiﬁcantly limiting applicability of their approach.

Original: We can also mention the work by Xu et al.
Revised: We can also mention the work by Xu et al.

Original: (2020), where the individual points are classiﬁed to different context during the episode.
Revised: (2020), where the individual points are classiﬁed to different context during the episode.

Original: However, the agent explicitly assumes that the context will not change during the execution of its plan.
Revised: However, the agent explicitly assumes that the context will not change during the execution of its plan.

Original: The Markovian context evolution can also be related to switching systems (cf.
Revised: The Markovian context evolution can also be related to switching systems (cf.

Original: Becker-Ehmck et al.
Revised: Becker-Ehmck et al.

Original: ; Fox et al.
Revised: ; Fox et al.

Original: (2008a); Dong et al.
Revised: (2008a); Dong et al.

Original: (2020)), where the context is representing the system’s mode.
Revised: (2020)), where the context is representing the system’s mode.

Original: While we take inspiration from these works, we improve the model in comparison to Becker-Ehmck et al.
Revised: While we take inspiration from these works, we improve the model in comparison to Becker-Ehmck et al.

Original: Fox et al.
Revised: Fox et al.

Original: (2008a)
Revised: (2008a)

Original: by using non-linear dynamics, in comparison to Fox et al.
Revised: by using non-linear dynamics, in comparison to Fox et al.

Original: (2008a)
Revised: (2008a)

Original: by proposing the distillation procedure and using deep learning, in comparison to Becker-Ehmck et al.
Revised: by proposing the distillation procedure and using deep learning, in comparison to Becker-Ehmck et al.

Original: Dong et al.
Revised: Dong et al.

Original: by using the hierarchical Dirichlet process prior.
Revised: by using the hierarchical Dirichlet process prior.

Original: Setting B.
Revised: Setting B.

Original: Here, an episodic non-stationarity constraint on the context probability distribution P C is introduced, that is, the context may change only (without loss of generality) at the start of the episode.
Revised: Here, an episodic non-stationarity constraint on the context probability distribution P C is introduced, that is, the context may change only (without loss of generality) at the start of the episode.

Original: Episodic non-stationarity constraint is widely adopted by meta-RL (Doshi-Velez & Konidaris, 2016;
Revised: Episodic non-stationarity constraint is widely adopted by meta-RL (Doshi-Velez & Konidaris, 2016;

Original: Finn et al., 2017; Al-Shedivat et al., 2018; Rothfuss et al., 2019; Clavera et al., 2018; Rakelly et al., 2019; Zintgraf et al., 2020; Xie et al., 2020) and continual RL (Chandak et al., 2020; Nagabandi et al., 2018; Khetarpal et al., 2020; Bou-Ammar et al., 2014; Riemer et al., 2019; Rolnick et al., 2019) communities.
Revised: Finn et al., 2017; Al-Shedivat et al., 2018; Rothfuss et al., 2019; Clavera et al., 2018; Rakelly et al., 2019; Zintgraf et al., 2020; Xie et al., 2020) and continual RL (Chandak et al., 2020; Nagabandi et al., 2018; Khetarpal et al., 2020; Bou-Ammar et al., 2014; Riemer et al., 2019; Rolnick et al., 2019) communities.

Original: We can distinguish optimization-based and context-based Meta-RL families of methods.
Revised: We can distinguish optimization-based and context-based Meta-RL families of methods.

Original: Optimization-based Meta-RL methods (Finn et al., 2017; Al-Shedivat et al., 2018; Rothfuss et al., 2019; Clavera et al., 2018) sample training contexts from a stationary distribution P C and optimize model or policy parameters, which can be quickly adapted to a test context.
Revised: Optimization-based Meta-RL methods (Finn et al., 2017; Al-Shedivat et al., 2018; Rothfuss et al., 2019; Clavera et al., 2018) sample training contexts from a stationary distribution P C and optimize model or policy parameters, which can be quickly adapted to a test context.

Original: In contrast, context-based Meta-RL methods (Rakelly et al., 2019; Zintgraf et al., 2020; Xie et al., 2020; Duan et al., 2016) attempt to infer a deterministic representation or a probabilistic belief on the context from the episode history.
Revised: In contrast, context-based Meta-RL methods (Rakelly et al., 2019; Zintgraf et al., 2020; Xie et al., 2020; Duan et al., 2016) attempt to infer a deterministic representation or a probabilistic belief on the context from the episode history.

Original: We also remark recent work on domain shift while controlling from pixels Hansen & Wang (2021);
Revised: We also remark recent work on domain shift while controlling from pixels Hansen & Wang (2021);

Original: Kostrikov et al.
Revised: Kostrikov et al.

Original: 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.
Revised: 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.

Original: Besides the episodic nature of the context change, this setting differs from ours on two fronts.
Revised: Besides the episodic nature of the context change, this setting differs from ours on two fronts.

Original: First, controlling from pictures constitutes a POMDP problem, where the true states (positions, velocities, accelerations) are observed through a proxy (through pixels) and hence need to be inferred.
Revised: First, controlling from pictures constitutes a POMDP problem, where the true states (positions, velocities, accelerations) are observed through a proxy (through pixels) and hence need to be inferred.

Original: Second, the underlying dynamics stay the same (Hansen & Wang, 2021; Kostrikov et al., 2020) and only observations change.
Revised: Second, the underlying dynamics stay the same (Hansen & Wang, 2021; Kostrikov et al., 2020) and only observations change.

Original: Our case is the opposite: the underlyinga dynamics change, but the observation  function stays the same.
Revised: Our case is the opposite: the underlying dynamics change, but the observation a function stays the same.

Original: In future work we aim to extend our methods to control from pixels.
Revised: In future work we aim to extend our methods to control from pixels.

Revised: We also note the work by Ball et al.

Revised: (2021), who proposed to learn a contextual policy by augmenting an ofﬂine world model.

Revised: In the online setting the algorithm had to learn only the current context thus improving efﬁciency of the learning procedure while delivering good performance.

Original: Continual (lifelong) RL mainly adopts a context incremental setting, where the agent is exposed to a sequence of contexts (Delange et al., 2021).
Revised: Continual (lifelong) RL mainly adopts a context incremental setting, where the agent is exposed to a sequence of contexts (Delange et al., 2021).

Original: While the agent’s goal is still to adapt efﬁciently to the unseen contexts, the emphasis is 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).
Revised: While the agent’s goal is still to adapt efﬁciently to the unseen contexts, the emphasis is 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).

Original: We further remark that the Markovian context model and episodic non-stationarity assumptions are often incompatible.
Revised: We further remark that the Markovian context model and episodic non-stationarity assumptions are often incompatible.

Original: Indeed, learning the transition model for P C in C-MDPs with an episodic non-stationarity constraint is somewhat redundant because P C can be assumed to be a stationary distribution rather than a Markov process.
Revised: Indeed, learning the transition model for P C in C-MDPs with an episodic non-stationarity constraint is somewhat redundant because P C can be assumed to be a stationary distribution rather than a Markov process.

Original: 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.
Revised: 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.

Original: This can lead not only to suboptimal policies, but also to policies not solving the task at all.
Revised: This can lead not only to suboptimal policies, but also to policies not solving the task at all.

Original: Xie et al.
Revised: Xie et al.

Original: (2020) tried to combine the two settings by taking a hierarchical view on non-stationary modeling.
Revised: (2020) tried to combine the two settings by taking a hierarchical view on non-stationary modeling.

Original: 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.
Revised: 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.

Original: 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 a fast time-scale prescribed by the steps.
Revised: 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 a fast time-scale prescribed by the steps.

Original: While this approach has its merits, it also has the some limitations, e.g., the hierarchical model is artiﬁcially imposed on the learning process.
Revised: While this approach has its merits, it also has the some limitations, e.g., the hierarchical model is artiﬁcially imposed on the learning process.

Original: Settings C and D.
Revised: Settings C and D.

Original: Many frameworks can ﬁt into the settings C and D, (c.f., Ong et al.
Revised: Many frameworks can ﬁt into the settings C and D, (c.f., Ong et al.

Original: Chandak et al.
Revised: Chandak et al.

Original: ; Zhang et al.
Revised: ; Zhang et al.

Original: , however, this makes it hard to compare the differences and similarities between them.
Revised: , however, this makes it hard to compare the differences and similarities between them.

Original: We will refer the reader to the review by Khetarpal et al.
Revised: We will refer the reader to the review by Khetarpal et al.

Original: for a further discussion.
Revised: for a further discussion.

Original: We make a few comments, however, on change point detection methods and on the relation to POMDP formulation (Astrom, 1965; Hauskrecht, 2000) in the case of unobservable contexts.
Revised: We make a few comments, however, on change point detection methods and on the relation to POMDP formulation (Astrom, 1965; Hauskrecht, 2000) in the case of unobservable contexts.

Original: Change-point detection methods can be readily used for the context estimation in an online fashion (da Silva et al., 2006; Hadoux et al., 2014; Banerjee et al., 2017; Padakandla et al., 2019;
Revised: Change-point detection methods can be readily used for the context estimation in an online fashion (da Silva et al., 2006; Hadoux et al., 2014; Banerjee et al., 2017; Padakandla et al., 2019;

Original: Li et al., 2019), however, these methods either track change-points in a heuristic way (da Silva et al., 2006;
Revised: Li et al., 2019), however, these methods either track change-points in a heuristic way (da Silva et al., 2006;

Original: Hadoux et al., 2014) or assume strong prior knowledge on the non-stationarity (Banerjee et al., 2017;
Revised: Hadoux et al., 2014) or assume strong prior knowledge on the non-stationarity (Banerjee et al., 2017;

Original: Padakandla et al., 2019).
Revised: Padakandla et al., 2019).

Original: A POMDP is deﬁned by tuple M po =
Revised: A POMDP is deﬁned by tuple M po =

Original: {X , A , O , P X , P O , R , γ d , p ( x 0 ) } , where X , A , O are the (unobservable) state, action, observation spaces, respectively; P X :
Revised: {X , A , O , P X , P O , R , γ d , p ( x 0 ) } , where X , A , O are the (unobservable) state, action, observation spaces, respectively; P X :

Original: X ×
Revised: X ×

Original: A ×
Revised: A ×

Original: X → [0 , 1] is the state transition probability; P O : X ×A×O → [0 , 1] is the conditional observation probability; and p ( x 0 ) is the initial state distribution.
Revised: X → [0 , 1] is the state transition probability; P O : X ×A×O → [0 , 1] is the conditional observation probability; and p ( x 0 ) is the initial state distribution.

Original: In our case, the state is x t = ( s t , z t ) and the observation is o t = s t .
Revised: In our case, the state is x t = ( s t , z t ) and the observation is o t = s t .

Original: A key step in POMDP literature is estimating or approximating the belief, i.e., the probability distribution of the current state using history of observations and actions (Hausknecht & Stone, 2015; Zhu et al., 2017; Igl et al., 2018).
Revised: A key step in POMDP literature is estimating or approximating the belief, i.e., the probability distribution of the current state using history of observations and actions (Hausknecht & Stone, 2015; Zhu et al., 2017; Igl et al., 2018).

Original: Since the belief is generally inﬁnite-dimensional, it is prudent to resort to sampling, point-estimates or other approximations (Hausknecht & Stone, 2015; Zhu et al., 2017; Igl et al., 2018).
Revised: Since the belief is generally inﬁnite-dimensional, it is prudent to resort to sampling, point-estimates or other approximations (Hausknecht & Stone, 2015; Zhu et al., 2017; Igl et al., 2018).

Original: In our case, this is not necessary, however, as the context is a discrete variable.
Revised: In our case, this is not necessary, however, as the context is a discrete variable.

Original: We further stress, that while general C-MDPs can be represented using POMDPs, this representation may not offer any beneﬁts at all due to complexity of belief estimation.
Revised: We further stress, that while general C-MDPs can be represented using POMDPs, this representation may not offer any beneﬁts at all due to complexity of belief estimation.

Original: Finally, we can classify the literature by cardinality of C .
Revised: Finally, we can classify the literature by cardinality of C .

Original: Some Meta-RL algorithms mainly consider a continuous contextual set C by implicitly or explicitly parametrizing C with real-valued vectors (DoshiVelez & Konidaris, 2016; Zintgraf et al., 2020;
Revised: Some Meta-RL algorithms mainly consider a continuous contextual set C by implicitly or explicitly parametrizing C with real-valued vectors (DoshiVelez & Konidaris, 2016; Zintgraf et al., 2020;

Original: Xie et al., 2020; Rakelly et al., 2019).
Revised: Xie et al., 2020; Rakelly et al., 2019).

Original: On the other hand, exploiting a discrete C can be motivated by controlling switching dynamic systems (Fox et al., 2008a) and may achieve better interpretability.
Revised: On the other hand, exploiting a discrete C can be motivated by controlling switching dynamic systems (Fox et al., 2008a) and may achieve better interpretability.

Original: Due to complexity, however, some works ﬁx the contextual cardinality |C| (Choi et al., 2000; Banerjee et al., 2017; Padakandla et al., 2019; Lee et al., 2019), or infer |C| from data in an online fashion (Xu et al., 2020; da Silva et al., 2006; Hadoux et al., 2014).
Revised: Due to complexity, however, some works ﬁx the contextual cardinality |C| (Choi et al., 2000; Banerjee et al., 2017; Padakandla et al., 2019; Lee et al., 2019), or infer |C| from data in an online fashion (Xu et al., 2020; da Silva et al., 2006; Hadoux et al., 2014).

Original: Besides, most works in continual RL adopt a setting, where there is no explicit speciﬁcation on C , but the agent can directly access the discrete contexts of all time steps/episodes (Rolnick et al., 2019; Riemer et al., 2019).
Revised: Besides, most works in continual RL adopt a setting, where there is no explicit speciﬁcation on C , but the agent can directly access the discrete contexts of all time steps/episodes (Rolnick et al., 2019; Riemer et al., 2019).

Original: In contrast, our approach learns the context evolution from data. 
Revised: In contrast, our approach learns the context evolution from data. a

Original: B P ROOFS B. P ROOF OF T HEOREM
Revised: B P ROOFS B. P ROOF OF T HEOREM

Original: For completeness, we restate the theorem below.
Revised: For completeness, we restate the theorem below.

Original: a Theorem A1 Consider a Markov chain p t = p t − 1 R with a stationary distribution p ∞ and distilled
Revised: Theorem A1 Consider a Markov chain p t = p t − 1 R with a stationary distribution p ∞ and distilled

Original: I 

Revised: I

Original: = { i | p ∞ i ≥ ε distil } and spurious I 

Revised: = { i | p ∞ i ≥ ε distil } and spurious I

Original: = { i | p ∞ i < ε distil } state indexes, respectively.
Revised: = { i | p ∞ i < ε distil } state indexes, respectively.

Original: Then a) the matrix (cid:98) R =
Revised: Then a) the matrix (cid:98) R =

Original: R
Revised: R

Original: I 1 , I 1 + R
Revised: I 1 , I 1 + R

Original: I 1 , I 2 ( I − R
Revised: I 1 , I 2 ( I − R

Original: I 2 , I 2 ) − 1 R
Revised: I 2 , I 2 ) − 1 R

Original: I 2 , I 1 is a valid probability transition matrix; b) the
Revised: I 2 , I 1 is a valid probability transition matrix; b) the

Original: Markov chain (cid:98) p t = (cid:98) p t − 1 (cid:98) R is such that its stationary distribution (cid:98) p ∞ ∝ p ∞I 1 .
Revised: Markov chain (cid:98) p t = (cid:98) p t − 1 (cid:98) R is such that its stationary distribution (cid:98) p ∞ ∝ p ∞I 1 .

Original: Let us ﬁrst provide an insight into our technical result.
Revised: Let us ﬁrst provide an insight into our technical result.

Original: In order to so consider the Markov chain evolution:
Revised: In order to so consider the Markov chain evolution:

Original: [Equation]
Revised: [Equation]

Original: Now assume that p t I 2 ≈ p ∞I 2 for all large enough t , which leads to
Revised: Now assume that p t I 2 ≈ p ∞I 2 for all large enough t , which leads to

Original: [Equation]
Revised: [Equation]

Original: If p ∞I 2 is small enough, then the Markov chain will rarely end up in the states with indexes I 2 .
Revised: If p ∞I 2 is small enough, then the Markov chain will rarely end up in the states with indexes I 2 .

Original: Meaning that we can remove these states, but we would need to take them into account while computing the new probability transitions.
Revised: Meaning that we can remove these states, but we would need to take them into account while computing the new probability transitions.

Original: Furthermore, we need to do so while obtaining a new
Revised: Furthermore, we need to do so while obtaining a new

Original: Markov chain in the process.
Revised: Markov chain in the process.

Original: The solution to this question is straightforward, i.e., we can solve for p ∞I 2 to obtain:
Revised: The solution to this question is straightforward, i.e., we can solve for p ∞I 2 to obtain:

Original: [Equation]
Revised: [Equation]

Original: Remarkably, the resulting process is also a Markov chain!
Revised: Remarkably, the resulting process is also a Markov chain!

Original: In order to verify this we need to prove the following:
Revised: In order to verify this we need to prove the following:

Original: a) the matrix I − R
Revised: a) the matrix I − R

Original: I 2 , I 2 is invertible ; b) the matrix ( I − R
Revised: I 2 , I 2 is invertible ; b) the matrix ( I − R

Original: I 2 , I 2 ) − 1 is nonnegative ; c) the right eigenvector of the matrix R I 1 , I 

Revised: I 2 , I 2 ) − 1 is nonnegative ; c) the right eigenvector of the matrix R I 1 , I

Original: + R
Revised: + R

Original: I 1 , I 

Revised: I 1 , I

Original: ( I − R I 2 , I 2 ) − 1 R
Revised: ( I − R I 2 , I 2 ) − 1 R

Original: I 2 , I 1 can be chosen to be the vector of ones 1 .
Revised: I 2 , I 1 can be chosen to be the vector of ones 1 .

Original: We will need to develop some mathematical tools in order to prove these statements.
Revised: We will need to develop some mathematical tools in order to prove these statements.

Original: For a matrix
Revised: For a matrix

Original: A ∈ R n × n the spectral radius r ( A ) is the maximum absolute value of the eigenvalues
Revised: A ∈ R n × n the spectral radius r ( A ) is the maximum absolute value of the eigenvalues

Original: λ i of A and r ( A ) = max i | λ i ( A ) | .
Revised: λ i of A and r ( A ) = max i | λ i ( A ) | .

Original: The matrix A ∈ R n × m is called nonnegative if all its entries are 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 (cid:29) 0 .
Revised: The matrix A ∈ R n × m is called nonnegative if all its entries are 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 (cid:29) 0 .

Original: The matrix A ∈ R n × n is called reducible if there exists a permutation matrix T such that T AT T = (cid:18) B
Revised: The matrix A ∈ R n × n is called reducible if there exists a permutation matrix T such that T AT T = (cid:18) B

Original: C
Revised: C

Original: D (cid:19) for some square matrices B and D .
Revised: D (cid:19) for some square matrices B and D .

Original: The matrix is called irreducible if such a permutation matrix does not exist.
Revised: The matrix is called irreducible if such a permutation matrix does not exist.

Original: The matrix A ∈ R n × n is called M-matrix, if all its off-diagonal entries are nonpositive and it can be represented as A
Revised: The matrix A ∈ R n × n is called M-matrix, if all its off-diagonal entries are nonpositive and it can be represented as A

Original: =
Revised: =

Original: s I − B with s ≥ r ( B ) (Berman & Plemmons, 1994).
Revised: s I − B with s ≥ r ( B ) (Berman & Plemmons, 1994).

Original: Interestingly the inverse of an M-matrix is always nonnegative (the opposite is generally false) (Berman & Plemmons, 1994).
Revised: Interestingly the inverse of an M-matrix is always nonnegative (the opposite is generally false) (Berman & Plemmons, 1994).

Original: We will also use the following results:
Revised: We will also use the following results:

Original: Proposition A1 (Perron-Frobenius Theorem) Let A ∈ R n × n be a nonnegative matrix, then the spectral radius r ( A ) is an eigenvalue of A and the corresponding left and right eigenvectors can be chosen to be nonnegative.
Revised: Proposition A1 (Perron-Frobenius Theorem) Let A ∈ R n × n be a nonnegative matrix, then the spectral radius r ( A ) is an eigenvalue of A and the corresponding left and right eigenvectors can be chosen to be nonnegative.

Original: If A is additionally irreducible then r ( A ) is a simple eigenvalue of A and the corresponding left and right eigenvectors can be chosen to be positive.
Revised: If A is additionally irreducible then r ( A ) is a simple eigenvalue of A and the corresponding left and right eigenvectors can be chosen to be positive.

Original: Proposition A2 Consider A , B ∈ R n × n and let A > B ≥ 0 , then r ( A ) ≥ r ( B ) .
Revised: Proposition A2 Consider A , B ∈ R n × n and let A > B ≥ 0 , then r ( A ) ≥ r ( B ) .

Original: If additionally
Revised: If additionally

Original: A is irreducible then the inequality is strict r ( A )
Revised: A is irreducible then the inequality is strict r ( A )

Original: > r ( B ) .
Revised: > r ( B ) .

Original:  Proof: This result is well-known, but for completeness we show the proof here (we adapt a similar technique to Noutsos (2006)).
Revised: a Proof: This result is well-known, but for completeness we show the proof here (we adapt a similar technique to Noutsos (2006)).

Original: 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 ( A ) .
Revised: 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 ( A ) .

Original: Now we have
Revised: Now we have

Original: [Equation]
Revised: [Equation]

Original: a If a T b is positive, then the ﬁrst part is shown.
Revised: If a T b is positive, then the ﬁrst part is shown.

Original: If a T b = 0 , then we can make a continuity argument by perturbing A and B to A (cid:48) and B (cid:48) in such a way that for their corresponding eigenvectors we have ( a (cid:48) )
Revised: If a T b = 0 , then we can make a continuity argument by perturbing A and B to A (cid:48) and B (cid:48) in such a way that for their corresponding eigenvectors we have ( a (cid:48) )

Original: T b (cid:48) > 0 .
Revised: T b (cid:48) > 0 .

Original: Hence the ﬁrst part of the proof is shown.
Revised: Hence the ﬁrst part of the proof is shown.

Original: Now consider the case when A is irreducible.
Revised: Now consider the case when A is irreducible.

Original: According to Proposition A1, since A is irreducible the spectral radius r ( A ) is a simple eigenvalue, the corresponding eigenvector a can be chosen to be positive.
Revised: According to Proposition A1, since A is irreducible the spectral radius r ( A ) is a simple eigenvalue, the corresponding eigenvector a can be chosen to be positive.

Original: Let us prove the second part by contradiction and assume that r ( A )
Revised: Let us prove the second part by contradiction and assume that r ( A )

Original: =
Revised: =

Original: r ( B ) , which implies that a T Ab = a T Bb and consequently Ab =
Revised: r ( B ) , which implies that a T Ab = a T Bb and consequently Ab =

Original: Bb since Ab ≥ Bb and a (cid:29) 0 .
Revised: Bb since Ab ≥ Bb and a (cid:29) 0 .

Original: Now since we also have A > B this means that the vector b has at least one zero entry.
Revised: Now since we also have A > B this means that the vector b has at least one zero entry.

Original: Since we also have
Revised: Since we also have

Original: Ab = r ( A ) b , the vector b is an eigenvector of A corresponding to r ( A ) and has zero entries by assumption above.
Revised: Ab = r ( A ) b , the vector b is an eigenvector of A corresponding to r ( A ) and has zero entries by assumption above.

Original: This contradicts Perron-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.
Revised: This contradicts Perron-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.

Original: Therefore, r ( A ) >
Revised: Therefore, r ( A ) >

Original: r ( B ) .
Revised: r ( B ) .

Original: (cid:3)
Revised: [Equation]

Original: Now we proceed with the proof.
Revised: Now we proceed with the proof.

Original: 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 A1.
Revised: 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 A1.

Original: It is also almost surely true if the transition model is learned.
Revised: It is also almost surely true if the transition model is learned.

Original: This is because the subset of reducible matrices is measure zero in the set of nonnegative matrices.
Revised: This is because the subset of reducible matrices is measure zero in the set of nonnegative matrices.

Original: Recall that we have a Markov chain p t =
Revised: Recall that we have a Markov chain p t =

Original: p t − 1 R with a stationary distribution p ∞ , that is p ∞ = p ∞ R .
Revised: p t − 1 R with a stationary distribution p ∞ , that is p ∞ = p ∞ R .

Original: Consider the index sets: the distilled context indexes I 

Revised: Consider the index sets: the distilled context indexes I

Original: = { i | p ∞ i ≥ ε distil } and the spurious context indexes I 2 = { i | p ∞
Revised: = { i | p ∞ i ≥ ε distil } and the spurious context indexes I 2 = { i | p ∞

Original: i < ε distil } .
Revised: i < ε distil } .

Original: First, we will show that the matrix (cid:98) R
Revised: First, we will show that the matrix (cid:98) R

Original: = R I 1 , I 1 + R I 1 , I 2 ( I − R
Revised: [Equation]

Original: I 2 , I 2 ) − 1 R

Original: I 2 , I 1 is nonnegative.
Revised: , I 1 is nonnegative.

Original: Since the matrix R
Revised: Since the matrix R

Original: I 2 , I 2 is a nonnegative submatrix of R , due to Proposition A2 we have that r ( R I 2 , I 2 ) <
Revised: I 2 , I 2 is a nonnegative submatrix of R , due to Proposition A2 we have that r ( R I 2 , I 2 ) <

Original: r ( R ) = 1 .
Revised: r ( R ) = 1 .

Original: This means that I − R
Revised: This means that I − R

Original: I 2 , I 2 is an M-matrix, which implies that it is invertible and its inverse is a nonnegative matrix.
Revised: I 2 , I 2 is an M-matrix, which implies that it is invertible and its inverse is a nonnegative matrix.

Original: Hence (cid:98) R is nonnegative (Berman & Plemmons, 1994).
Revised: Hence (cid:98) R is nonnegative (Berman & Plemmons, 1994).

Original: Since R describes a Markov chain we have R 1 = 1 (where 1 is the vector of ones) and p ∞ R = p ∞
Revised: Since R describes a Markov chain we have R 1 = 1 (where 1 is the vector of ones) and p ∞ R = p ∞

Original: [Equation]
Revised: [Equation]

Original: Now using elementary algebra we can establish that (cid:98) R 1 I 1 = 1 I 1 :
Revised: Now using elementary algebra we can establish that (cid:98) R 1 I 1 = 1 I 1 :

Original: [Equation]
Revised: [Equation]

Original: due to
Revised: to

Original: [Equation]
Revised: [Equation]

Original: due to 
Revised: due to (A2a

Original: [Equation]
Revised: [Equation]

Original:  Similarly for p ∞I 1 (cid:98) R = p ∞I 1 we have:
Revised: 1 I 1 . Similarly for p ∞I 1 (cid:98) R = p ∞I 1 we have:

Original: [Equation]
Revised: [Equation]

Original: Since p ∞ I 1 (cid:98) R = p ∞ I 1 , the normalized vector p ∞ I 1 is the stationary distribution of the distilled Markov chain and hence (cid:98) p ∞ ∝ p ∞I 1 .
Revised: Since p ∞I 1 (cid:98) R = p ∞I 1 , the normalized vector p ∞I 1 is the stationary distribution of the distilled Markov chain and hence (cid:98) p ∞ ∝ p ∞ I 1 .

Original: This completes the proof.
Revised: This completes the proof.

Original:  B. D ERIVATION OF DYNAMIC PROGRAMMING PRINCIPLE AND PROOF OF T HEOREM
Revised: a B. D ERIVATION OF DYNAMIC PROGRAMMING PRINCIPLE AND PROOF OF T HEOREM

Original: For completeness we reproduce the theorem formulation belowa
Revised: For completeness we reproduce the theorem formulation below

Original: Theorem A2 a)
Revised: Theorem A2 a)

Original: The belief of z can be computed as p ( z t +1 | I Ct )
Revised: The belief of z can be computed as p ( z t +1 | I Ct )

Original: =
Revised: =

Original: b zt +1 , where ( b zt +1 )
Revised: b zt +1 , where ( b zt +1 )

Original: i ∝ N i = (cid:80) j p ( s t +1 | s t , θ i , a t ) ρ ji ( b zt ) j , where ( b zt )
Revised: i ∝ N i = (cid:80) j p ( s t +1 | s t , θ i , a t ) ρ ji ( b zt ) j , where ( b zt )

Original: i are the entries of b zt ; b) the optimal policy can be computed as π ( s , b z ) =
Revised: i are the entries of b zt ; b) the optimal policy can be computed as π ( s , b z ) =

Original: argmax a Q ( s , b z , a ) , where the value function satisﬁes the dynamic programming principle Q ( s t , b zt , [Equation]
Revised: argmax a Q ( s , b z , a ) , where the value function satisﬁes the dynamic programming principle Q ( s t , b zt , a t )

Original: =
Revised: =

Original: [Equation]
Revised: r ( s t , b zt , a t ) + γ (cid:82) (cid:80) i N i max a t +1 Q ( s t +1 , b zt +1 , a t +1 ) d s t +1 .

Original: The following deﬁnitions and derivations are in line with previous work by Hauskrecht (2000);
Revised: The following deﬁnitions and derivations are in line with previous work by Hauskrecht (2000);

Original: Porta et al.
Revised: Porta et al.

Original: We introduce the complete information state at the time t as follows:
Revised: We introduce the complete information state at the time t as follows:

Original: [Equation]
Revised: [Equation]

Original: where b 0 = p ( z 0 ) δ o 0 ( s 0 ) .
Revised: where b 0 = p ( z 0 ) δ o 0 ( s 0 ) .

Original: In order to tackle the intractability of the complete information state, one can use any information state that is sufﬁcient in some sense:
Revised: In order to tackle the intractability of the complete information state, one can use any information state that is sufﬁcient in some sense:

Original: Deﬁnition A1 Consider a partially observable Markov decision process {X , A , O , P X , P O , R , b 0 } .
Revised: Deﬁnition A1 Consider a partially observable Markov decision process {X , A , O , P X , P O , R , b 0 } .

Original: Let I be an information state space and ξ : I × A × O → I be an update function deﬁning an information process I t = ξ ( I t − 1 , a t − 1 , o t ) .
Revised: Let I be an information state space and ξ : I × A × O → I be an update function deﬁning an information process I t = ξ ( I t − 1 , a t − 1 , o t ) .

Original: We say that I St is a sufﬁcient information process with regard to the optimal control if it is an information process and for any time step t , it satisﬁes
Revised: We say that I St is a sufﬁcient information process with regard to the optimal control if it is an information process and for any time step t , it satisﬁes

Original: [Equation]
Revised: [Equation]

Original: As a sufﬁcient information state of the process I St , we will use the belief b t deﬁned as follows:
Revised: As a sufﬁcient information state of the process I St , we will use the belief b t deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: Effectively, we introduce the belief b s t = p ( s t | o ≤ t , a <t ) of the state s t and the belief b zt = p ( z t | o ≤ t , a <t ) of the state z t given an observation o t .
Revised: Effectively, we introduce the belief b s t = p ( s t | o ≤ t , a <t ) of the state s t and the belief b zt = p ( z t | o ≤ t , a <t ) of the state z t given an observation o t .

Original: However, since o t = s t we can consider only the belief of z t :
Revised: However, since o t = s t we can consider only the belief of z t :

Original: [Equation]
Revised: [Equation]

Original: One of the features of our model is that the updates of the beliefs can be analytically computed.
Revised: One of the features of our model is that the updates of the beliefs can be analytically computed.

Original: Lemma A

Revised: Lemma A

Original: The belief b zt = p ( z t | I Ct ) is a sufﬁcient state information with the updates:
Revised: The belief b zt = p ( z t | I Ct ) is a sufﬁcient state information with the updates:

Original: [Equation]
Revised: [Equation]

Original: Proof: First, let us check that the belief b t satisﬁes both conditions of Deﬁnition A1.
Revised: Proof: First, let us check that the belief b t satisﬁes both conditions of Deﬁnition A1.

Original: By deﬁnition of our belief we have p ( x t | I
Revised: By deﬁnition of our belief we have p ( x t | I

Original: Ct − 1 )
Revised: Ct − 1 )

Original: =
Revised: =

Original: b t = p ( x t | I St − 1 ) , which satisﬁes the ﬁrst condition, as for the second condition we have:
Revised: b t = p ( x t | I St − 1 ) , which satisﬁes the ﬁrst condition, as for the second condition we have:

Original: [Equation]
Revised: [Equation]

Original: Now straightforward derivations yield:
Revised: Now straightforward derivations yield:

Original: [Equation]
Revised: [Equation]

Original: completing the proof.
Revised: completing the proof

Original: [Equation]
Revised: [Equation]

Original:  As discussed by Porta et al.
Revised: a As discussed by Porta et al.

Original: [Equation] POMDP in continuous state [Equation] action and observation spaces also satisfy Bellman dynamic programming principle, albeit in a different space.
Revised: (2006), POMDP in continuous state, action and observation spaces also satisfy Bellman dynamic programming principle, albeit in a different space.

Original: Recall that the control problem is typically formulated as
Revised: Recall that the control problem is typically formulated as

Original: [Equation]
Revised: [Equation]

Original: a where ˜ τ = { x 0 , a 0 , . . 
Revised: where ˜ τ = { x 0 , a 0 , . . .

Original: , x T − 1 , a T − 1 , x T } and the horizon T can be inﬁnite.
Revised: , x T − 1 , a T − 1 , x T } and the horizon T can be inﬁnite.

Original: As we do not have access to the state transitions we need to rewrite the problem in the observation or the belief spaces.
Revised: As we do not have access to the state transitions we need to rewrite the problem in the observation or the belief spaces.

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: where τ = { b 0 , o 0 , a 0 , . .
Revised: where τ = { b 0 , o 0 , a 0 , . .

Original: , o T − 1 , a T − 1 , o T } and:
Revised: , o T − 1 , a T − 1 , o T } and:

Original: [Equation]
Revised: [Equation]

Original: Given this reparameterization we can introduce the value functions and derive the Bellman equation similarly to Porta et al.
Revised: Given this reparameterization we can introduce the value functions and derive the Bellman equation similarly to Porta et al.

Original: (2006), which can be written in terms of the Q-function as follows:
Revised: (2006), which can be written in terms of the Q-function as follows:

Original: [Equation]
Revised: [Equation]

Original: We, however, have an additional structure that allows for simpliﬁed value functions.
Revised: We, however, have an additional structure that allows for simpliﬁed value functions.

Original: First note that ˜ r t ( b t , a t ) =
Revised: First note that ˜ r t ( b t , a t ) =

Original: r t ( s t , b zt , a t ) , i.e., our reward depends directly on the observation and the belief of z .
Revised: r t ( s t , b zt , a t ) , i.e., our reward depends directly on the observation and the belief of z .

Original: Now we need to estimate E max a t +1 Q ( b s t +1 , b zt +1 , a t +1 ) , where the expectation is taken over new observation, the probability distribution of which can be computed as:
Revised: Now we need to estimate E max a t +1 Q ( b s t +1 , b zt +1 , a t +1 ) , where the expectation is taken over new observation, the probability distribution of which can be computed as:

Original: [Equation]
Revised: [Equation]

Original: This allows rewriting the second part of the Bellman equation as follows:
Revised: This allows rewriting the second part of the Bellman equation as follows:

Original: [Equation]
Revised: [Equation]

Original: where N i , b zt +1 depend on s t +1 .
Revised: where N i , b zt +1 depend on s t +1 .

Original: Finally, we have:
Revised: Finally, we have:

Original: [Equation]
Revised: [Equation]

Original: which completes the proof.
Revised: which completes the proof

Revised: [Equation]

Original:  B. P ERFORMANCE GAIN FOR OBSERVABLE CONTEXTS
Revised: a B. P ERFORMANCE GAIN FOR OBSERVABLE CONTEXTS

Original: It is not surprising that observing the ground truth of the contexts should improve the maximum expected return.
Revised: It is not surprising that observing the ground truth of the contexts should improve the maximum expected return.

Original: In particular, even knowing the ground truth context model we can correctly estimate the context z t +1 only a posteriori , i.e., after observing the next state s t +1 .
Revised: In particular, even knowing the ground truth context model we can correctly estimate the context z t +1 only a posteriori , i.e., after observing the next state s t +1 .

Original: This means that with every context switch we will mislabel a context with a high probability.
Revised: This means that with every context switch we will mislabel a context with a high probability.

Original: This will lead to a sub-optimal action and performance loss, which the following result quantiﬁes using the value functions.
Revised: This will lead to a sub-optimal action and performance loss, which the following result quantiﬁes using the value functions.

Original: a Theorem A3 Assume we know the true transition model of the contexts and states and consider two settings: we observe the ground truth z t
Revised: Theorem A3 Assume we know the true transition model of the contexts and states and consider two settings: we observe the ground truth z t

Original: and we estimate it using b zt .
Revised: and we estimate it using b zt .

Original: Assume we computed the optimal model-based policy π ( ·| s t , b zt ) with the return R and the optimal ground-truth policy π gt ( ·| s t , z t +1 ) with the corresponding optimal value functions V gt ( s , z ) and Q gt ( s , z, a ) , then:
Revised: Assume we computed the optimal model-based policy π ( ·| s t , b zt ) with the return R and the optimal ground-truth policy π gt ( ·| s t , z t +1 ) with the corresponding optimal value functions V gt ( s , z ) and Q gt ( s , z, a ) , then:

Original: E
Revised: E

Original: [Equation]

Original: [Equation]

Original: s 0 , z

Original: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where M is the number of misidentiﬁed context switches in a trajectory τ .
Revised: where M is the number of misidentiﬁed context switches in a trajectory τ .

Original: Proof: Let us consider the best case scenario.
Revised: Proof: Let us consider the best case scenario.

Original: As we assume that we have access to the ground truth model for computing the policy π ( ·| s t , b zt ) , we can assume that the ground truth value of z t has the highest probability mass in the vector b zt .
Revised: As we assume that we have access to the ground truth model for computing the policy π ( ·| s t , b zt ) , we can assume that the ground truth value of z t has the highest probability mass in the vector b zt .

Original: That is, we can assume that we can identify the correct context a posteriori .
Revised: That is, we can assume that we can identify the correct context a posteriori .

Original: In effect, we can use a priori estimate of z t +1 by transitioning to the next time step, i.e., using the vector b zt R .
Revised: In effect, we can use a priori estimate of z t +1 by transitioning to the next time step, i.e., using the vector b zt R .

Original: We can also assume that the action distributions of the policy π and the ground truth policy π gt are identical provided our a priori estimate of the context and the ground truth context are the same.
Revised: We can also assume that the action distributions of the policy π and the ground truth policy π gt are identical provided our a priori estimate of the context and the ground truth context are the same.

Original: 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.
Revised: 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.

Original: Now if we can estimate the effect of this mismatch on the performance, this will provide us with a lower bound on the performance gain.
Revised: Now if we can estimate the effect of this mismatch on the performance, this will provide us with a lower bound on the performance gain.

Original: Let V gt ( s ) , Q gt ( s, a ) be the optimal value functions for the ground truth policy π gt satisfying the
Revised: Let V gt ( s ) , Q gt ( s, a ) be the optimal value functions for the ground truth policy π gt satisfying the

Original: Bellman equation:
Revised: Bellman equation:

Original: [Equation]
Revised: [Equation]

Original: Consider a particular realization of the stochastic context variable z t (which is independent of s t , a t ) and assume the context switched only once at t 1 .
Revised: Consider a particular realization of the stochastic context variable z t (which is independent of s t , a t ) and assume the context switched only once at t 1 .

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: In effect
Revised: In effect

Original: [Equation]
Revised: [Equation]

Original: we are using the Q gt function to estimate the performance loss of one mistake.
Revised: we are using the Q gt function to estimate the performance loss of one mistake.

Original: The same procedure can be repeated for context realizations with M misidentiﬁed switches, where the number
Revised: The same procedure can be repeated for context realizations with M misidentiﬁed switches, where the number

Original: M depends on the realization of the context variable
Revised: M depends on the realization of the context variable

Original: [Equation]
Revised: [Equation]

Original: Now averaging over the context realizations proves the result.
Revised: Now averaging over the context realizations proves the result.

Original: (cid:3)
Revised: [Equation] a

Original: C A LGORITHM D ETAILS
Revised: C A LGORITHM D ETAILS

Original: 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.
Revised: 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.

Original: We ﬁrstlya brieﬂy comment on each on these components to give an overview of the results and then explain our main contributions to each.
Revised: We ﬁrstly brieﬂy comment on each on these components to give an overview of the results and then explain our main contributions to each.

Original: 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.
Revised: 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.

Original: (2008a;b).
Revised: (2008a;b).

Original: We improve the model by proposing a context spuriosity measure allowing for reconstruction of ground truth contexts.
Revised: We improve the model by proposing a context spuriosity measure allowing for reconstruction of ground truth contexts.

Original: We then derive a model learning algorithm using probabilistic inference.
Revised: We then derive a model learning algorithm using probabilistic inference.

Original: Having a model, we can take off-the-shelf frameworks such as (Pineda et al., 2021), which can include a Model Predictive Control (MPC) approach using Cross-Entropy Minimization (CEM) (cf.
Revised: Having a model, we can take off-the-shelf frameworks such as (Pineda et al., 2021), which can include a Model Predictive Control (MPC) approach using Cross-Entropy Minimization (CEM) (cf.

Original: Chua et al.
Revised: Chua et al.

Original: (2018) and Appendix C.5), or a policy-gradient approach Soft-actor critic (cf.
Revised: (2018) and Appendix C.5), or a policy-gradient approach Soft-actor critic (cf.

Original: Haarnoja et al.
Revised: Haarnoja et al.

Original: (2018) and Appendix C.6), which are both well-suited for model-based reinforcement learning.
Revised: (2018) and Appendix C.6), which are both well-suited for model-based reinforcement learning.

Original: 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 in the main text and its proof in Appendix B.1).
Revised: 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 in the main text and its proof in Appendix B.1).

Original: We summarize our model-based approach in Algorithm A1.
Revised: We summarize our model-based approach in Algorithm A1.

Original: Algorithm A1: Learning to Control HDP-C-MDP Input
Revised: Input: ε distill - distillation threshold, N warm - number of trajectories for warm start, N trajnumber of newly collected trajectories per epoch, N epochs - number of training epochs, AGENTpolicy gradient or MPC agent

Original: Initialize AGENT with RANDOMAGENT , D =
Revised: Initialize AGENT with RANDOM , D =

Original: ; for i = 1 , .  , N epochs do
Revised: ; for i = 1 , . . . , N epochs do

Original: Sample
Revised: Sample

Original: = trajectories from the environmentAGENT ;
Revised: the environment

Original: Set Dnew =
Revised: Set D

Original: ( s i , a i ) } N traj i =1 , where s i
Revised: , a i ) } N traj i =1 , where s i

Original: = { s it } Tt = − 1 and a i = { a it } Tt = − 1 are the state and action sequences in the i -th trajectory.
Revised: = { s it } Tt = − 1 and a i = { a it } Tt = − 1 are the state and action sequences in the i -th trajectory.

Original: Set D
Revised: Set D

Original: = D ∪ 
Revised: = D ∪ D new ; Update generative model parameters by gradient ascent on ELBO in Equation 4; Perform context distillation with ε distill ; 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

Original: C. S LIGHTLY MORE DETAILS ON THE H IERARCHICAL D IRICHLET P ROCESSES
Revised: C. S LIGHTLY MORE DETAILS ON THE H IERARCHICAL D IRICHLET P ROCESSES

Original: A Dirichlet process (DP), denoted as DP ( γ, H ) , is characterized by a concentration parameter γ and a base distribution H ( λ ) deﬁned over the parameter space Θ .
Revised: A Dirichlet process (DP), denoted as DP ( γ, H ) , is characterized by a concentration parameter γ and a base distribution H ( λ ) deﬁned over the parameter space Θ .

Original: A sample G from DP ( γ, H ) is a probability distribution satisfying ( G ( A 1 ) , ..., G
Revised: A sample G from DP ( γ, H ) is a probability distribution satisfying ( G ( A 1 ) , ..., G

Original: ( A r )) ∼ Dir ( γH ( A 1 ) , ..., γH ( A r )) for every ﬁnite measurable partition
Revised: ( A r )) ∼ Dir ( γH ( A 1 ) , ..., γH ( A r )) for every ﬁnite measurable partition

Original: A 1 , ..., A r of Θ , where Dir denotes the Dirichlet distribution.
Revised: A 1 , ..., A r of Θ , where Dir denotes the Dirichlet distribution.

Original: Sampling G is often performed using the stick-breaking process (Sethuraman, 1994) and constructed by randomly mixing atoms independent and identically distributed samples θ k from H :
Revised: Sampling G is often performed using the stick-breaking process (Sethuraman, 1994) and constructed by randomly mixing atoms independent and identically distributed samples θ k from H :

Original: [Equation]
Revised: [Equation]

Original: where δ θ k is the Dirac distribution at θ k , and the resulting distribution of β = ( β 1 , ...β ∞ ) is called
Revised: where δ θ k is the Dirac distribution at θ k , and the resulting distribution of β = ( β 1 , ...β ∞ ) is called

Original: GEM ( γ ) for Grifﬁths-Engen-McCloskey (Teh et al., 2006).
Revised: GEM ( γ ) for Grifﬁths-Engen-McCloskey (Teh et al., 2006).

Original: The discrete nature of G motivates the application of DP as a non-parametric prior for mixture models with an inﬁnite number of atoms θ k .
Revised: The discrete nature of G motivates the application of DP as a non-parametric prior for mixture models with an inﬁnite number of atoms θ k .

Original: We note that the stick-breaking procedure assigns progressively smaller values to β k for large k , thus encouraging a smaller number of meaningful atoms.
Revised: We note that the stick-breaking procedure assigns progressively smaller values to β k for large k , thus encouraging a smaller number of meaningful atoms.

Original: The Hierarchical Dirichlet Process (HDP) is a group of DPs sharing a base distribution, which itself is a sample from a DP: G ∼ DP ( γ, H ) , G j ∼ DP ( α, G ) for all j = 0 , 1 , 2 , . . .
Revised: The Hierarchical Dirichlet Process (HDP) is a group of DPs sharing a base distribution, which itself is a sample from a DP: G ∼ DP ( γ, H ) , G j ∼ DP ( α, G ) for all j = 0 , 1 , 2 , . . .

Original: (Teh et al., 2006).
Revised: (Teh et al., a

Original: The distribution G guarantees that all G j inherit the same set of atoms, i.e., atoms of G , while keeping the beneﬁts of DPs in the distributions G j .
Revised: The distribution G guarantees that all G j inherit the same set of atoms, i.e., atoms of G , while keeping the beneﬁts of DPs in the distributions G j .

Original: HDPs have received a signiﬁcant attention in the literature (Teh et al., 2006; Fox et al., 2008b;a) with various applications including Markov chain modeling.a
Revised: HDPs have received a signiﬁcant attention in the literature (Teh et al., 2006; Fox et al., 2008b;a) with various applications including Markov chain modeling.

Original: In our case, the atoms { θ k } forming the context set (cid:101) C are sampled from H ( λ ) .
Revised: In our case, the atoms { θ k } forming the context set (cid:101) C are sampled from H ( λ ) .

Original: It can be shown that a random draw G j from DP ( α, G ) can be done using (cid:101) ρ j ∼ GEM ( α ) and (cid:101) θ k ∼ G .
Revised: It can be shown that a random draw G j from DP ( α, G ) can be done using (cid:101) ρ j ∼ GEM ( α ) and (cid:101) θ k ∼ G .

Original: However, since (cid:101) θ k is sampled from (cid:101) C , G j is also a distribution over (cid:101) C and
Revised: However, since (cid:101) θ k is sampled from (cid:101) C , G j is also a distribution over (cid:101) C and

Original: [Equation]
Revised: [Equation]

Original: for some ρ jk , which can be sampled using another stick-break construction (Teh et al., 2006).
Revised: for some ρ jk , which can be sampled using another stick-break construction (Teh et al., 2006).

Original: We consider its modiﬁed version introduced by Fox et al.
Revised: We consider its modiﬁed version introduced by Fox et al.

Original: [Equation]
Revised: [Equation]

Original: where k ≥ 1 , j ≥ 0 , ˜ δ jk is the Kronecker delta, the parameter κ ≥ 0 , called the sticky factor, modiﬁes the transition matrix priors encouraging self-transitions.
Revised: where k ≥ 1 , j ≥ 0 , ˜ δ jk is the Kronecker delta, the parameter κ ≥ 0 , called the sticky factor, modiﬁes the transition matrix priors encouraging self-transitions.

Original: The sticky factor serves as another measure of regularization reducing the average number of transitions.
Revised: The sticky factor serves as another measure of regularization reducing the average number of transitions.

Original: Thus DP ( α, G ) can serve as the prior for the initial context distribution ρ 0 and each row ρ j in the transition matrix R . 
Revised: Thus DP ( α, G ) can serve as the prior for the initial context distribution ρ 0 and each row ρ j in the transition matrix R . Figure

Revised: A3: A probabilistic model for C-MDP with Markovian context

Original: In summary, our probabilistic model is constructed in Equations A1,A3,A4 and illustrated in Figure A3 as a graphical model.
Revised: In summary, our probabilistic model is constructed in Equations A1,A3,A4 and illustrated in Figure A3 as a graphical model.

Original: We stress that the HDP in its stick-breaking construction assumes that | (cid:101) C| is inﬁnite and countable.
Revised: We stress that the HDP in its stick-breaking construction assumes that | (cid:101) C| is inﬁnite and countable.

Original: In practice, however, we make an approximation and set | (cid:101) C|
Revised: In practice, however, we make an approximation and set | (cid:101) C|

Original: =
Revised: =

Original: K with a large enough K .
Revised: K with a large enough K .

Original: C. V ARIATIONAL I NFERENCE FOR P ROBABILISTIC M ODELING
Revised: C. V ARIATIONAL I NFERENCE FOR P ROBABILISTIC M ODELING

Original: Recall that our context MDP is represented as follows
Revised: Recall that our context MDP is represented as follows

Original: [Equation]
Revised: [Equation]

Original: and we depict our generative model as a graphical one in Figure A3.
Revised: and we depict our generative model as a graphical one in Figure A3.

Original: Also recall that the distributions ρ j have the following priors:
Revised: Also recall that the distributions ρ j have the following priors:

Original: [Equation]
Revised: [Equation]

Original: |

Original: [Equation] ∼ Beta [Equation] where k ∈ N ≥ 1 , j ∈ N ≥ 0 , t ∈ N ≥ 0 and ˜ δ jk is the Kronecker delta function.
Revised: A5) a where k ∈ N ≥ 1 , j ∈ N ≥ 0 , t ∈ N ≥ 0 and ˜ δ jk is the Kronecker delta function.

Original: We aim to ﬁnd a variational distribution q ( ν , µ , θ ) to approximate the true posterior p ( ν , µ , θ |D ) , where D = { ( s i , a i ) } N i =1 is a data set, s i
Revised: We aim to ﬁnd a variational distribution q ( ν , µ , θ ) to approximate the true posterior p ( ν , µ , θ |D ) , where D = { ( s i , a i ) } Ni =1 is a data set, s i

Original: = { s it } T t = − 1 and a i = { a it } T t = − 1 are the state and actiona sequences in the i -th trajectory.
Revised: = { s it } Tt = − 1 and a i = { a it } Tt = − 1 are the state and action sequences in the i -th trajectory.

Original: This is achieved by minimizing KL ( q ( ν , µ , θ )
Revised: This is achieved by minimizing KL ( q ( ν , µ , θ )

Original: || p ( ν , µ , θ |D )) , or equivalently, maximizing the evidence lower bound (ELBO):
Revised: || p ( ν , µ , θ |D )) , or equivalently, maximizing the evidence lower bound (ELBO):

Original: [Equation]
Revised: [Equation]

Original: The variational distribution above involves inﬁnite-dimensional random variables ν , µ , θ .
Revised: The variational distribution above involves inﬁnite-dimensional random variables ν , µ , θ .

Original: To reach a tractable solution, we set | ˜ C|
Revised: To reach a tractable solution, we set | ˜ C|

Original: =
Revised: =

Original: K and exploit a mean-ﬁeld truncated variational distribution (Blei et al., 2006; Hughes et al., 2015; Bryant & Sudderth, 2012).
Revised: K and exploit a mean-ﬁeld truncated variational distribution (Blei et al., 2006; Hughes et al., 2015; Bryant & Sudderth, 2012).

Original: We construct the following variational distributions:
Revised: We construct the following variational distributions:

Original: [Equation]
Revised: [Equation]

Original: where the hatted symbols represent free parameters.
Revised: where the hatted symbols represent free parameters.

Original: Random variables not shown in (A6) are conditionally independent of the data, and thus can be discarded from the problem.
Revised: Random variables not shown in (A6) are conditionally independent of the data, and thus can be discarded from the problem.

Original: We maximize ELBO using stochastic gradient ascent.
Revised: We maximize ELBO using stochastic gradient ascent.

Original: In particular, given a sub-sampled batch
Revised: In particular, given a sub-sampled batch

Original: B = { ( s i , a i ) } Bi =1 , the gradient of ELBO is estimated as:
Revised: B = { ( s i , a i ) } Bi =1 , the gradient of ELBO is estimated as:

Original: [Equation]
Revised: [Equation]

Original: where we apply implicit reparameterization method (Figurnov et al., 2018; Jankowiak & Obermeyer, 2018) for gradients with respect to the expectations over Beta distributions.
Revised: where we apply implicit reparameterization method (Figurnov et al., 2018; Jankowiak & Obermeyer, 2018) for gradients with respect to the expectations over Beta distributions.

Original: For computing the gradient with respect to the likelihood term log p ( s i | a i , µ , ˆ θ ) , we exploit a message passing algorithm to integrate out the context indexes
Revised: For computing the gradient with respect to the likelihood term log p ( s i | a i , µ , ˆ θ ) , we exploit a message passing algorithm to integrate out the context indexes

Original: z
Revised: z

Original: i 1: T .
Revised: i 1: T .

Original: We present the details of the gradient computations in what follows.
Revised: We present the details of the gradient computations in what follows.

Original: Gradient of log p ( s i | a i , µ , ˆ θ ) .
Revised: Gradient of log p ( s i | a i , µ , ˆ θ ) .

Original: We drop the dependency on µ , ˆ θ in the following derivations.
Revised: We drop the dependency on µ , ˆ θ in the following derivations.

Original: We have:
Revised: We have:

Original: [Equation]
Revised: [Equation]

Original: where z i = { z it } Tt =1 is the context index sequence.
Revised: where z i = { z it } Tt =1 is the context index sequence.

Original: Since the second term equals zero, we have:
Revised: Since the second term equals zero, we have:

Original: [Equation]
Revised: [Equation]

Original: Context index posteriors p ( z i 

Revised: Context index posteriors p ( z i

Original: | s i , a i ) and p ( z it − 1 , z it | s i , a i ) required to compute the above expectation can be obtained by the message passing algorithm.
Revised: | s i , a i ) and p ( z it − 1 , z it | s i , a i ) required to compute the above expectation can be obtained by the message passing algorithm.

Original: The forward pass can be written as:
Revised: The forward pass can be written as:

Original: [Equation]
Revised: [Equation]

Original: a
Revised: a

Original: The backward pass can be written as:
Revised: The backward pass can be written as:

Original: [Equation]
Revised: [Equation]

Original: Combining the forward and backward messages, we have:
Revised: Combining the forward and backward messages, we have:

Original: [Equation]
Revised: [Equation]

Revised: The forward pass estimates the posterior context distribution at time t give the past observations and actions (i.e., for k ≤ t ), which is similar to a ﬁltering process.

Original: The forward pass estimates theposterior context distribution at time t give the past observations and actions ( i
Revised: The backward pass estimates the context distribution at time t give the future observations and actions (i

Original: [Equation]
Revised: [Equation]

Original: e., for k ≤ t ), which is similar to a ﬁltering process.
Revised: e., for k ≥ t ).

Original: The backward pass estimates the context distribution at time t give the future observations and actions (i.e., for k ≥ t ).

Original: Combining both passes allows to compute the context distribution at time t given the whole trajectory.
Revised: Combining both passes allows to compute the context distribution at time t given the whole trajectory.

Original: Gradient of ELBO
Revised: Gradient of ELBO

Original: [Equation]
Revised: [Equation]

Original: where the terms in blue involve differentiating an expectation over Beta distributions, which we compute by adopting the implicit reparameterization (Figurnov et al., 2018; Jankowiak & Obermeyer, 2018).
Revised: where the terms in blue involve differentiating an expectation over Beta distributions, which we compute by adopting the implicit reparameterization (Figurnov et al., 2018; Jankowiak & Obermeyer, 2018).

Original: Considering a general case where x ∼ p φ ( x ) and the cumulative distribution function (CDF) of p φ ( x ) is F φ ( x ) , it has been shown that:
Revised: Considering a general case where x ∼ p φ ( x ) and the cumulative distribution function (CDF) of p φ ( x ) is F φ ( x ) , it has been shown that:

Original: [Equation]
Revised: [Equation]

Original: J USTIFICATION FOR THE VARIATIONAL DISTRIBUTIONS
Revised: J USTIFICATION FOR THE VARIATIONAL DISTRIBUTIONS

Original: Here, we provide both intuitive and empirical justiﬁcations for the choice of variational distributions in (A6).
Revised: Here, we provide both intuitive and empirical justiﬁcations for the choice of variational distributions in (A6).

Original: The mean-ﬁeld approximation is mainly based on the tractability consideration where a reparameterizable variational distribution is required for gradient estimation (Blei et al., 2017).
Revised: The mean-ﬁeld approximation is mainly based on the tractability consideration where a reparameterizable variational distribution is required for gradient estimation (Blei et al., 2017).

Original: Besides, the truncation level is set to K , which reduces an inﬁnite dimensional problem to a ﬁnite one.
Revised: Besides, the truncation level is set to K , which reduces an inﬁnite dimensional problem to a ﬁnite one.

Original: The intuition of choosing the point estimation for q ( ν ) is the following: The q ( µ ) in (A6) induces a variational distribution q ( ρ ) , following ρ jk = µ jk (cid:81) k − 1 i =1 (1 − µ ji ) in (A5).
Revised: The intuition of choosing the point estimation for q ( ν ) is the following: The q ( µ ) in (A6) induces a variational distribution q ( ρ ) , following ρ jk = µ jk (cid:81) k − 1 i =1 (1 − µ ji ) in (A5).

Original: When observing reasonable amount of trajectories, the optimal q ∗ ( ρ ) should center around the ground-truth initial context distribution and the context transition.
Revised: When observing reasonable amount of trajectories, the optimal q ∗ ( ρ ) should center around the ground-truth initial context distribution and the context transition.

Original: The HDP prior in our generative model speciﬁes that ρ j | α, β ∼ DP ( α, β ) (Teh et al., 2006), which means β serves as the expectation of the initial context distribution and each row in the context transition.
Revised: The HDP prior in our generative model speciﬁes that ρ j | α, β ∼ DP ( α, β ) (Teh et al., 2006), which means β serves as the expectation of the initial context distribution and each row in the context transition.

Original: Intuitively, the optimal q ∗ ( β ) , which is induced from q ∗ ( ν ) , should center around the stationary distribution of the context chain.
Revised: Intuitively, the optimal q ∗ ( β ) , which is induced from q ∗ ( ν ) , should center around the stationary distribution of the context chain.

Original: Therefore, each factor q ∗ ( ν k ) in the optimal q ∗ ( ν ) is supposed to be uni-modal, and thus a point estimation could be a reasonable simpliﬁcation.
Revised: Therefore, each factor q ∗ ( ν k ) in the optimal q ∗ ( ν ) is supposed to be uni-modal, and thus a point estimation could be a reasonable simpliﬁcation.

Original: This intuition was supported by our computing β and stationary distributions during training resulting in similar distilled Markov chains.
Revised: This intuition was supported by our computing β and stationary distributions during training resulting in similar distilled Markov chains.

Original: We then conduct a simple numerical study to verify the intuition above, which reveals another problem when using a full Beta variational distribution for ν .
Revised: We then conduct a simple numerical study to verify the intuition above, which reveals another problem when using a full Beta variational distribution for ν .

Original: Consider the following context Markov chain:
Revised: Consider the following context Markov chain:

Original: [Equation]
Revised: [Equation]

Original: a [Figure] We set hyper-parameters in (A5) as γ = α = 1 , κ = 0 , K = 2 and assume that, with sampled trajectories D , the learned optimal q ∗ ( µ ) is given by:
Revised: We set hyper-parameters in (A5) as γ = α = 1 , κ = 0 , K = 2 and assume that, with sampled trajectories D , the learned optimal q ∗ ( µ ) is given by:

Original: [Equation]
Revised: [Equation]

Revised: a

Revised: [Figure]

Original: Recall that ρ j 1 = µ j 1 in (A5).
Revised: Recall that ρ j 1 = µ j 1 in (A5).

Original: The assumed q ∗ ( µ jk ) has the mean at its ground-truth value and the variance 0 .
Revised: The assumed q ∗ ( µ jk ) has the mean at its ground-truth value and the variance 0 .

Original: According to (A7), maximizing ELBO w.r.t q ( ν ) is equivalent to:
Revised: According to (A7), maximizing ELBO w.r.t q ( ν ) is equivalent to:

Original: [Equation]
Revised: [Equation]

Original: where q ( µ ) is ﬁxed to the assumed optima q ∗ ( µ ) .
Revised: where q ( µ ) is ﬁxed to the assumed optima q ∗ ( µ ) .

Original: We investigate two conﬁgurations of q ( ν ) : (1) a point estimation where q ( ν | ˆ ν ) =
Revised: We investigate two conﬁgurations of q ( ν ) : (1) a point estimation where q ( ν | ˆ ν ) =

Original: δ ( ν 1 | ˆ ν 1 ) ; (2) a full
Revised: δ ( ν 1 | ˆ ν 1 ) ; (2) a full

Original: Beta variational distribution where q ( ν | ˆ ν )
Revised: Beta variational distribution where q ( ν | ˆ ν )

Original: =
Revised: =

Original: Beta ( ν 1 | ˆ ν 11 , ˆ ν 12 ) .
Revised: Beta ( ν 1 | ˆ ν 11 , ˆ ν 12 ) .

Original: The results are shown in Figure A4.
Revised: The results are shown in Figure A4.

Original: The optimal point estimation is q ∗ ( ν ) =
Revised: The optimal point estimation is q ∗ ( ν ) =

Original: δ ( ν 1 | 0 . 493) (labeled by the red star in Figure A4(a)), while the optimal Beta distribution in Figure A4(c) is q ∗ ( ν )
Revised: δ ( ν 1 | 0 . 493) (labeled by the red star in Figure A4(a)), while the optimal Beta distribution in Figure A4(c) is q ∗ ( ν )

Original: =
Revised: =

Original: Beta ( ν 1 | 4 . 86 , 4 . 78) .
Revised: Beta ( ν 1 | 4 . 86 , 4 . 78) .

Original: The Beta optimum is uni-modal and its mode 0.505 is very close to the point estimation 0 .
Revised: The Beta optimum is uni-modal and its mode 0.505 is very close to the point estimation 0 .

Original: 493 , which is consistent with our intuition.
Revised: 493 , which is consistent with our intuition.

Original: Comparing Figure A4(b) with A4(d), we observe that the point estimation can provide gradients better suited for optimization while using a Beta variational distribution potentially suffers from vanishing gradients.
Revised: Comparing Figure A4(b) with A4(d), we observe that the point estimation can provide gradients better suited for optimization while using a Beta variational distribution potentially suffers from vanishing gradients.

Original: We observed this phenomenon in our model learning experiments as well.
Revised: We observed this phenomenon in our model learning experiments as well.

Original: Since the transition function p ( s t | s t − 1 , a t − 1 , θ z t ) is modeled by neural networks, it is generally hard to predict any property of the true posterior of θ and choose an appropriate variational distribution.
Revised: Since the transition function p ( s t | s t − 1 , a t − 1 , θ z t ) is modeled by neural networks, it is generally hard to predict any property of the true posterior of θ and choose an appropriate variational distribution.

Original: Bayesian neural network literature attempt to tackle this problem (Welling & Teh, 2011; Blundell et al., 2015; Kingma et al., 2015; Gal & Ghahramani, 2016; Ritter et al., 2018).
Revised: Bayesian neural network literature attempt to tackle this problem (Welling & Teh, 2011; Blundell et al., 2015; Kingma et al., 2015; Gal & Ghahramani, 2016; Ritter et al., 2018).

Original: However, most methods have considerably high computational complexity and it is not trivial to evaluate the quality of generated probabilistic prediction.
Revised: However, most methods have considerably high computational complexity and it is not trivial to evaluate the quality of generated probabilistic prediction.

Original: In this work, we explicitly assume the distribution of s t whose parameters are ﬁtted by neural networks.
Revised: In this work, we explicitly assume the distribution of s t whose parameters are ﬁtted by neural networks.

Original: The transition model is still capable of generating probabilistic prediction with a point estimation of θ .
Revised: The transition model is still capable of generating probabilistic prediction with a point estimation of θ .

Original: This assumption/simpliﬁcation is followed by many model-based RL works.a
Revised: This assumption/simpliﬁcation is followed by many model-based RL works.

Original: C. C ONTEXT D ISTILLATION
Revised: C. C ONTEXT D ISTILLATION

Revised: At every iteration of model learning, we can extract MAP parameter estimates { θ k } Kk =1 and approximated posteriors of ρ 0 and R =

Revised: [ ρ 1 , ..., ρ K ] , which are induced from q ( µ ) .

Revised: Let us also deﬁne

Original:  Inputs: ε distil - distillation threshold; ¯ R - expected context transition matrix; ˆ β - weights of
Revised: a Inputs: ε distil - distillation threshold; ¯ R - expected context transition matrix; ˆ β - weights of

Original: HDP’s base distribution
Revised: HDP’s base distribution

Original: Determine the distillation vector v using one of the following choices:
Revised: Determine the distillation vector v using one of the following choices:

Original: (a) stationary distribution of the chain, v such that v = v ¯ R ; (b) weights of HDP’s base distribution, v = ˆ β .
Revised: (a) stationary distribution of the chain, v such that v = v ¯ R ; (b) weights of HDP’s base distribution, v = ˆ β .

Original: Determine distilled context indexes I 1 and spurious context indexes I 2 as follows:
Revised: Determine distilled context indexes I 1 and spurious context indexes I 2 as follows:

Original: [Equation]
Revised: [Equation]

Original: Compute ˆ R as follows: ifAGENT is MPC then ˆ
Revised: Compute ˆ R as follows: if

Original: [Equation]
Revised: [Equation]

Original: endelse if AGENT is POLICY then
Revised: endelse if AGENT is POLICY then

Original: [Equation]
Revised: [Equation]

Original: end return ˆ R - distilled probability transition matrix.
Revised: end return ˆ R - distilled probability transition matrix.

Original: At every iteration of model learning, we can extract MAP parameter estimates { θ k } Kk =1 and approximated posteriors of ρ 0 and R =

Original: [ ρ 1 , ..., ρ K ] , which are induced from q ( µ ) .

Original: Let us also deﬁne the expected context initial distribution ¯ ρ [-
-] = E q ( µ ) [ ρ 0 ] and the expected context transition matrix ¯
Revised: the expected context initial distribution ¯ ρ [-
-]0 = E q ( µ ) [ ρ 0 ] and the expected context transition matrix ¯

Original: R = E q ( µ ) [ R ] .
Revised: R = E q ( µ ) [ R ] .

Original: These MAP estimates are used during training as well as testing for sampling the values z .
Revised: These MAP estimates are used during training as well as testing for sampling the values z .

Original: Hence distilling ¯ R has an effect on training as well as testing.
Revised: Hence distilling ¯ R has an effect on training as well as testing.

Original: Our distillation criterion is based on the values of the stationary distribution of the context Markov chain.
Revised: Our distillation criterion is based on the values of the stationary distribution of the context Markov chain.

Original: Recall that one can compute the stationary distribution ρ ∞ by solving ρ ∞ = ρ ∞ ¯ R .
Revised: Recall that one can compute the stationary distribution ρ ∞ by solving ρ ∞ = ρ ∞ ¯ R .

Original: Now the meaningful context indexes I 

Revised: Now the meaningful context indexes I

Original: = { i | ρ ∞ i ≥ ε distil } and spurious context indexes I 2 = { i | ρ ∞
Revised: = { i | ρ ∞ i ≥ ε distil } and spurious context indexes I 2 = { i | ρ ∞

Original: i < ε distil } can be chosen using a distillation threshold ε distil .
Revised: i < ε distil } can be chosen using a distillation threshold ε distil .

Original: Then, we distill the learned contexts by simply discarding ˆ θ I 2 .
Revised: Then, we distill the learned contexts by simply discarding ˆ θ I 2 .

Original: Meanwhile, the context Markov chain also needs to be reduced.
Revised: Meanwhile, the context Markov chain also needs to be reduced.

Original: For ¯ ρ 0 , we gather those dimensions indexed by I 1 into a new vector ˆ ρ 0 and re-normalize ˆ ρ 0 .
Revised: For ¯ ρ 0 , we gather those dimensions indexed by I 1 into a new vector ˆ ρ 0 and re-normalize ˆ ρ 0 .

Original: In addition, ¯ R can be reduced to ˆ R following the Theorem 1 in the main text.
Revised: In addition, ¯ R can be reduced to ˆ R following the Theorem 1 in the main text.

Original: Perhaps, a less rigorous, but deﬁnitely a simpler approach is choosing the index sets I 

Revised: Perhaps, a less rigorous, but deﬁnitely a simpler approach is choosing the index sets I

Original: and I 2 using ˆ β — a MAP estimation of β computed using ˆ ν .
Revised: and I 2 using ˆ β — a MAP estimation of β computed using ˆ ν .

Original: Since optimizing the KL ( q ( µ ) || p ( µ | ˆ ν ))] term in
Revised: Since optimizing the KL ( q ( µ ) || p ( µ | ˆ ν ))] term in

Original: ELBO is essentially driving the posterior of ρ 0: K toward ˆ β .
Revised: ELBO is essentially driving the posterior of ρ 0: K toward ˆ β .

Original: Therefore, ˆ β can be seen as a ‘summary’ distribution over contexts and we can consider the k -th context as a redundancy if ˆ β k is small.
Revised: Therefore, ˆ β can be seen as a ‘summary’ distribution over contexts and we can consider the k -th context as a redundancy if ˆ β k is small.

Original: It is not clear if ˆ β has a direct relation to the stationary distribution of the Markov chain with the transition probability ˆ R .
Revised: It is not clear if ˆ β has a direct relation to the stationary distribution of the Markov chain with the transition probability ˆ R .

Original: However, we have observed that the magnitudes of the entries of ˆ β and p ∞ are correlated.
Revised: However, we have observed that the magnitudes of the entries of ˆ β and p ∞ are correlated.

Original: Hence, in order to avoid computing an eigenvalue decomposition at every context estimation one can employ distillation using ˆ β .
Revised: Hence, in order to avoid computing an eigenvalue decomposition at every context estimation one can employ distillation using ˆ β .

Original: Both approaches are summarized in Algorithm A2.
Revised: Both approaches are summarized in Algorithm A2.

Original: For policy optimization, we actually need to keep the number of contexts constant as dealing with changing state-belief space can be challenging during training.
Revised: For policy optimization, we actually need to keep the number of contexts constant as dealing with changing state-belief space can be challenging during training.

Original: 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.
Revised: 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.

Original: We can still remove the spurious contexts after training both from the model and the policy.
Revised: We can still remove the spurious contexts after training both from the model and the policy.

Original: a
Revised: a

Original: [Figure]
Revised: [Figure]

Original: This procedure is gradient-free, which beneﬁts lowerdimensional settings, but could be less efﬁcient in higher dimensional environments.
Revised: This procedure is gradient-free, which beneﬁts lowerdimensional settings, but could be less efﬁcient in higher dimensional environments.

Original: It is also known to be more efﬁcient than random shooting methods (Chua et al., 2018).
Revised: It is also known to be more efﬁcient than random shooting methods (Chua et al., 2018).

Original: The idea of this approach is quite simple and sketched in Algorithm A3.
Revised: The idea of this approach is quite simple and sketched in Algorithm A3.

Original: At the state s t with an action plan { a k } t + H k = t , we sample plan updates { δ ik } t
Revised: At the state s t with an action plan { a k } t + H k = t , we sample plan updates { δ ik } t

Original: + H k = t .
Revised: + H k = t .

Original: We then roll-out trajectories and compute the average returns for each plan { a k + δ ik } t + H k = t .
Revised: We then roll-out trajectories and compute the average returns for each plan { a k + δ ik } t + H k = t .

Original: We pick N elite best performing action plans {{ a k
Revised: We pick N elite best performing action plans {{ a k

Original: + δ i j k } t
Revised: + δ i j k } t

Original: + H k = t } N elite j =1 , compute the empirical mean ˜ µ k and variance ˜ Σ
Revised: + H k = t } N elite j =1 , compute the empirical mean ˜ µ k and variance ˜ Σ

Original: k of the elite plan updates δ i j k , and then compute the updates on the action plan distribution as follows:
Revised: k of the elite plan updates δ i j k , and then compute the updates on the action plan distribution as follows:

Original: [Equation]
Revised: [Equation]

Original: where l r is the learning rate.
Revised: where l r is the learning rate.

Original: Algorithm A3:
Revised: Algorithm A3:

Original: MPC based on CEM
Revised: MPC based on CEM

Original: Inputs: { a k } t + H k = t , { µ k } t + H k = t , { Σ k } t + H k = t , N epochs , N elite , N traces , N pops , l r , s , H for j = 0 , . .
Revised: Inputs: { a k } t + H k = t , { µ k } t + H k = t , { Σ k } t + H k = t , N epochs , N elite , N traces , N pops , l r , s , H for j = 0 , . .

Original: . , N epochs do
Revised: . , N epochs do

Original: Sample action plan updates {{ δ ik } t
Revised: Sample action plan updates {{ δ ik } t

Original: + H k = t } N pops i
Revised: + H k = t } N pops i

Original: =1 , where δ ik ∼ N ( µ k , Σ k ) ;
Revised: =1 , where δ ik ∼ N ( µ k , Σ k ) ;

Original: Roll-out N traces for each update plan;
Revised: Roll-out N traces for each update plan;

Original: Compute the returns 1 /N traces (cid:80) N traces p = (cid:80) t
Revised: Compute the returns 1 /N traces (cid:80) N traces p = (cid:80) t

Original: + H k = t r ( s pk , a k + δ ik ) with s pt =
Revised: + H k = t r ( s pk , a k + δ ik ) with s pt =

Original: s for all p ;
Revised: s for all p ;

Original: Pick N elite best performing 
Revised: Pick N elite best performing action

Original: C. S OFT -A CTOR C RITIC
Revised: C. S OFT -A CTOR C RITIC

Original: We reproduce the summary of the soft-actor critic algorithm by Achiam (2018), which we found very accessible.
Revised: We reproduce the summary of the soft-actor critic algorithm by Achiam (2018), which we found very accessible.

Original: The soft-actor critic algorithm aims at solving a modiﬁed RL problem with an entropy-regularized objective:
Revised: The soft-actor critic algorithm aims at solving a modiﬁed RL problem with an entropy-regularized objective:

Original: [Equation]
Revised: [Equation]

Original: where H ( P ) =
Revised: where H ( P ) =

Original: − E x ∼ P [log( P ( x )] and α is called the temperature parameter.
Revised: − E x ∼ P [log( P ( x )] and α is called the temperature parameter.

Original: The entropy regularization modiﬁes the Bellman equation for this problem as follows:
Revised: The entropy regularization modiﬁes the Bellman equation for this problem as follows:

Original: [Equation]
Revised: [Equation]

Original: The algorithm largely follows the standard actor-critic framework for updating value functions and policy, with a few notable changes.
Revised: The algorithm largely follows the standard actor-critic framework for updating value functions and policy, with a few notable changes.

Original: First, two Q functions are used in order to avoid overestimation of the value functions .
Revised: First, two Q functions are used in order to avoid overestimation of the value functions .

Original: In particular, the loss for value learning is as follows:
Revised: In particular, the loss for value learning is as follows:

Original: [Equation]
Revised: [Equation]

Original: For policy updates the reparameterization trick is used allowing for differentiation of the policy.
Revised: For policy updates the reparameterization trick is used allowing for differentiation of the policy.

Original: Namely, the policy loss function is as follows:
Revised: Namely, the policy loss function is as follows:

Original: [Equation]
Revised: [Equation]

Original: a
Revised: a

Revised: Inputs: N epochs - number of epochs, N upd - number of gradient updates per epochs,

Revised: N target − freq - target value function update frequency, N samples - number of steps per epoch, l w learning rates

Revised: N total − upd = 0 .

Original: Initialize parameters ψ , φ i , φ target , i = φ i ; for j = 
Revised: Initialize parameters ψ , φ i , φ target , i = φ i ; for j = 0 , . . .

Revised: , N epochs do

Revised: Sample N samples steps from the environment with a ∼ π ψ ( ·| s ) resulting in the buffer update

Original:  N do
Revised: D new = ( s i , a i , s (cid:48) i , r i , d i ) } N samples i = ;

Revised: Set D =

Revised: D new ∪ D ;

Revised: Sample a batch B from the buffer D ; for k = 0 , . . .

Revised: , N upd do N total

Revised: N

Revised: ; Update parameters of the value functions

Original: Sample steps from the environment with a ∼ π ψ ( ·| resulting in the buffer update do if mod ( N total − upd , N target − freq ) {+= 
 then Update parameters of the target value function φ+}
Revised: Update parameters of the policy ψ ψ ) ; if mod ( N total − upd , N target − freq ) {+= 
 then Update parameters of the target value function φ+}

Revised: , i ← w φ

Original: [-= 
-]
Revised: [-= 
-]+

Original: then end end endreturn 
Revised: 1 − w ) φ i ; end end endreturn π ψ

Original: D E XPERIMENT D ETAILS D. L EARNING ALGORITHMS
Revised: D E XPERIMENT D ETAILS D. L EARNING ALGORITHMS

Original: Model learning We implemented the model learning using the package Pyro (Bingham et al., 2018), which is designed for efﬁcient probabilistic programming.
Revised: Model learning We implemented the model learning using the package Pyro (Bingham et al., 2018), which is designed for efﬁcient probabilistic programming.

Original: Pyro allows for automatic differentiation, i.e., we do not need to explicitly implement message passing and reparametrized gradients for the
Revised: Pyro allows for automatic differentiation, i.e., we do not need to explicitly implement message passing and reparametrized gradients for the

Original: ELBO gradient computation.
Revised: ELBO gradient computation.

Original: We still need, however, a forward message pass to compute the belief estimate, e.g., to perform ﬁltering on the variable z t when needed.
Revised: We still need, however, a forward message pass to compute the belief estimate, e.g., to perform ﬁltering on the variable z t when needed.

Original: PPO with an RNN model We modiﬁed an implementation of PPO by Kostrikov (2018) to account for our belief model.
Revised: PPO with an RNN model We modiﬁed an implementation of PPO by Kostrikov (2018) to account for our belief model.

Original: In our implementation, the RNN with a hidden state h at time t is taking the inputs h t − 1 , s t − 1 , a t − 1 , while producing the output h t .
Revised: In our implementation, the RNN with a hidden state h at time t is taking the inputs h t − 1 , s t − 1 , a t − 1 , while producing the output h t .

Original: What is left is to project the hidden state onto the belief space using a decoder, which we have chosen as (cid:98) b t = softmax( W h t ) , where the length of the vector (cid:98) b t is equal to the number of contexts.
Revised: What is left is to project the hidden state onto the belief space using a decoder, which we have chosen as (cid:98) b t = softmax( W h t ) , where the length of the vector (cid:98) b t is equal to the number of contexts.

Original: The architecture is depicted in Figure A6(a).
Revised: The architecture is depicted in Figure A6(a).

Original: Note that one can see the RNN and the decoder architecture as a model for the sufﬁcient information state for the POMDP.
Revised: Note that one can see the RNN and the decoder architecture as a model for the sufﬁcient information state for the POMDP.

Original: We have experimented with different architectures, e.g., projecting to a larger belief space to account for spurious contexts, removing the decoder altogether etc.
Revised: We have experimented with different architectures, e.g., projecting to a larger belief space to account for spurious contexts, removing the decoder altogether etc.

Original: These architectures, however, did not yield reasonable results.
Revised: These architectures, however, did not yield reasonable results.

Original: GPMM.
Revised: GPMM.

Original: We took the implementation by Xu et al.
Revised: We took the implementation by Xu et al.

Original: (2020), which is able to swing-up the pole attached to the cart and adapt to environments with different parameters.
Revised: (2020), which is able to swing-up the pole attached to the cart and adapt to environments with different parameters.

Original: SAC.
Revised: SAC.

Original: We based our implementation largely on (Tandon, 2018) with some inspiration from (Yarats & Kostrikov, 2020).
Revised: We based our implementation largely on (Tandon, 2018) with some inspiration from (Yarats & Kostrikov, 2020).

Original: We use two architectures: full information policy (see Figure A6(c)) and model-based policy (see Figure A6(b)).
Revised: We use two architectures: full information policy (see Figure A6(c)) and model-based policy (see Figure A6(b)).

Original: The full information policy is using one hot encoded true context and is, therefore, used as a reference for the best case performance only.
Revised: The full information policy is using one hot encoded true context and is, therefore, used as a reference for the best case performance only.

Original: CEM-MPC.
Revised: CEM-MPC.

Original: We implemented the algorithm from scratch in PyTorch.
Revised: We implemented the algorithm from scratch in PyTorch.

Original: D. E NVIRONMENTS AND THEIR M ODELS
Revised: D. E NVIRONMENTS AND THEIR M ODELS

Original: Cart-Pole Swing Up
Revised: Cart-Pole Swing Up

Original: We largely followed the setting introduced by Xu et al.
Revised: We largely followed the setting introduced by Xu et al.

Original: (2020), that is we set the maximum force magnitude to 20 , time interval to 0 . 04 , and time horizon to 100 .
Revised: (2020), that is we set the maximum force magnitude to 20 , time interval to 0 . 04 , and time horizon to 100 .

Original: We took a
Revised: We took a

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Revised: implementation by Lovatto (2019) and modiﬁed it to ﬁt our context MDP setting

Original: modiﬁed 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 ( θ )  angular velocity ( ˙ θ )
Revised: 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 ( θ ) and angular velocity ( ˙ θ )

Original: For GPMM we replaced θ with sin( θ ) and cos( θ ) as was done by Xu et al  set the reward function to cos( θ ) where θ is the deviation from the top position of model is the successor is a multi perceptron predicting
Revised: For GPMM we replaced θ with sin( θ ) and cos( θ ) as was done by Xu et al We set the reward function to cos( θ ) where θ is the deviation from the top position Our transition model predicts the mean the the its as in of network model for (cid:48)

Revised: +

Revised: ( s a )

Original:  Drone Take-off The drone environment (Panerati et al., 2021) has 12 states: position in the world frame ( p x , p y and p z ), yaw, pitch, roll angles ( ψ , θ and φ ), velocities in the world frame ( v x , v y and v z ) and angular velocities in the world frame ( ω x , ω y and ω z ).
Revised: (cid:48) is the successor and is a multi-layer perceptron predicting (cid:48) The variance in the model is a depicted Drone Take-off The drone environment (Panerati et al., 2021) has 12 states: position in the world frame ( p x , p y and p z ), yaw, pitch, roll angles ( ψ , θ and φ ), velocities in the world frame ( v x , v y and v z ) and angular velocities in the world frame ( ω x , ω y and ω z ).

Original: 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.
Revised: 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.

Original: Note that we can estimate spacial positions, roll, pitch and yaw a
Revised: Note that we can estimate spacial positions, roll, pitch and yaw a

Original: angles given position and angular velocities using crude but effective formulae:
Revised: angles given position and angular velocities using crude but effective formulae:

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where the formula for angular velocities can be found, for example, in Hover & Triantafyllou (2009)
Revised: where the formula for angular velocities can be found, for example, in Hover & Triantafyllou (2009)

Original: [Table]
Revised: [Table]

Original: We will refer to the matrix R as the rotation matrix with a slight abuse of notation.
Revised: We will refer to the matrix R as the rotation matrix with a slight abuse of notation.

Original: We also choose special features for the MLP: angular velocities, velocities, sines and cosines of the Euler angles ( E ), actions and actions squared.
Revised: We also choose special features for the MLP: angular velocities, velocities, sines and cosines of the Euler angles ( E ), actions and actions squared.

Original: Using these expression we impose the structure on the neural network depicted in Figure A7(c).
Revised: Using these expression we impose the structure on the neural network depicted in Figure A7(c).

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Left turn on the Intersection in Highway Environment We take the environment by Leurent (2018), but use the modiﬁcations made by Xu et al.
Revised: Left turn on the Intersection in Highway Environment We take the environment by Leurent (2018), but use the modiﬁcations made by Xu et al.

Original: (2020) including the overall reward function structure.
Revised: (2020) including the overall reward function structure.

Original: We, however, do not penalize the collisions.
Revised: We, however, do not penalize the collisions.

Original: and we increase the episode time from 40 to a
Revised: and we increase the episode time from 40 to a

Original: [Table]
Revised: [Table]

Original: ZZ 1 Z
Revised: Z 1 Z

Original: Z 3 Zp ( z 0 ) p ( z ) 81.0 0.0 9.1 9.9 0.1 5.9 1.5 58.9 32.4 1.
Revised: Z 3 Zp ( z 0 ) p ( z ) 81.0 0.0 9.1 9.9 0.1 5.9 1.5 58.9 32.4 1.

Original: P r o b a b ili t y ( % ) (c) MLE
Revised: P r o b a b ili t y ( % ) (c) MLE

Original: Z 0 Z 1 Z 2 Z 3 Z 4 Z
Revised: Z 0 Z 1 Z 2 Z 3 Z 4 Z

Original: Z 1 Z
Revised: Z 1 Z

Original: Z 3 Zp ( z 0 ) p ( z )
Revised: Z 3 Zp ( z 0 ) p ( z )

Original: P r o b a b ili t y ( % ) (d) HDP w distillation
Revised: P r o b a b ili t y ( % ) [Figure]

Original: 100 time steps.
Revised: 100 time steps.

Original: We again predicted the difference between current and next steps for the mean, and used the simpliﬁed model for the position, i.e., ∆ p x ≈ dtv x , ∆ p y
Revised: We again predicted the difference between current and next steps for the mean, and used the simpliﬁed model for the position, i.e., ∆ p x ≈ dtv x , ∆ p y

Original: ≈ dtv y for both social and ego vehicles.
Revised: ≈ dtv y for both social and ego vehicles.

Original: D. H YPER - PARAMETERS
Revised: D. H YPER - PARAMETERS

Original: All the hyper-parameters are presented in Tables A1, A2 and A3.
Revised: All the hyper-parameters are presented in Tables A1, A2 and A3.

Original: For model learning experiments we used 500 trajectory roll-outs and 500 epochs for optimization.
Revised: For model learning experiments we used 500 trajectory roll-outs and 500 epochs for optimization.

Original: In the cart-pole environment we used the higher learning rate for hard failure experiments when χ < 0 and used the lower learning rate for the soft failure experiments χ > 0 .
Revised: In the cart-pole environment we used the higher learning rate for hard failure experiments when χ < 0 and used the lower learning rate for the soft failure experiments χ > 0 .

Original: We use the weight decay to avoid gradient explosion in the value functions and the policies.
Revised: We use the weight decay to avoid gradient explosion in the value functions and the policies.

Original: Similarly, Clipped Adam optimizer (available in Pyro) was used to avoid gradient explosion in model learning.
Revised: Similarly, Clipped Adam optimizer (available in Pyro) was used to avoid gradient explosion in model learning.

Original: E A DDITIONAL E XPERIMENTS E. HDP IS AN EFFECTIVE PRIOR FOR LEARNING AN ACCURATE AND INTERPRETABLE MODEL
Revised: E A DDITIONAL E XPERIMENTS E. HDP IS AN EFFECTIVE PRIOR FOR LEARNING AN ACCURATE AND INTERPRETABLE MODEL

Original: We plot the time courses of the context evolution and the ground truth context evolution in Figure A9.
Revised: We plot the time courses of the context evolution and the ground truth context evolution in Figure A9.

Original: 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.
Revised: 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.

Original: The difference was the choice of the distillation threshold, which had to be signiﬁcantly higher for the DP prior.
Revised: The difference was the choice of the distillation threshold, which had to be signiﬁcantly higher for the DP prior.

Original: This experiment indicates that DP prior can be a good tool for modeling context transitions, but HDP provides sharper model ﬁt and a more interpretable model.
Revised: This experiment indicates that DP prior can be a good tool for modeling context transitions, but HDP provides sharper model ﬁt and a more interpretable model.

Original: For completeness, we performed the same experiment, but with a different seed and setting ε distil = 0 . 1 during training.
Revised: For completeness, we performed the same experiment, but with a different seed and setting ε distil = 0 . 1 during training.

Original: We plot the results in Figures A10 and A11.
Revised: We plot the results in Figures A10 and A11.

Original: In this case, all models (including the MLE method) coupled with distillation provided an accurate estimate of the context evolution.
Revised: In this case, all models (including the MLE method) coupled with distillation provided an accurate estimate of the context evolution.

Original: This suggests that the optimization proﬁle for the MLE method has many local minima (peaks and troughs in ELBO), which we can be trapped in given an unlucky seed.
Revised: This suggests that the optimization proﬁle for the MLE method has many local minima (peaks and troughs in ELBO), which we can be trapped in given an unlucky seed.

Original: E. D ISTILLATION ACTS AS A REGULARIZER

Original: a
Revised: a

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: [Figure]
Revised: [Figure]

Original: P r o b a b ili t y ( % ) 
Revised: P r o b a b ili t y ( % ) (

Original: [Figure]
Revised: [Figure]

Revised: Dirichlet

Original: Z 0 Z 1 Z 2 Z 3 Z 4 Z
Revised: Z 0 Z 1 Z 2 Z 3 Z 4 Z

Original: Z
Revised: Z

Original: Z
Revised: Z

Original: Z
Revised: Z

Original: Zp ( z 0 ) p ( z )[-0.
 [Figure] [Figure] [Figure]-]
Revised: Zp ( z 0 ) p ( z )[-0.
 [Figure] [Figure] [Figure]-]

Original: [Figure]

Original: [Figure] [Figure] [Figure] P r o b a b ili t y 
Revised: P r o b a b ili t y ( % ) (

Original: [Figure]
Revised: [Figure]

Revised: MLE

Original: [Figure]
Revised: [Figure]

Original: a
Revised: a

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: (cid:101) C| = 8 .
Revised: C| = 8 .

Original: [Table]
Revised: [Table]

Revised: E. D ISTILLATION ACTS AS A REGULARIZER

Original: After the ﬁrst experiment
Revised: After the ﬁrst experiment

Original: [Table]
Revised: [Table]

Original: we noticed that the context Z 2 has a low probability mass in stationarity, but a high probability of self-transition (see Figure A8 - reproduction of Figure 2 in the main text)
Revised: we noticed that the context Z 2 has a low probability mass in stationarity, but a high probability of self-transition (see Figure A8 - reproduction of Figure 2 in the main text)

Original: [Table]
Revised: [Table]

Original: This suggest that spurious transitions can happen, while highly unlikely.
Revised: This suggest that spurious transitions can happen [Table] while highly unlikely.

Original: We speculate that the learning algorithm tries to ﬁt the uncertainty in the model (e
Revised: We speculate that the learning algorithm tries to ﬁt the uncertainty in the model (e

Original: [Table]
Revised: [Table]

Original: g., due to unseen data) to one context.
Revised: g., due to unseen data) to one context.

Original: This can lead to over-ﬁtting and unwanted side-effects.
Revised: This can lead to over-ﬁtting and unwanted side-effects.

Original: Results in Figure A8 (reproduction of Figurein the main text) suggest that distillation during training can act as a regularizer when we used a high enough threshold ε distil = 0 .
Revised: Results in Figure A8 (reproduction of Figurein the main text) suggest that distillation during training can act as a regularizer when we used a high enough threshold ε distil = 0 .

Original: We further validate our ﬁndings by changing the context number upper bound | (cid:101) C| between 4 and 20 .
Revised: We further validate our ﬁndings by changing the context number upper bound | (cid:101) C| between 4 and 20 .

Original: In particular, we proceed by presenting the results for varying |
Revised: In particular, we proceed by presenting the results for varying |

Original: (cid:101) C| (taking values 4 , 5 , 6 , 8 , 10 , 20 and 100 ) and the distillation threshold ε distil (taking values 0 , 0 . 01 , and 0 . 1 ).
Revised: (cid:101) C| (taking values 4 , 5 , 6 , 8 , 10 , 20 and 100 ) and the distillation threshold ε distil (taking values 0 , 0 . 01 , and 0 . 1 ).

Original: Note that we distill during training and we refer to the transition matrix for the distilled Markov chain as the distilled transition matrix.
Revised: Note that we distill during training and we refer to the transition matrix for the distilled Markov chain as the distilled transition matrix.

Original: First, consider the results in Figure A12, where we plot the learned MAP estimates of the transition matrices with | (cid:101) C| = 8 .
Revised: First, consider the results in Figure A12, where we plot the learned MAP estimates of the transition matrices with | (cid:101) C| = 8 .

Original: Note that using distillation with both thresholds prevents overﬁtting to one context.
Revised: Note that using distillation with both thresholds prevents overﬁtting to one context.

Original: One could argue that instead of context distillation during training one could simply use distillation after training.
Revised: One could argue that instead of context distillation during training one could simply use distillation after training.

Original: This approach, however, can lead to emergence of a spurious context with a large probability of a self-transition raising the possibility of inaccurate model predictions.
Revised: This approach, however, can lead to emergence of a spurious context with a large probability of a self-transition raising the possibility of inaccurate model predictions.

Original: Furthermore, the large number of spurious contexts can lead to a large probability mass concentrated in one of them in stationarity.
Revised: Furthermore, the large number of spurious contexts can lead to a large probability mass concentrated in one of them in stationarity.

Original: Indeed, consider Table A4, where we plot the context with third largest probability mass.
Revised: Indeed, consider Table A4, where we plot the context with third largest probability mass.

Original: In particular, for | (cid:101) C|
Revised: In particular, for | (cid:101) C|

Original: = 20 the probability mass values for this context are larger than 0 . 01 .
Revised: = 20 the probability mass values for this context are larger than 0 . 01 .

Original: This indicates a small but not insigniﬁcant possibility of a transition to this context, if the distillation does not remove it.
Revised: This indicates a small but not insigniﬁcant possibility of a transition to this context, if the distillation does not remove it.

Original: We verify that the distilled transition matrices for various cardinalities | (cid:101) C| are close to each other.
Revised: We verify that the distilled transition matrices for various cardinalities | (cid:101) C| are close to each other.

Original: Let ˆ R K denote the estimated transition matrix for | (cid:101)
Revised: Let ˆ R K denote the estimated transition matrix for | (cid:101)

Original: C|
Revised: C|

Original: =
Revised: =

Original: K with only two contexts chosen a posteriori.
Revised: K with only two contexts chosen a posteriori.

Original: We use the following metric for comparison
Revised: We use the following metric for comparison

Original: [Equation]
Revised: [Equation]

Original: That is, we compare all the transition matrices to the case of | (cid:101) C|
Revised: That is, we compare all the transition matrices to the case of | (cid:101) C|

Original: = 5 .
Revised: = 5 .

Original: In Table A5 we present the results, which indicate that the estimated transition matrices are quite close to each other.
Revised: In Table A5 we present the results, which indicate that the estimated transition matrices are quite close to each other.

Original: Furthermore, the distillation during training helps to recover the true context regardless of the upper bound | (cid:101) C| .
Revised: Furthermore, the distillation during training helps to recover the true context regardless of the upper bound | (cid:101) C| .

Original: To summarize, our experiments suggest that the context set cardinality | (cid:101) C| can be conﬁdently overestimated and the context set can be reconstructed using our distillation procedure.

Original: a
Revised: a

Original: [Table]
Revised: [Table]

Original: We vary
Revised: [Table] To summarize [Table] our experiments suggest that the context set cardinality | (cid:101) C| can be conﬁdently overestimated and the context set can be reconstructed using our distillation procedure

Original: [Table]
Revised: [Table]

Original:  C ONTEXT CARDINALITY VS MODEL COMPLEXITY
Revised: E. C ONTEXT CARDINALITY VS MODEL COMPLEXITY

Original: [Figure]
Revised: [Figure]

Original: various distillation thresholds where
Revised: 4 Z Z 6 Z 7 Z Z 9 Z 10 Z(c) (cid:15) distill = 0 . where

Original: C = {{− 1 } , {− 0 .
Revised: C = {{− 1 } , {− 0 .

Original: We deduce that the learned context sets: {+(cid:98) C 
+}
Revised: We deduce that the learned context sets: {+(cid:98) C 
+}

Revised: = {{− 1 } , {− 0 .

Revised: 1 } , { 0 . 5 } , { 1 }} , (cid:98) C

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Revised: 5 } , { 1 }} , (cid:98)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: = 
Revised: = { [Equation] [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: Next we demonstrate that our algorithm can merge similar context automatically using an appropriate distillation threshold during training.
Revised: Next we demonstrate that our algorithm can merge similar context automatically using an appropriate distillation threshold during training.

Original: We increased the force magnitude by two-fold in order to create a large number of distinct contexts and chosen the context set as C = {{− 1 } , {− 0 .
Revised: We increased the force magnitude by two-fold in order to create a large number of distinct contexts and chosen the context set as C = {{− 1 } , {− 0 .

Original: 5 } , { 0 . 5 } , { 0 . 05 } , { 0 . 1 } , { 0 . 5 } , { 1 }} and vary the distillation threshold (cid:15) distill setting it to 0 , 0 . 05 and 0 .
Revised: 5 } , { 0 . 5 } , { 0 . 05 } , { 0 . 1 } , { 0 . 5 } , { 1 }} and vary the distillation threshold (cid:15) distill setting it to 0 , 0 . 05 and 0 .

Original: 1 during training.
Revised: 1 during training.

Original: In order to get the estimated context sets we distilled one more time after training with (cid:15) distill = 0 . 02 .
Revised: In order to get the estimated context sets we distilled one more time after training with (cid:15) distill = 0 . 02 .

Original: We set a sufﬁciently high context cardinality estimate K = 12 .
Revised: We set a sufﬁciently high context cardinality estimate K = 12 .

Original: Results in Figure A13 suggest that the learned contexts sets for various distillation threshold are as follows:
Revised: Results in Figure A13 suggest that the learned contexts sets for various distillation threshold are as follows:

Original: [Equation]
Revised: [Equation]

Revised: C 0 .

Revised: 1 = {{− 1 , − 0 .

Original: This indicates the ability to change the model structure using the distillation threshold.
Revised: This indicates the ability to change the model structure using the distillation threshold.

Original: Furthermore, the trade-off between continuous and discrete dynamics is decided automatically using the distillation threshold.
Revised: Furthermore, the trade-off between continuous and discrete dynamics is decided automatically using the distillation threshold.

Original: E. B REAKING THE ASSUMPTIONS IN H IDDEN M ARKOV M ODELS
Revised: E. B REAKING THE ASSUMPTIONS IN H IDDEN M ARKOV M ODELS

Original: One of the major assumptions in our framework is the Markovian nature of the switches, which can limit its applicability.
Revised: One of the major assumptions in our framework is the Markovian nature of the switches, which can limit its applicability.

Original: We conduct two simple experiments to assess if these assumptions can be broken without detrimental effects.
Revised: We conduct two simple experiments to assess if these assumptions can be broken without detrimental effects.

Original: First, we model the context transitions as a process, where the next context depends on the previous context rather than the current context.
Revised: First, we model the context transitions as a process, where the next context depends on the previous context rather than the current context.

Original: This process is clearly  non-Markovian.
Revised: This process is clearly a non-Markovian.

Original: However, results in Figure A14 indicate that our model successfully predicts the correct contexts in the context realization.
Revised: However, results in Figure A14 indicate that our model successfully predicts the correct contexts in the context realization.

Original: Second, we make the context state dependent.
Revised: Second, we make the context state dependent.

Original: In particular a we have: z = 0 , if pos ∈ [0 , 0 . 1) , [ − 0 .
Revised: In particular, we have: c = 0 , if pos ∈ [0 , 0 . 1) , [ − 0 .

Original: 1 , − 0 .
Revised: 1 , − 0 .

Original: 2) , [0 . 2 , 0 . 3) , [ − 0 .
Revised: 2) , [0 . 2 , 0 . 3) , [ − 0 .

Original: 3 , − 0 .
Revised: 3 , − 0 .

Original: 4) , · · · z = 1 , if pos ∈ [0 , − 0 . 1) , [0 . 1 , 0 . 2) , [ − 0 .
Revised: 4) , · · · c = 1 , if pos ∈ [0 , − 0 . 1) , [0 . 1 , 0 . 2) , [ − 0 .

Original: 2 , − 0 .
Revised: 2 , − 0 .

Original: 3) , [0 . 3 , 0 . 4) , · · · , where pos is the position of the cart.
Revised: 3) , [0 . 3 , 0 . 4) , · · · , where pos is the position of the cart.

Original: Again our model handles this case as the results in Figure A suggest.
Revised: Again our model handles this case as the results in Figure A suggest.

Original: So how our Markovian model handles predictions for non-Markovian models?
Revised: So how our Markovian model handles predictions for non-Markovian models?

Original: In both cases, the key is comparing context estimations/predictions rather than context models, which are almost surely incorrect.
Revised: In both cases, the key is comparing context estimations/predictions rather than context models, which are almost surely incorrect.

Original: However, the predictive power of our model is still preserved since it relies heavily on the observed state history.
Revised: However, the predictive power of our model is still preserved since it relies heavily on the observed state history.

Original: This allowed us to correctly estimate the contexts in these experiments.
Revised: This allowed us to correctly estimate the contexts in these experiments.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: E. L EARNING NEW CONTEXTS
Revised: E. L EARNING NEW CONTEXTS