Original: However, our model is ﬂexible enough to add new (unseen) contexts without learning the state transition model from scratch.
Revised: However, our model is ﬂexible enough to add new (unseen) contexts without learning the state transition model from scratch.

Original: We design the following experimental procedure to demonstrate this a ability: (1) train the model on dataset D 1 which contains two contexts ( C 0 and C 1 ); (2) reset free parameters in the variational distribution q ( ν | ˆ ν ) and q ( µ | ˆ µ ) while preserving q ( θ | ˆ θ ) ; (3) re-train the model on dataset D 2 which contains the two original contexts and two new contexts ( C 2 and
Revised: We design the following experimental procedure to demonstrate this a ability: (1) train the model on dataset D 1 which contains two contexts ( C 0 and C 1 ); (2) reset free parameters in the variational distribution q ( ν | ˆ ν ) and q ( µ | ˆ µ ) while preserving q ( θ | ˆ θ ) ; (3) re-train the model on dataset D 2 which contains the two original contexts and two new contexts ( C 2 and

Original: C 3 ).
Revised: C 3 ).

Original: We present the training results on the sets D 1 and D 2 in Figure A16 and A17, respectively.
Revised: We present the training results on the sets D 1 and D 2 in Figure A16 and A17, respectively.

Original: These results suggest that our model is able to learn new contexts ( C 2 and C 3 in Figure A17), while preserving the original contexts ( C 0 and C 1 in both Figures A16 and A17).
Revised: These results suggest that our model is able to learn new contexts ( C 2 and C 3 in Figure A17), while preserving the original contexts ( C 0 and C 1 in both Figures A16 and A17).

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Z Z
Revised: Z Z

Original: Z Z
Revised: Z Z

Original: Zp ( z 0 )
Revised: Zp ( z 0 )

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: E.
Revised: E.

Original: C OMPARING TO POMDP AND C
Revised: C OMPARING TO POMDP AND C

Original: -MDP METHODS
Revised: -MDP METHODS

Original: We chosen the context set as C = {− 1 , χ } for χ ∈ {− 1 , 0 . 1 , 0 . 3 , 0 . 5 } .
Revised: We chosen the context set as C = {− 1 , χ } for χ ∈ {− 1 , 0 . 1 , 0 . 3 , 0 . 5 } .

Original: However, for χ > 0 , we again increased the force magnitude by two-fold.
Revised: However, for χ > 0 , we again increased the force magnitude by two-fold.

Original: For SAC-based algorithms we present the statistics for 50 episodes and 3 seeds.
Revised: For SAC-based algorithms we present the statistics for 50 episodes and 3 seeds.

Original: To avoid unfair comparison to POMDP and continual RL methods, we pick their best performance and explain why these methods are not well-suited for our problem.
Revised: To avoid unfair comparison to POMDP and continual RL methods, we pick their best performance and explain why these methods are not well-suited for our problem.

Original: For
Revised: For

Original: GPMM we run ten separate experiments with 15 episode each and picked three best runs and the best episode performance.
Revised: GPMM we run ten separate experiments with 15 episode each and picked three best runs and the best episode performance.

Original: For RNN-PPO we run three experiments, but picked the best learned policy over time and over the runs.
Revised: For RNN-PPO we run three experiments, but picked the best learned policy over time and over the runs.

Original: Furthermore, we use the RNN in a more favorable setting as we assume that the number of contexts is known and can be hard coded in the RNN architecture.
Revised: Furthermore, we use the RNN in a more favorable setting as we assume that the number of contexts is known and can be hard coded in the RNN architecture.

Original: We present our experimental results in Table A6.
Revised: We present our experimental results in Table A6.

Original: While it seems that RNN-PPO learns to swing-up correctly for soft failures and GPMM almost learns to swing-up for χ = 0 .
Revised: While it seems that RNN-PPO learns to swing-up correctly for soft failures and GPMM almost learns to swing-up for χ = 0 .

Original: 5 , the learned belief models indicate a
Revised: 5 , the learned belief models indicate a

Original: [Table]
Revised: [Table]

Original: that this not so.
Revised: that this not so.

Original: In fact, plotting the evolution of the belief for RNN-PPO and the context evolution for GPMM illustrates that the algorithms do not learn the context model (see Figure 4 in the main text).
Revised: In fact, plotting the evolution of the belief for RNN-PPO and the context evolution for GPMM illustrates that the algorithms do not learn the context model (see Figure 4 in the main text).

Original: In order to illustrate that it is not Gaussian Processes that cause failure in GPMM, we replace the Gaussian Process mixture with Deep Neural Network and Attentive Neural Process (Qin et al., 2019;
Revised: In order to illustrate that it is not Gaussian Processes that cause failure in GPMM, we replace the Gaussian Process mixture with Deep Neural Network and Attentive Neural Process (Qin et al., 2019;

Original: Kim et al., 2019) mixtures (DNNMM and ANPMM, respectively) using the code from Xu et al.
Revised: Kim et al., 2019) mixtures (DNNMM and ANPMM, respectively) using the code from Xu et al.

Original: The results in this case are similar to the GPMM case, i.e., we manage to get reasonable rewards for α = 0 .
Revised: The results in this case are similar to the GPMM case, i.e., we manage to get reasonable rewards for α = 0 .

Original: 5 , but fail for other cases.
Revised: 5 , but fail for other cases.

Revised: E. E XPERIMENTS WITH A LARGER NUMBER OF CONTEXTS

Revised: We further test our approach by introducing a larger number of contexts.

Original: a
Revised: We have the following parameter sets [Equation] [Equation] ] , with 16 parameters in total.

Revised: The contexts for our experiment are as follows:

Revised: C 0 : c friction = 0 .

Revised: 1 , c gravity = 9 .

Revised: 82 , c cart mass = 0 .

Revised: 5 , c max force = 40 ,

Revised: C 1 : c friction = 0 .

Revised: 1 , c gravity = 9 .

Revised: 82 , c cart mass = 0 .

Revised: 5 , c max force = 20 ,

Revised: C 2 : c friction = 0 .

Revised: 1 , c gravity = 9 .

Revised: 82 , c cart mass = 5 , c max force = 40 ,

Revised: C 3 : c friction = 0 .

Revised: 1 , c gravity = 9 .

Revised: 82 , c cart mass = 5 , c max force = 20 ,

Revised: C 4 : c friction = 0 .

Revised: 1 , c gravity = 50 , c cart mass = 0 .

Revised: 5 , c max force = 40 ,

Revised: C 5 : c friction = 0 .

Revised: 1 , c gravity = 50 , c cart mass = 0 .

Revised: 5 , c max force = 20 ,

Revised: C 6 : c friction = 0 .

Revised: 1 , c gravity = 50 , c cart mass = 5 , c max force = 40 ,

Revised: C 7 : c friction = 0 .

Revised: 1 , c gravity = 50 , c cart mass = 5 , c max force = 20 ,

Revised: C 8 : c friction = 0 , c gravity = 9 .

Revised: 82 , c cart mass = 0 .

Revised: 5 , c max force = 40 ,

Revised: C 9 : c friction = 0 , c gravity = 9 .

Revised: 82 , c cart mass = 0 .

Revised: 5 , c max force = 20 ,

Revised: C 10 : c friction = 0 , c gravity = 9 .

Revised: 82 , c cart mass = 5 , c max force = 40 ,

Revised: C 11 : c friction = 0 , c gravity = 9 .

Revised: 82 , c cart mass = 5 , c max force = 20 ,

Revised: C 12 : c friction = 0 , c gravity = 50 , c cart mass = 0 .

Revised: 5 , c max force = 40 ,

Revised: C 13 : c friction = 0 , c gravity = 50 , c cart mass = 0 .

Revised: 5 , c max force = 20 ,

Revised: C 14 : c friction = 0 , c gravity = 50 , c cart mass = 5 , c max force = 40 ,

Revised: C 15 : c friction = 0 , c gravity = 50 , c cart mass = 5 , c max force = 20 .

Revised: We learn a model with K = 20 .

Revised: We further add a structure on the contexts transition matrix: form every contexts we can switch only to and from two contexts, i.e, from the context i we can switch to/from the context i − 1 and the context i +1 , where operations on i should be understood as modulo a

Revised: K (i.e., − 1 (cid:44) K − 1 , K (cid:44) 0 ), i.e.: p ( z i | z j ) =

Revised:   0 .

Revised: 6 i = j, 0 .

Revised: 2 j = mod ( i + 1 , K ) or j = mod ( i − 1 , K ) , otherwise

Revised: Similarly to our previous settings we have a transition cool-off period of 5 time steps.

Revised: As the results in Figures A18(a) and A18(b) suggest, we identify the following meaningful contexts:

Revised: Z 1 = { C 1 , C 9 } , Z 4 = { C 6 , C 7 , C 14 , C 15 } , Z 10 = { C 4 , C 5 , C 12 , C 13 } , Z 14 =

Revised: { C 2 , C 3 , C 10 , C 11 } ,

Revised: Z 18 = { C 0 , C 4 , C 8 } .

Revised: Our algorithm does not distinguish the difference in the values of the friction parameter.

Revised: The difference in maximum force is not identiﬁed for larger masses and stronger gravity, but it can be identiﬁed for cart mass 0 . 5 and gravity 9 . 82 .

Revised: The other contexts are characterized by different gravity values and cart masses.

Revised: The only exception is the ground truth context C 4 , which is present in both Z 10 and Z 18 .

Revised: This, however, is bound to happen if the ground truth contexts are hard to separate.

Revised: [Figure]