An Adaptive Deep RL Method for Non-Stationary Environments with Piecewise Stable Context
Authors: Xiaoyu Chen
Xiangming Zhu
Yufeng Zheng
Pushi Zhang
Li Zhao
Wenxue Cheng
Peng CHENG
Yongqiang Xiong
Tao Qin
Jianyu Chen
Tie-Yan Liu
Venue: NeurIPS 2022 Accept (2022)

Abstract: One of the key challenges in deploying RL to real-world applications is to adapt to variations of unknown environment contexts, such as changing terrains in robotic tasks and fluctuated bandwidth in congestion control. Existing works on adaptation to unknown environment contexts either assume the contexts are the same for the whole episode or assume the context variables are Markovian. However, in many real-world applications, the environment context usually stays stable for a stochastic period and then changes in an abrupt and unpredictable manner within an episode, resulting in a segment structure, which existing works fail to address. To leverage the segment structure of piecewise stable context in real-world applications, in this paper, we propose a \textit{\textbf{Se}gmented \textbf{C}ontext \textbf{B}elief \textbf{A}ugmented \textbf{D}eep~(SeCBAD)} RL method. Our method can jointly infer the belief distribution over latent context with the posterior over segment length and perform more accurate belief context inference with observed data within the current context segment. The inferred belief context can be leveraged to augment the state, leading to a policy that can adapt to abrupt variations in context. We demonstrate empirically that SeCBAD can infer context segment length accurately and outperform existing methods on a toy grid world environment and Mujuco tasks with piecewise-stable context.

An Adaptive Deep RL Method for Non-Stationary Environments with Piecewise Stable Context
Xiaoyu Chen
Xiangming Zhu
Yufeng Zheng
Pushi Zhang
Li Zhao
Wenxue Cheng
Peng Cheng
Yongqiang Xiong
Tao Qin
Jianyu Chen o
Tic-Yan Liu
> Tsinghua University
Shanghai Jiao Tong University University of California, Berkeley I Microsoft Research Asia
Abstract
One of the key challenges in deploying RL real-world applications adapt variations of unknown environment contexts, such as changing terrains in robotic tasks and fluctuated bandwidth congestion control. Existing works on adaptation to unknown environment contexts either assume the contexts are the same for the whole episode or assume the context variables are Markovian. However, many real-world applications, the environment context usually stays stable for a stochastic period and then changes in an abrupt and unpredictable manner within an episode, resulting in segment structure, which existing works fail to address. To leverage the segment structure of piecewise stable context in real-world applications, in this paper, we propose Segmented Context Belief Augmented Deep (SeCBAD) RL method. Our method can jointly infer the belief distribution over latent context with the posterior over segment length and perform more accurate belief context inference with observed data within the current context segment. The inferred belief context can be leveraged to augment the state, leading to a policy that adapt to abrupt variations context. We demonstrate empirically that SeCBAD can infer context segment length accurately and outperform existing methods toy grid world environment and MuJoCo tasks with piecewise-stable context.
1 Introduction
Deep reinforcement learning has achieved great success in a wide range of challenging environments such as Atari games (Mnih et al. 2013) Bellemare et al., 2012) Hessel et al. 2017) or continuous control tasks (Schulman et al.] 2015 (017a). However, in stark contrast with this trend, applying RL to real-world applications remains a great challenge. In most real-world settings, there could be variations environmental factors, such as changing terrains in robotic tasks, fluctuated bandwidth in congestion control, and dynamic traffic patterns in autonomous driving. We refer to such environment factors as situation or context. The changes in context are not neglectable since context usually has substantial impact transition and reward functions. When the context is fixed and known to us, the problem is stationary and easier to solve. However, in most realistic settings, the context is usually dynamic within episode and unknown at test time. Therefore, detecting and adapting
Equul contribution. This work is conducted at Microsoft. Corresponding author.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Bandwidth (kbps)
800
600
400
200-
30
40
50
Time (s)
Figure 1: typical trace of network available bandwidth. The trace can be approximately divided into several segments of different length g², for i E 21,2,- ... 6}. The network condition changes abruptly at the end of each segment, while in each segment the network condition stays stable.
to variations in context is very important for RL agents to make a real impact in a wide range of real-world applications.
In real-world environments with varied unknown contexts, we find a typical evolving context pattern particularly interesting. The context usually stays the same for stochastic period until it changes abruptly and unpredictably into another context value which sampled i.i.d. from some prior context distribution. What's more, we usually do not have access to context directly but instead have access to some noisy observation H; sampled from some distribution p(zz(c) only at training time (e.g., auxiliary information from the simulator). We take fluctuated bandwidth in congestion control as an example. In Figure T we show a typical trace of network available bandwidth. The fluctuated bandwidth is usually modeled by multiple non-overlapping segments (Akhtar et al., 2018) Zhang and Duffield 2001). Within each segment, the network condition is stationary, and thus the bandwidth approximately follows the same distribution. At the end of segment, the network condition changes, and the bandwidth changes abruptly into another distribution. Here the context e represents the parameters that determine the distribution of fluctuated bandwidth (network condition), while the observations 2; represents the observed bandwidth which follows the above distribution p(zz/e). While the contexts are piecewise-stable, there could be slight variations in F, within each segment. This piecewise-stable pattern of context dynamics is of particular interest to us since it can capture a wide range of stochastic context processes in real-world applications, such as changing terrains in robotic tasks and fluctuating bandwidth in congestion control. Since the change in context abrupt and unpredictable, we cannot predict the future context in advance. The best we can do is detect and adapt to changes when they happen.
Although adaptation to varied unknown contexts has been studied under non-stationary RL and meta RL, few existing works look into piecewise stable context with abrupt changes within an episode.
Most works in meta RL as task inference (Rakelly al. 2019 Zintgraf et al., 2020 Zhao et al. 2020
Poiani et al. 2021) and some works in non-stationary RL (Chandak et al. 2020a Xie et al. 2021
assume the context stays the same for the whole episode and infer the context based the entire episode fc.f. Figure ((a)), therefore cannot quickly adapt to context changes within an episode. Other works on non-stationary RL assume intra-episode context changes and model c at each time step, but few study the piccewise-stable context as we do. Nagabandi et al. (2018) directly predict the context and thus cannot capture the prior that the context tends to stay the same for stochastic period. Feng et al. (2022); Ren et al.] (2022) model Markovian discrete context for each time step (c.f. Figure b)), therefore failing to model the non-mare non-markov property of context and the prior over context segment length in our setting. Compared with existing works, our setting is distinctive since the segment structure is latent to us (c.f. Figure
Therefore, the unique challenge in our setting is that we need to infer the segment structure, which can be further leveraged to infer belief context by only incorporating the relevant observed data in the current segment.
This paper studies how to infer segment structure to detect abrupt context changes and infer belief context accordingly to adapt to the piecewise-stie stable context in RL environments. We first introduce latent situational MDP which models RL environments with the stochastic situation/contex pro- cess (Section 2. Then, we introduce how to infer the belief context from observed data. To address the challenge above, we propose to infer context segment structure and belief context jointly from ob- served data (Section[3.1) Then we augment the state with inferred belief context so that the RL agent can automatically trade-off between acting optimally conditioned on inferred context and gathering

---

QL-2
R--1
Ri-2
a1-1
QL-2
61-1
(a) One Context
(6) Markovian Contexts
(c) Our setting
Figure 2: Probabilistic Graphical Models (PGMs) for different problem settings, with shaded circles for observable variables and white circles for latent variables. (a) One context: assume the context remains unchanged for the whole episode. (b) Markovian context: assuming the context is Markovian for each time step. (c) Our setting: the context remains unchanged in each segment, but the segment structure, illustrated by the red segment is unknown and needs to be inferred. We show two possible examples of PGM corresponding to G, = 2 ((c), left) and Ge = 1 ((c), right), where G3 measures the length of the current segment up to time step t.
more information about the current context (Section 3.2). Finally, we combine the training objectives for RI. and inference and present all the details for our proposed deep RI. algorithm (Section 3.3 We
evaluate our algorithm on a gridworld environment with dynamic goals and MuloCo tasks (Todoroy
al. (012) with varied contexts. Experiments demonstrate that our algorithm can quickly detect and
adapt to abrupt changes in piece-wise stable contexts and outperform existing methods(Seetionn 4
Our contributions can be listed as follows:
We introduce latent situational MDP with piecewise-stable context, which can capture a wide range of real-world applications (Section
We propose SeCBAD, an adaptive deep RL method for non-stationary environments with giecewise-stable context. Our method can infer the context segment structure and the belief context accordingly from observed data, which can be leveraged to detect and adapt to context changes (Section
Experiments on a gridworld environment and Mujuco tasks with piecewise-stable context demonstrate that our method can quickly detect and adapt to abrupt context changes and outperform existing methods (Section
2 Problem Formulation
In this section, we define a latent situational Markov Decision Process (LS-MDP) as a tuple M A, X, G, T, R, 7), where S is the set of states, A is the set of actions, and y E (0, 1] is the discount factor. To formulate the variations in environment factors, we introduce C, the set of latent contexts, and X, the set of observable contexts. C refers to the set of underlying hidden contexts that contains all necessary information but are left unobservable to the agent, while X' refers to the set of contexts with only partial information. G refers to the segment length which will be described in detail in the next paragraph. X are observable only during training and unobservable during deployment. In example environment setting where robot walking over changing terrain, C represents the terrain features containing perfect information, and X' represents noisy and imperfect information like mechanical metrics of the current step from virtual sensors in simulator and thus is only accessible during training.
In our setting, focus the case where the environmental changes are abrupt and irregular, in contrast to the smoothly changing assumption on context/task in existing works in non non-stationary RL (Chandak et al. 2020a To better model the generative process of the contexts, segment length G. Each episode is composed of several stationary segments with different segment lengths. For the ease of notation, we also introduce G₁, which measures the length of the current segment up to time step t. Then, the generative process can be described as follows: at the beginning of the h-th segment, the environment samples Ghh (G) and ch Pe(c) from prior distributions PG and Pe- Then, in the next Gh steps, the latent context c remains unchanged. For each time step t
introduce the
3

---

in this segment, the current segment length accumulates as G, G + 1 (we define G₁ = 1 for the first time step in this segment), and the current latent context satisfies that e, = c². The agent can observe p(Za(ct) for each time step if during training. The stationarity lasts until the end of the current segment, then the environment resamples Gh+1 and c+1 and the process repeats. We show Iwo examples of graphical models of our setting given different G, in Tigure[(e)
The transition function T p( 1 8t, 0t, c²) and the reward function conditioned on current latent context c². Therefore the changes in C lead to the changes in transition and reward functions. At test time, A' is no longer accessible, so we need to infer context changes from observed transitions and rewards. Since the latent context remains unobservable and changes silently, the environment is no longer stationary for the agents. To act optimally, it is important to keep track of the environment to recognize changes in time, and rapidly adapt to those changes.
p(ra|s at,c") c²)) are all
3 Methodology
In this section, we present our Segmented Context Belief Augmented Deep (SeCBAD) RL method and elaborate on how SeCBAD solves the challenges discussed above. Our method consists two main components:
Joint inference of the belief distribution over the latent context and the segment structure from observed data.
Policy optimization with inferred belief context under the belief MDP framework.
We first introduce the latent context inference part in Section 3.1 especially how infer the segment structure jointly with belief context and leverage the segment structure to remove irrelevant observed data. After the belief context is approximated, policy. We detail the policy optimization part under the belief MDP framework in Section 3.2 And finally, in Section 3.3
then incorporated into the state as the input of the
we describe how these parts constitute a practical algorithm.
3.1 Joint Inference for Belief Context and Segment Structure
In this part, we perform joint inference over latent context e and current segment length G, from observed trajectory Tis as to remove irrelevant data in Tizt for belief context inference. This can be formally expressed by the following equation:
p(a,, G. = p(c Ge+Lit G1, To- Ge- (pG
(1)
where Tip: (Sto: Oto: Pto ... 8t 04-1 1t-1, and ct G refers the latent context for the whole segment from t -- G₁ + to f. In Equation 1 we separately estimate the posterior
p(c' D(c'-G,+l:']G, Ti T--G,t) of latent context given known segment structure G2, and the posterior of segment length (T..)
3.1.1 Approximate Inference for Latent Context under Known Segment Structure
In this part, we focus on estimating p(c--G Mld-GG++lll TC-G.)), which is the posterior of the latent context under known segment structure GG. We use the variational inference framework to approxi- mate the true posterior. To be specific, we use an posterior inference network gø to infer the belief context in segment [t G2 + 1 : tl: go(ct-Gr+id G.+1: Gr, G.a) The variational lower bound for the log-likelihood of the current segment is given by:
log - 8t- Gt)
> 0e (log [Togpo(Ti- G 1. -(--
DKL - GG, TC
(.t))))) = (
(Gt) (2)
where Po denotes the decoder and T. x see Appendix A.
T-G (Tt-G,t, Tt- For the detailed derivation, please
The reconstruction term of Equation can be factorized as:
log Po ( G.a Gr, G-G,+ . O1-G.t-1, 8t
log po 8i, Gr, - si-11 a;-1)
(3)
G++1
This definition makes Tt-G.:t only contain observed data that belong the current segment.

---

Ps (Tg-gs
-t--g-1188-g)))
8o Co(c--0t:0 (--
@o 2o(c (T1))
Decoder
G
2o
Pe-
(e d
Te
Encoder
4o % ( m7-g-11t)
(a) Encoder-Decoder Architecture
(6) Belief Context as Mixture Distribution
An overview of SeCBAD. (a) peccoder-decoder architecture for belief context inference given current segment length G. = g: the encoder 9s takes the recent g steps the current red segment as input, and infer the belief context (o(c' (o(c"-9+1:f - and the decoder takes sampled context e from Is and decode the trajectory segment Tg-gut (the green nodes).
Figure 3:
(b) Belief context co(ce/Taa) as mixture distribution by considering belief context (c Qo(c--g+1: 1:: Tr-g:t)) for all possible structures according the posterior p(Gr/Tia)
which is the sum of log-probablity of transitions under context sampled from qa- The term Dxz(g=ll)) is the KL-divergence between our variational posterior ds and the prior over the belief context. For the prior p(ct-G++1:)) we use previous posterior at timestep / 1 if the context remains unchanged at timestep t, or (o, otherwise.
Unlike most previous methods which assume invariant context within an episode or markovian context each time step, our method 1) takes only observed data within current segment Tt-G.t as input for encoder Qs. 2) reconstructs only data within current segment 77 G,st assumption on context dynamics. By removing irrelevant data outside the current segment that are generated by past unassociated context, our method can estimate the current latent context more accurately.
Our method naturally motivated by the piecewise-stabe 7-GG,t at the output for decoder log po.
3.1.2 Iterative Inference for the Segment Length
In this part, we focus on estimating p(G,/Tıx), which is the posterior for current segment length, given (T..)) the posterior of context under given segment structure. To compute this posterior distribution, we first compute the joint distribution p(G,, Tia) recursively based on p(G,-1,,TIee as F
p(G. = -i, Σ -, p(G. i|G,-1 = k) p(st, -- (4)
The three major components in Equation 4 in turn are the previous joint distribution, the evolution prior, and the observation probability. The previous joint distribution is iteratively provided the beginning of each time step. The evolution prior p(G, = ilG k) measures the prior knowledge on the segment length G, given the segment length of the previous time step G, where either G. G + 1 or G, = 1 holds.
As for the observation probability which is the third term of Equation 4 it measures how likely the observations show given the history of the segment. To be specific, we have
p(st, Ft-.,,,,1-11-----.. = KE P ( i. aa p(s) r 1 c(tlit)) (5)
'For notation simplicity. omit the condition G, = é in the term p(s 0r 1, ie. i the start seement.
In Equation
K p(a--1/sc-) 115 - is constant with respect to and has no impact on the posterior
p(G, d)):)).
5

---

In Equation 5 the term p(s 1t-1|st-1, a,,,, state-reward pair, where e is drawn from the posterior distribution given the data the segment before timestep L. To compute the RHS of Equation 5 we sample from 2s(c/=i+1--1 G. i, Ti:it 1 which is the belief context inferred from the whole history of the segment before timestep f, and use the sampled the decoder to compute the data likelihood.
estimates the data likelihood for the next
Given the joint distribution, the posterior distribution of G, can be derived as
p(G, i, p i|Tr.)) E,P(G, = Ti:))
(6)
We also incorporate observable contexts in the observation probability to improve accuracy during training. See Appendix AA.3 for more details.
3.1.3 Belief Context as a Mixture Distribution
Given inferred posterior of Ge in Section 3.1.2 the belief of the latent context time step can be derived as
be(c) ( (1a) = Col 9 set = Sc-et))G ge F1:e)
This mixed probability of t which has taken all possible segment structures into consideration, can now represent the current belief by of the latent context c.
3.2 Policy Optimization with Belief Context
Inspired by the belief MDP (Kaclbling et al., 1998). Bayes Adaptive MDP (BAMDP) (Dufi 2002 and recent works on meta RL as task inference (Zingrraf et al. 2020). we incorporate the inferred belief context into the augmented state. At each time step t, the belief latent context is approximated via 9o using Equation 7 Therefore, we define the augmented state as (s.b) E S x B, where S is the same state space as in LS-MDP and B is the set of belief latent contexts. Accordingly, we have transition (s 1, Tr. be+ |sc, Ge. b,) = p($++ F' Pels,, o,, (bb ar, Tr, S24 6,) and reward R (ri)ss,, b,, ax) This definition brings advantages in the sense that the information gathering and exploitation tradeoff is no longer problem under such augmented states, since the transition and reward functions are no longer conditioned on exact c. Now, the policy is defined as x x (a|s, b) mapping from the augmented state space to the action space, and the agent's objective is to maximize
H Σ
JRL = E Egn.do TT
r't (relse, by,
3.3 Algorithm and Implementations of SeCBAD
In this section, we describe the overall algorithm and implementation details of SeCBAD. See Figure 3 for overview of our framework. As shown in Figure 3a), we use GRU (Cho et al. 2014)
parameterized by Φ as the recurrent encoder qo. and distributions in latent context space assumed to be diagonal Gaussians with mean and variance parameterized by Qs- The decoder includes transition model pa ($141|$s (s, ae c), reward model po f (t, a,, c) and observable context model po(xx|() (c). The output of all the decoders are Gaussian distributions with mean parameterized by foed-forward neural networks and fixed identity covariance. Then, we estimate p(C pG, using Ya and pe as described in Section 5.1.2 As shown in Figure (b) we combine the belief context based on different segment according to p(Ga|TT:d to get the belief (x(c) = ( [T.t), where one approach is to provide a total of mean, covariance and weights as policy input. For simplicity, we choose G2 with highest probability in p(Gr TT:))) and use the corresponding 9 (c [T= the belief. Empirically, we find this approximation leads to little performance loss. We build our RI. algorithm on the top of PPO (Schulman al. 2017b) to learn the policy nhe (ag (st, (c)), where v denotes the parameters in the actor and the critic. We use the objective described in Section 3.2 to optimize the policy.
When i p(c -+ G. = i, 1.) For the case where i = 1, the posterior distribution p(c" i, 71- actually the prior distribution Pe(c) = ( 1). and can sample from the prior distribution.
can use the belief go(c - --t------------------------
Tt-it-1) to approximate the posterior

---

10
Decrioatoo
belef dwnl CT0M
Figure 4: Experiment on grid world with dynamic goal state. We show case study on the inferred belief. The leftmost figure depicts the ground truth goal position, and the next two figure depicts the belief estimated by vanilla inference and SeCBAD. where the color of one grid measures how the agent believes 8 locates at that grid. The rightmost figure is G? estimated by SeCBAD. The result shows that our belief matches the ground truth more closely.
During deployment, ls and pe fixed. We first compute the segment posterior p(Ge(Ti:)) using encoder ls and the transition and reward model pp- Then, we estimate the belief (, and feed into the policy as input. For more implementation details, please refer the Appendix A.6
4 Experiments
In this section, we empirically evaluate SeCBAD two tasks. We first demonstrate our algorithm in a grid world environment in Section [4T illustrate the significance of incorporating segment structure during inference. For large-scale experiments, we test the proposed algorithm in Section 4.2
The results multiple challenging tasks show that SeCBAD outperforms baselines in terms of performance and sample efficiency. In Section 4.3 we further provide case studies of agent behaviors and learned latent to gain more insights into the results.
4.1 Motivating Example
In this experiment, we use a motivating example to show that it is vital for the agent to consider the segment structure to adapt to variations of unknown environment contexts. The agent is in 4 x 5 grid world containing goal state s* and 19 other states. In s", the agent can gain + reward with probability po. and 0 reward with probability 1 Po. For other states, the agent can gain +1 reward with probability P1 and reward with probability 1 - Pr, where pa > Pr. The actions include up, down, left, right, and none. The specific location of s* unobservable the agent, and may change after a stochastic period.
We use the analytical belief update rule instead of variational inference in this experiment. To be specific, we assume the qo and pe is known and fixed and then infer p(G. and the belief by (c), which represents the belief distribution of the location of s The belief update rule is detailed in Appendix A.55 With the analytical belief update rule, a theoretically optimal belief by is computed for each environment step t. We concatenate (st, be) together as the augmented state. We may expect that a policy trained upon this augmented state will achieve better performance b, is an accurate enough guess the true location of the goal state. In such a way, aim to focus the effectiveness of inferring segment structure, in terms of how inferring segment structure can arrive at meaningful belief state. We compare algorithm with vanilla inference. where the agent updates belief without considering segment structure, like in Zintgraf et al. (202)); Rakelly et al. (2019). For both of two routines of updating belief state, we let the agent move along same path. By moving along same path, we mean the agent starts from the same location and then take the same action in each environment step for belief update rules.
Figure illustrates how SeCBAD and the vanilla inference baseline update their belief over states. It is shown that the belief (c) of SeCBAD closely matches with the actual goal state s*. Also, compared with SeCBAD, the vanilla inference baseline needs more steps to detect the changes since it needs to correct the deviated prior belief. This result supports that incorporating segment structure into inference like SeCBAD is significant for estimating accurate belief in fast varying environments.
4.2 Locomotion Control Tasks with Varying Contexts
In this section, we show that SecBAD is able to tackle more complex tasks by testing the algorithm on several challenging control tasks with varying contexts. We conduct the experiments on modified MuJoCo (Tooorv et al., 2012) tasks. These environments are commonly used in previous works
7

---

Ant Direction
Ant Velocity
Cheetah Direction
Cheetah Velocity
Cheetah Goal
400
-30-
Episode
-00
-300
2011
600
200
-2000
100
40
126
124
Steps(M)
Steps M)
Steps(M)
(teps(M)
Steps(
SeCBAD
VariBAD
FANS-RL
PPO-RNN
Figure 5: Performance curves on 5 MuJoCo environments with variations in contexts. SeCBAD achieves better performance and sample efficiency in various challenging control tasks.
(Zintgraf et al. 2020 Rakelly al. 2019), and we further modify the contexts provide challenges corresponding to LS-MDP as mentioned in Section As for Ant Direction and Cheetah Direction, the policy needs to move towards the target direction fast as possible. As for Ant Velocity and Cheetah Velocity, the policy then needs to additionally move at the speed of the target velocity besides the target direction. For Cheetah Goal, the policy needs navigate varying goals and stay there. For Ant environments, the direction and velocity are specified on 2-dimensional plane, and for half cheetah environments, the direction and velocity are I-dimensional. In these environments, the contexts refer to the target velocity, target direction, or location of the goal. To make it more challenging, we further add noises to contexts within each segment. Please refer to Appendix 6A.6 for more details. As for baselines, we select one representative algorithm for each type in Figure 2 to compare since to the best of our knowledge, few existing works study the same setting as we do. We use VariBAD (Zintgraf et al.] 2220) to represent those methods assuming one context in the episode, and use FANS-RL (Feng et al. 2022) to represent those methods assuming Markovian contexts. Since LS-MDP is special case of POMDP, we include a POMDP baseline PPO-RNN (Hausknecht and Stone 2015) for comparison. The results are provided in Figure
The experiment results empirically show that SeCBAD achieves superior performance and sample efficiency on challenging control tasks. As illustrated in Figure SeCBAD achieves higher scores than baselines. We provide some insights into the results. For the method assuming that contexts stay the same within an episode, (Zintgraf et al. 2020) uses the learned latent contexts decode the whole trajectory including transitions and rewards in other segments. This reconstruction mismatch may lead to averaged latent contexts so that the policy cannot act correspondingly (see detailed analysis and the ablation study in Appendix (A.733) For Markovian contexts baseline (Feng et al.) (022) and POMDP baseline (Hausknecht and Stone 2015), they can perform relatively well on tasks where the contexts can be inferred from only one step transition (i.e., Cheetah Direction and Cheetah Velocity). However, for more complex tasks where the contexts need more steps to infer, SeCBAD significantly outperforms these two baselines (Feng et al. 20222 Hausknecht and Stone 2015), which proves that the specially designed joint inference component in SeCBAD is effective and can help improve the performance.
To better understand the proposed algorithm, we conduct series of ablation studies. In Appendix A.7.2 we study the effects of inaccurate prior on SeCBAD. In Appendix AA.7.4] we study the implementation choices on how to use p(G))TT.). In Appendix A.7.5 we further test SeCBAD against different levels of noises within each segment.
4.3 Case Studies on the Segment Structure
In order to provide more insights into the results, present study on Ant Direction in this section. We show the behavior and learned latent along with the inferred segment of a randomly selected episode during testing of the SeCBAD algorithm on Ant Direction in Figure
In the first row, we exhibit the task direction in orange as well as the agent's actual direction in blue. It shows that SeCBAD can rapidly detect and adapt to context changes. We also show the learned latent contexts in the second row and the inferred segment structure the third row. From the third row, it can be seen that the agent can detect the segment in time as the inferred segment structure closely matches the ground truth segment structure. As for the segment starting from the 430th environment step, the change in goal direction is negligible so that the agent automatically merges
8

---

Direction
Agent
Task direction
249
390
400
Latent meon
14
Inference
106
269
399
400
Environment step
Figure 6: case study on Ant Direction of SeCBAD. the first row, we show the behavior of the agent with the goal direction in orange and the actual agent direction in blue. In the second row, show the learned latent contexts with different colors corresponding to different dimensions. In the last row, we plot the G; inferred by SeCBAD.
the two segments. Since the latent contexts are calculated according to the inferred G7, the latent contexts in the second-row change the segment changes while staying stable within each segment. This proves that the agent can be aware of the changes in unknown contexts and quickly adapt to the changes. We provide case studies of baselines and detailed analysis in Appendix A.7.I
4.4 Bandwith Control Tasks for Real Time Communication
To further illustrate that the proposed LS-MDP setting can boost the deployment of RL in many real-world applications, we test SeCBAD on real-world bandwidth control task for real-time communications (RTC) (OpenNetLab) 202T) in this section. The most critical goal in RTC is to provide high Quality of Experience (QoE) for users. To achieve this, bandwidth control module is needed, ie.. the RTC sender needs to decide the bitrate of outstreaming audio/video based on the network status towards the receiver. However, the network conditions constantly change and the changes occur in multiple items, which makes this problem intricate and hard to solve.
In our reinforcement learning formulation, we use 7-tuple of current network statistics that is visible to the agent as states 8t, consisting of sending rate, short-term and long-term receiving rate, loss, and delay. The action a, is the estimated bandwidth. We provide the detailed settings in Appendix A.88 The latent context C, here refers to the fluctuated network condition, in this section, we consider v, as the ground truth bandwidth capacity C, and the RTT. In this experiment, with VariBAD (Zintgraf et al. 2020), FANS-RI (Feng et al. 2022), vanilla PPO (Schulman et al.. 2017ab and PPO-RNN (Hausknecht and Stone 2015). We also incorporate oracle PPO scores by incorporating the unobservable contexts into the observable states. All the methods are trained for 10 million steps and the shaded area is across 3 random seeds.
compare our method
As illustrated in Figure 7, SeCBAD achieves better performance than other baselines and is very close to the oracle PPO baseline score. The performances of VariBAD (Zintgraf et al., 2020) and FANS-RL (Fngg et al. 2022) are better than PPO-RNN (Hausknecht and Stone 2015), but SeCBAD outperforms both of these methods. The results suggest that SeCBAD is able to detect and adapt to the varying contexts more rapidly, which allows the policy to precisely control. For detailed descriptions and results, please refer to Appendix A.8
Related works
Our work is closely related to non-stationary RL, where the transition and reward functions may change over time. Most existing works on non-stationary RL focus on inter-episode non- stationarity (Xie et al. 2021 Chandak et al. 2020a bl Xie et al. [022 [odhanii et al. [021 Alegre jet al.) 2021; Poiani et al. 2021 AI-Shedivat et al. (017), some of which adopt contextual MDP as formulation (Hallak et al.) 2015). Recently there are some works considering intra-episode non-

---

Bandwidth Control
Epicode Return
Steps(M)
SeCBAD VarBAD
TAKS-RI
PPO
PPO RAN
Oraci PPO
Figure 7: Experiment results on bandwidth control for RTC.
stationarity (Ren al. 2022 Kamienny et al. 2020 Kumar al. 2021 Nagabandi et al. 2018
Feng et al. 2022). Ren et al. (2022) assume the latent context to be finite and Markovian, while Feng al. (2022) assume the latent context is Makovian and the environment can be modeled as factored MDP. In contrast to existing works on non-stationary MDP, we assume piecewise-stable context with abrupt changes within episode, which is
realistic and capture wide range of real-world applications. To rapidly adapt to dynamic context, the agent needs to continuously perform information gathering behavior. Bayesian RL (Duff) 2002 Zintgraf al. 2020 Fellows al. [021) is an elegant framework to optimally tradeoff the exploration and exploitation in an unknown and stationary MDP. As a special type of Belief MDP (Kaelbling et al. 1998). Bayes Adaptive MDP (BAMDP) (Dfff 2002) maintains a belief over the environment and uses this belief to augment the state. Our model can viewed as a special case of belief MDP, where we only maintain belief over latent context to trade off between information gathering and exploitation. To accurately infer belief context, we adopt the variational inference, which has been adopted by many existing works in RL for task inference (Rakelly et al., 2019 Zhao et al. 20200 Humplik et al. 2019 Poiani et al. 2021 Zintgraf et al. 2020) or context inference (Xie al. [021 Feng al.) 2022 Ren et al. 2022). However, none of these methods suit our setting. We perform joint inference over latent context and segment structure structure from observed data, so as to remove irrelevant data for more accurate belief context inference. LS-MDP can also be viewed as special case of POMDP. Recently, progress has been made in learning the latent dynamics model (Krishnan et al. 2013 Karl et al. [016]] Doerr et al.) 2018 Buesing al. 2018: Ha and Schmidhuber] 2018 Han et al. 2019 Hafner al. 2019| a). Theoretically. it is possible to perform optimally only using recurrent neural networks (RNNs) like Hausknecht and Stone (2015) since the whole history has been taken into consideration. However, it has been shown that (Hafner et al. 2019b) introducing more structured information will significantly enhance the performance. In this paper, exploit the addtional assumption of LS-MDP to infer the context based on segment structure. The experiments show that SeCBAD can achieve better performance compared with existing methods.
6 Conclusion
In this paper, propose SeCBAD, a Segmented Context Belief Augmented Deep RL method to deal with piecewise-stable context in non-stationary environments. Piecewise-stable context is quite common in wide range of real-world applications. Compared with existing methods, our method can automatically detect the segment structure, which reflects when the context changes abruptly. The detected segment structure can be further used compute context belief with only relevant observed data. To the best of our knowledge, this is the first method that can model and leverage piecewise-stable context reinforcement learning to help the agent adapt to environment change. With inferred belief context, our RL agents can quickly detect and adapt to abrupt changes in gridwold environment and mujoco tasks with piecewise-stale leverage various deep learning techniques to improve SeCBAD, which includes replacing the GRU encoder by Transformer to better capture the long-term dependency in the input trajectory segment.
context. For future work, we plan to
References
Akhtar, 7.,, Nam, Y. S., Govindan, R., Rao, S., Chen, J., Katz-Bassert E., Ribeiro, B., Zhan, J., and Zhang, H. (2018). Oboe: Auto-tuning video abr algorithms to network conditions. In Proceedings
10

---

of the 2018 Conference of the ACM Special Interest Croup on Data Communication, pages 44-58.
Al-Shedivat, M., Bansal, T., Burda, Y., Sutskever, I., Mordatch, I., and Abbeel, (2017). Continuous adaptation via meta-learning in nonstationary and competitive environments. arXiv preprint 17X1::1100.06411
Alegre, L. N.. Bazzan, A. L., and da Silva, B. C. (2021). Minimum-delay adaptation in non- stationary reinforcement learning via online high-confidence shange-point detection. arXiv preprint arXiv:2105.09452.
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. (2012). The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, Vol. 47:253- 279. cite arxiv: 1207.4708
Buesing, L., Weber, T., Racaniere, S., Eslami, S., Rezende, D., Reichert, D. P., Viola, F., Besse, F. Gregor, K., Hassabis, D., et al. (2018). Learning and querying fast generative models for reinforcement learning. arXiv preprint arXiv:8802.00 1e0..0002000000
Chandak, Y., Jordan, S., Theocharous, G. White, M., and Thomas, P. S. (2020a). Towards safe policy improvement for non stationary mdps. Advances in Neural Information Processing Systems, 33:9156-9168.
Chandak, Y., Theocharous, G., Shankar, S., White, M., Mahadevan, S., and Thomas, (2020b). Optimizing for the future in non-stationary mdps. In International Conference on Machine Learning, pages 1414-1425 PMLR.
Cho, K., Van Merriënboer, B., Gulcehre, C., Bahdanau, D., Bougares, F., Schwenk, H., and Bengio, Y. (2014). Learning phrase representations using encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1406. 1078.
Doemr A., Daniel, C., Schiegg, M., Duy, N.-T., Schaal, S., Toussaint, M., and Sebastian, T. (2018). Probabilistic recurrent state-space models. In International Conference on Machine Learning. pages 1280-1289. PMLR.
Duff, M. O. (2002). Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. University of Massachusetts Amherst.
Fellows, M., Hartikainen, K., and Whiteson, S. (2021). Bayesian bellman operators. Advances in Neural Information Processing Systems, 34.
Feng, F., Huang, B., Zhang, K., and Magliacane, S. (2022). Factored adaptation for non-stationary reinforcement learning. arXiv preprint arXiv:2203.1658
Ha, D. and Schmidhuber, J. (2018). World models. arXiv preprint 1803.1280...112..
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. (2019a). Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.0166
Hafner, D., Lillicrap, T., Fischer, I,, Villegas, R., Ha, D., Lee, H., and Davidson, J. (2019b). Learning latent dynamics for planning from pixels.
Hallak, A., Di Castro, D., and Mannor, (2015). Contextual markov decision processes, arXiv preprint arXiv:1502.022559
Han, D., Doya, K., and Tani, J. (2019). Variational recurrent models for solving partially observable control tasks.
Hausknecht, M. and Stone, P. (2015). Deep recurrent q-learning for partially observable mdps. In 2015 aaai fall symposium series.
Hessel, M., Modayil, J., van Hasselt, H., Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., and Silver, D. (2017). Rainbow: Combining improvements in deep reinforcement learning, cite wtt.1177002298Commeet Under review conference paper at AAAI 2018.
11

---

Humplik, J., Galashov, A., Hasenclever, L.,, Ortega, P. P.A., Teh, Y. W. and Heess, N. (2019). Meta reinforcement learning as task inference. arXiv preprint arXiv: Nttt::9905006444
Kaelbling, L. P., Littmann M. L., and Cassandra, A. R. (1998). Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2)999-134
Kamienny, P.-A., Pirotta, M., Lazaric, A., Lavril, T., Usumier, N., and Denoyer, L. (2020). Learning adaptive exploration strategies in dynamic environments through informed policy regularization. arXiv preprint ar arX1v:2005.02934.
Karl, M., Soelch, M., Bayer, J., and Van der Smagt, P. (2016). Deep variational bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint arXiv:1605.064: 06432.
Krishnan, R. G., Shalit, U., and Sontag, D. (2015). Deep kalman filters. arXiv:1511.05121.
arXiv preprint
Kumar, A., Fu, Z., Pathak, D., and Malik, J. (2021). Rma: Rapid motor adaptation for legged robots.
Mnih, V., Kavukcuoglu, K., Silver, D., Graves, A., Antonoglou, I., Wierstra, D., and Riedmiller, M.
(2013). Playing atari with deep reinforcement learning. cite 1xxx:1312.5622Commen NIPS Deep Learning Workshop 2013.
Nagabandi, A., Finn, C., and Levine, S. (2018). Deep online learning via neta-learning: Continual adaptation for model-based rl. arXiv preprint
OpenNetLaab (2021). Alpharte. In hutps//giiaacccmmmmpeevvv//aa//a
Poiani, R., Tirinzoni, A., and Restelli, M. (2021). Meta-reinfforcemen learning by tracking task non-stationarity. arXiv preprint arXiv:2105.08834
Rakelly, K., Zhou, A., Finn, C., Levine, S., and Quillen, D. (2019). Efficient off-policy meta- reinforcement learning via probabilistic context variables. International conference on machine learning, pages 5331-5340. PMLR.
Ren, H., Sootla, A., Jafferjee, T., Shen, J., Wang. J., and Bou-Ammar, H. (2022). Reinforcement learning in presence of discrete markovian context evolution. arXiv preprint arXiw:2202.06553.
Schulman, J., Levine, S., Abbeel, P., Jordan, M. I., and Moritz, P. (2015). Trust region policy optimization. In Bach, F. R. and Blei. D. M., editors, ICML, volume 37 of JMLR Workshop and Conference Proceedings, pages 1889-1897. JMLR.org.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017a). Proximal policy optimization algorithms. CoRR, 000/1707706377
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017b). Proximal policy optimization algorithms. arXiv preprint arXiv:1707.0634
Sodhani, S., Meier, F., Pineau, J., and Zhang, A. (2021). Block contextual mdps for continual learning. arXiv preprint arXiv:21 10.06972.
Todorov, E., Erez, T., and Tassa, Y. (2012). Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference intelligent robots and systems, pages 5026-5033. IEEE.
Xie, A., Harrison, J., and Finn, C. (2021). Deep reinforcement learning amidst continual structured non-stationarity. In International Conference on Machine Learning, pages 1393-11443 PMLR.
Xie, A., Sodhani, S., Finn, C., Pineau, J., and Zhang, A. (2022). Robust policy learning over multiple uncertainty sets. arXiv preprint arXiv:2202.07013.
Zhang, Y. and Duffield, N. (2001). On the constancy of internet path properties. In Proceedings of the 1st ACM SIGCOMM Workshop on Internet Measurement, pages 197-211.
Zhao, T. Z.., Nagabandi, A., Rakelly, K., Finn, C., and Levine, (2020). Meld: Meta-reinforcement learning from images via latent state models, arXiv preprint 13357
12

---

Zintgraf, L.. Shiarlis, K., Igl. M., Schulze, S., Gal, Y., Hofmann, K., and Whiteson, S. (2020). Varibad: very good method for bayes-adaptive deep via meta-learning In International Conference on Learning Representation (ICLR).
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "kuHDbe1tQQI", "reviewer": "Reviewer_ehim", "summary": "This paper tacles RL on non-stationary MDP. Based on the core assumtion that the environment context is piecewise stable, the authors prospose a Segmented Context Belief Augmented Deep (SeCBAD) RL method, which can jointly infer the belief context and segment structure from observed data. Experiments on a gridworld environment and Mujuco tasks show that the proposed method can quickly detect and adapt to abrupt context changes and outperform existing methods.", "questions": "In principle, if the segment of the piecewise stable MDP lasts long enough, variBAD and its variants should be able to fully adapt to each segmented context. For instance, in Figure 4 of the original paper [1], variBAD adapts to a stationary MDP within 1 episode. However, both Figure 5 and Figure 7 show that they learn almost nothing, which is strange.\n\n[1]: Zintgraf, L., Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hofmann, K., and Whiteson, S. (2020). Varibad: A very good method for bayes-adaptive deep rl via meta-learning. InInternational Conference on Learning Representation (ICLR).\n", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper tackles a relavant problem, non-stationary MDP, and picks a simplified setting based on piecewise stable assumption, which has merits in some real-world scenarios (robotic and congestion control).\n- The presentation of the paper is generally clear, despite some typos (e.g. line 5).\n\n\nWeaknesses:\n\n- **Limited scope**. The contributions of the proposed method depend entirely on the assumption that the non-stationary MDPs are piecewise stable, which inevitably limits the scope of real-world applications, e.g., greenhouse control given stochastic outdoor climate and quantitative trading in unpredictable stock market.  \n\n- **Limited novelty**. The proposed algorithm employs the exact same formulation of approximate inference and training objectives of variBAD and makes incremental modifications to the framework to incorporate iterative inference for the segment length. The contributions to both the theoretical and algorithmic aspects of BAMDP are limited.\n\n- **Lack of experiments**. The proposed method is only evaluated on 3 relatively simple environments (1 tabular + 2 MuJoCo) and compared to 1 baseline, variBAD, which is designed for stationary MDP. More baselines targeting non-strationary MDP [1][2] should be included to validate the effectivenss of the proposed method. Also more focused studies such as the robustness of SeCBAD to noise within segmented context could bring more insight to the readers.\n\nMinor issues\n\n- The paper seems a bit rushed. More details of the setup and experiments are needed, for example:\n  - How many random seed/trials were performed to produce Figure 5 and 7?\n  - What is $G$ defined in line 89?\n\n[1]: Ren, H., Sootla, A., Jafferjee, T., Shen, J., Wang, J., and Bou-Ammar, H. (2022).  Reinforcement learning in presence of discrete markovian context evolution.arXiv preprint arXiv:2202.06557\n\n[2]: Feng, F., Huang, B., Zhang, K., and Magliacane, S. (2022). Factored adaptation for non-stationary reinforcement learning.arXiv preprint arXiv:2203.16582.\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "OLbLJeD7h-r", "reviewer": "Reviewer_PKbi", "summary": "Authors investigate the non-stationary RL problems defined as scenario MDP where there exists an unobservable context variable that could change abruptly and unpredictably but remain stable within a stochastic period.  The general idea of this paper is to develop a joint inference framework for the segment structure and the belief context. A variational inference method similar to variational auto-encoder (VAE) is proposed to infer the context given a segment of the past trajectory. An iterative inference method is also proposed for the segment length.  The infered context is then augmented to the observed state variable as inputs to RL algorithms. Experiments on grid world and Mujuco are conducted to verify the performance of the proposed method. ", "questions": "1. What is the complexity of the proposed algorithms? According to Eq. 5, the posterior of the belief context has to sum over all possible segmentation and if it is the case, then the computational complexity seems to large for larget t.\n\n2. If we look at the motivating example in Fig.1, it seems that the change of the context can be easily detected using the state variables, e.g. we can apply well developed \"change point\" detection algorithms to the observables to detect the change of the context. Do we really need such a complex method? I would like authors to further comment on this question.\n\n3. In this paper, is the context variable a scaler? If we extend it to a high dimensional latent variable, whether the proposed method is able to solve the problem is not clear. \n\n", "limitations": "The major concern of this paper is as stated \"Authors do not compare the proposed method to an important brunch of research on Partial Observable MDP (POMDP) in the main text and also in the experiment section. \"", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n1. RL for non-stationary environment is an important research topic and advanced methods on this topic could potentially improve the applicability of RL algorithms to real world problems. \n\n2. The scenario MDP formulation proposed in this paper is interesting and could be a more realistic formulation than existing Markovian ones. The joint inference method of the segment structure and the belief context is also non-trivial and thus novel, to the best of the reviewer's knowledge. \n\n3. In general the paper is well motivated and clearly written and techincally sound.\n\nWeakness:\n\n1. Authors do not compare the proposed method to an important brunch of research on Partial Observable MDP (POMDP) in the main text and also in the experiment section. \n\n2. The RL part is fairly weak as the proposed method simply augment the learned context to the state and then feed them to any existing RL algorithms.  To me, the whole framework seems a two-stage approach where the connection between the two stages are quite weak.\n\n3. The experiment section is weak. As stated as the motivation of the paper, non-stationarity could be the major issue for RL to be applied to real tasks, and I expect authors to demonstrate the performance of their algorithms on real tasks. However, experiments on real tasks are not presented.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FDTGIbjlxuL", "reviewer": "Reviewer_5PEc", "summary": "The paper introduces the problem setting where the environment has a “piecewise-stable context” (i.e. the context is constant for some stochastic length of time then changes abruptly to another context several times in an episode), called LS-MDP. Previous works tackle either extremes where either (1) only a single constant context is sampled for the episode or (2) there is a Markovian context at each timestep. The paper then proposed a method, SeCBAD, which performs joint inference for the segment structure (i.e. how long each context occurs for) the belief context given segment structure using variational inference techniques, and combines them in a mixture distribution to marginalize over the possible segment structures. The policy then is learned with the belief context as additional information (instead of conditioning on the exact context). \n\nExperiments were conducted both in a didactic grid-world example, as well as in larger school Mujoco environments, i.e. inferring ideal movement direction for ant locomotion and half-cheetah with changing wind direction. The inferred segment structure is also visualized in the half-cheetah environment to indicate agreement with the underlying context changes.", "questions": "Please see my questions below:\n1. The experiments would be stronger with additional baselines in addition to the VariBAD baseline. For example, one of the other works assumes intra-episode context changes and models $c_t$ at each timestep. This will help to strengthen the result that the piecewise-stable prior helps over Markovian context type approaches in your environment settings. \n2. It would be good to provide empirical data to back up the claim that using $G_t^*$ with the highest probability in $p(G_t|\\tau_{1:t}))$ and then using the corresponding belief lead to little performance loss. \n3. What is the sensitivity of the training schedule between training the encoder/decoder/policies? In Algorithm 1, it seems that they are all updated at the same pace, but I am wondering if you have also explored for example training the encoder/decoder for more steps relative to the policy optimization step, etc. \n4. More hypothetical, not asking for experimental data: Could the context transition also depend on the states reached? E.g. reaching a new terrain, etc. How would your method behave in those types of environments?\n", "limitations": "As mentioned above, the paper did not describe much limitations and potential negative societal impact of their work. Perhaps it is also worth to discuss about the potential sample complexity required to collect enough data to learn the encoder/decoder, or what happens if the underlying MDP has markovian contexts / one context, then would the approach still be able to also solve those situations just as well as other baselines (i.e. if there is a mismatch in the MDP assumption and the method). ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Originality**: The paper was effective in contextualizing its problem formulation (pun intended) in the sea of non-stationary RL works, while clarifying that LS-MDP is a special case of POMDP. Given these additional assumptions in the problem setting, then we can exploit them in the solution. The proposed approach applies well-known tools like variational inference methods for this problem setting, but there is merit in showing that these existing tools can provide significant gains over methods that have more general/different problem assumptions (such as VariBAD). \n\n**Quality**: The paper mainly provides empirical results to analyze the proposed method to demonstrate that it can incorporate the piecewise-stable context for adapting to abrupt environmental changes. However, the limitations of the method was not discussed in detail, nor was there a section on potential negative societal impacts the work.  \n\n**Clarity**: Overall the paper is relatively clear to read, with helpful conceptual diagrams for the method. There are several minor comments for improving the draft:\n\n1. In the abstract, second sentence seems incomplete. \n2. Figure 4: should label the y-axis as ‘time step’ or something along that lines to help quickly understand how to read the plots. Based on the last paragraph on page 7 I was also expecting a plot for the vanilla inference behaviour but I do not see that in there. \n3. For the G_t plots, perhaps it’s also informative to plot the oracle $G_t^*$ value (i.e. the perfect zig-zag lines) to help compare that with the estimated $G_t^*$ by the proposed algorithm\n\n**Significance**: The paper proposes a relevant problem setting and demonstrated that their method can yield noticeable improvement. There are still a lot of architecture improvements to the method (as mentioned about using Transformer instead of GRU in the recurrent encoder), and more likely also new approaches for the given problem setting. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6GCaJoJuH7k", "reviewer": "Reviewer_3SWX", "summary": "This paper provides a new method called SeCBAD (Segmented Contect Belief Augmented Deep RL), to deal with the segment structure of piecewise-stable context. Experiment demonstrates that this method can automatically detect the segment structure, which reflects when the context changes abruptly.", "questions": "1. How does this problem compare/relate to distribution shift? Should cite/discuss some distribution shift works.\n2. Although it's termed context in this work, the formulation of the method looks quite similar to goal-conditioned RL. Can the authors provide some explanations on how the two are different?\n3. I'm a bit confused by the Half-Cheetah Velocity case study. From my understanding, velocity should be the target/reward in this environment. However, it seems to be treated as the context in this paper?\n4. Just curious, if this method can detect and adapt to abrupt changes, then can it also adapt to slow and insidious context changes?", "limitations": "1. Seems to be a straightforward extension of variational inference framework.\n2. Experiment is done in specifically modified environments. Therefore, its objectiveness and fairness remain to be confirmed. \n3. Clarity of experiment set up can be improved.\n4. The paper claims that 'Piecewise-stable context is quite common in a wide range of real-world applications'. However, the authors did not give a concrete example of this.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe discussed latent situational MDP with piecewise-stable context seems to be potentially useful to model certain stochastic environments. The approach proposed is neat.\n\nWeaknesses:\n1. Writing incomplete sentences, e.g. line 5: or assume Markovian context variables for?\n2. Lack of theoretical analysis to compare to distribution shift in RL which seems to be a very similar topic.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept