An Adaptive Deep RL Method for Non-Stationary Environments with Piecewise Stable Context
Abstract
One of the key challenges in deploying RL to real-world applications is to adapt tovariations of unknown environment contexts, such as changing terrains in robotictasks and fluctuated bandwidth in congestion control.
Existing works on adaptationto unknown environment contexts either assume the contexts are the same for thewhole episode or assume the context variables are Markovian. However, in manyreal-world applications, the environment context usually stays stable for a stochasticperiod and then changes in an abrupt and unpredictable manner within an episode,resulting in a segment structure, which existing works fail to address.
To leveragethe segment structure of piecewise stable context in real-world applications, in thispaper, we propose a Se gmented C ontext B elief A ugmented D eep (SeCBAD) RLmethod.
Our method can jointly infer the belief distribution over latent contextwith the posterior over segment length and perform more accurate belief contextinference with observed data within the current context segment.
The inferredbelief context can be leveraged to augment the state, leading to a policy that canadapt to abrupt variations in context.
We demonstrate empirically that SeCBADcan infer context segment length accurately and outperform existing methods on atoy grid world environment and Mujuco tasks with piecewise-stable context.
Introduction
Deep reinforcement learning has achieved great success in a wide range of challenging environmentssuch as Atari games (Mnih et al., 2013; Bellemare et al., 2012; Hessel et al., 2017) or continuouscontrol tasks (Schulman et al., 2015, 2017a).
However, in stark contrast with this trend, applying
RL to real-world applications remains a great challenge.
In most real-world settings, there could bevariations in environmental factors, such as changing terrains in robotic tasks, fluctuated bandwidth incongestion control and dynamic traffic patterns in autonomous driving.
We refer to such environmentfactors as situation or context .
The changes in context are not neglectable since context usually hasa substantial impact on transition and reward functions.
When the context is fixed and known tous, the problem is stationary and easier to solve.
However, in most realistic settings, the context isusually dynamic within an episode and unknown to us at test time.
Therefore, detecting and adaptingto variations in context is very important for RL agents to make a real impact in a wide range ofreal-world applications.
In real-world environments with varied unknown contexts, we find a typical evolving context patternparticularly interesting.
The context c usually stays the same for a stochastic period until it changesabruptly and unpredictably into another context value c ′ which is sampled i.i.d.
from some priorcontext distribution.
What’s more, we usually do not have access to context c directly but insteadhave access to some noisy observation x t sampled from some distribution p ( x t | c ) only at trainingtime (e.g., as auxiliary information from the simulator).
We take fluctuated bandwidth in congestioncontrol as an example.
In Figure 1, 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 thusthe bandwidth approximately follows the same distribution.
At the end of a segment, the networkcondition changes, and the bandwidth changes abruptly into another distribution.
Here the context crepresents the parameters that determine the distribution of fluctuated bandwidth (network condition),while the observations x t represents the observed bandwidth which follows the above distributionp ( x t | c ) .
While the contexts c are piecewise-stable, there could be slight variations in x t within eachsegment.
This piecewise-stable pattern of context dynamics is of particular interest to us since it cancapture a wide range of stochastic context processes in real-world applications, such as changingterrains in robotic tasks and fluctuating bandwidth in congestion control.
Since the change in contextis abrupt and unpredictable, we cannot predict the future context in advance.
The best we can do isdetect 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 et 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,b; Xie et al., 2021)assume the context stays the same for the whole episode and infer the context based on the entireepisode (c.f. Figure 2(a)), therefore cannot quickly adapt to context changes within an episode.
Otherworks on non-stationary RL assume intra-episode context changes and model c t at each time step, butfew study the piecewise-stable context as we do.
Nagabandi et al.
(2018) directly predict the contextand thus cannot capture the prior that the context tends to stay the same for a stochastic period.
Fenget al.
; Ren et al.
(2022) model Markovian discrete context for each time step (c.f. Figure 2(b)),therefore failing to model the non-markov property of context and the prior over context segmentlength in our setting.
Compared with existing works, our setting is distinctive since the segmentstructure is latent to us (c.f. Figure 2(c)).
Therefore, the unique challenge in our setting is that weneed to infer the segment structure, which can be further leveraged to infer belief context by onlyincorporating the relevant observed data in the current segment.
This paper studies how to infer segment structure to detect abrupt context changes and infer beliefcontext accordingly to adapt to the piecewise-stable context in RL environments.
We first introducelatent situational MDP which models RL environments with the stochastic situation / context process (Section 2).
Then, we introduce how to infer the belief context from observed data.
To addressthe challenge above, we propose to infer context segment structure and belief context jointly from observed data (Section 3.1).
Then we augment the state with inferred belief context so that the RL agentcan automatically trade-off between acting optimally conditioned on inferred context and gatheringmore information about the current context (Section 3.2).
Finally, we combine the training objectivesfor RL and inference and present all the details for our proposed deep RL algorithm (Section 3.3).
Weevaluate our algorithm on a gridworld environment with dynamic goals and MuJoCo tasks (Todorovet al., 2012) with varied contexts.
Experiments demonstrate that our algorithm can quickly detect andadapt to abrupt changes in piece-wise stable contexts and outperform existing methods(Section 4).
Our contributions can be listed as follows:
• We introduce latent situational MDP with piecewise-stable context, which can capture awide range of real-world applications (Section 2). • We propose SeCBAD, an adaptive deep RL method for non-stationary environments withpiecewise-stable context.
Our method can infer the context segment structure and the beliefcontext accordingly from observed data, which can be leveraged to detect and adapt tocontext changes (Section 3). • 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 4).
Problem Formulation
In this section, we define a latent situational Markov Decision Process (LS-MDP) as a tuple M =( S , A , C , X , G, T, R, γ ) , where S is the set of states, A is the set of actions, and γ ∈ (0 , 1] is thediscount factor.
To formulate the variations in environment factors, we introduce C , the set of latentcontexts, and X , the set of observable contexts.
C refers to the set of underlying hidden contexts thatcontains all necessary information but are left unobservable to the agent, while X refers to the setof contexts with only partial information. G refers to the segment length which will be describedin detail in the next paragraph.
X are observable only during training and unobservable duringdeployment.
In an example environment setting where a robot walking over changing terrain, Crepresents the terrain features containing perfect information, and X represents noisy and imperfectinformation like mechanical metrics of the current step from virtual sensors in a simulator and thus isonly accessible during training.
In our setting, we focus on the case where the environmental changes are abrupt and irregular, incontrast to the smoothly changing assumption on context/task in existing works in non-stationary
RL (Chandak et al., 2020a,b).
To better model the generative process of the contexts, we introduce thesegment length G .
Each episode is composed of several stationary segments with different segmentlengths.
For the ease of notation, we also introduce G t , which measures the length of the currentsegment up to time step t .
Then, the generative process can be described as follows: at the beginningof the h -th segment, the environment samples G h ∼ p G ( G ) and c h ∼ p c ( c ) from prior distributionsp G and p c .
Then, in the next G h steps, the latent context c h remains unchanged.
For each time step tin this segment, the current segment length accumulates as G t =
G t − 1 + 1 (we define G t = 1 forthe first time step in this segment), and the current latent context satisfies that c t = c h .
The agent canobserve x t ∼ p ( x t | c t ) for each time step if during training.
The stationarity lasts until the end of thecurrent segment, then the environment resamples G h +1 and c h +1 , and the process repeats.
We showtwo examples of graphical models of our setting given different G t in Figure 2(c).
The transition function T : p ( s t +1 | s t , a t , c h ) and the reward function R : p ( r t | s t , a t , c h ) are allconditioned on current latent context c h .
Therefore the changes in C lead to the changes in transitionand reward functions.
At test time, X is no longer accessible, so we need to infer context changesfrom observed transitions and rewards.
Since the latent context C remains unobservable and changessilently, the environment is no longer stationary for the agents.
To act optimally, it is important tokeep track of the environment to recognize changes in time, and rapidly adapt to those changes.
Methodology
In this section, we present our Segmented Context Belief Augmented Deep (SeCBAD) RL method andelaborate on how SeCBAD solves the challenges discussed above.
Our method consists of two maincomponents:
• Joint inference of the belief distribution over the latent context and the segment structurefrom 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 to infer the segmentstructure jointly with belief context and leverage the segment structure to remove irrelevant observeddata.
After the belief context is approximated, it is then incorporated into the state as the input of thepolicy.
We detail the policy optimization part under the belief MDP framework in Section 3.2.
Andfinally, in Section 3.3, 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 c t and current segment length G t fromobserved trajectory τ 1: t , so as to remove irrelevant data in τ 1: t for belief context inference.
This canbe formally expressed by the following equation:
[Equation] [Equation]
where τ t 0 : t = ( s t 0 , a t 0 , r t 0 · · · , s t − 1 , a t − 1 , r t − 1 , s t ) 1 , and c t − G t +1: t refers to the latent contextfor the whole segment from t − G t + 1 to t .
In Equation 1, we separately estimate the posteriorp ( c t − G t +1: t | G t , τ t − G t : t ) of latent context given known segment structure G t , and the posterior ofsegment length p ( G t | τ 1: t ) .
3.1.1 Approximate Inference for Latent Context under Known Segment Structure
In this part, we focus on estimating p ( c t − G t
+1: t | G t , τ t − G t : t ) , which is the posterior of the latentcontext under known segment structure G t .
We use the variational inference framework to approximate the true posterior.
To be specific, we use an posterior inference network q ϕ to infer the beliefcontext in segment [ t − G t + 1 : t ] : q ϕ ( c t − G t +1: t | G t , τ t − G t : t ) .
The variational lower bound for thelog-likelihood of the current segment is given by:
[Equation]
where p θ denotes the decoder and τ Xt − G t :
t = ( τ t − G t : t , x t − G t +1: t ) .
For the detailed derivation, please 144 see Appendix A.1.
The reconstruction term of Equation 2 can be factorized as:
[Equation]
which is the sum of log-probablity of transitions under context sampled from q ϕ .
The term
D KL ( q ϕ || p ) is the KL-divergence between our variational posterior q ϕ and the prior over the beliefcontext.
For the prior p ( c t − G t +1: t ) , we use previous posterior at timestep t − 1 if the context remainsunchanged at timestep t , or N (0 , I ) otherwise.
Unlike most previous methods which assume invariant context within an episode or markovian contextat each time step, our method 1) takes only observed data within current segment τ t − G t : t as input forencoder q ϕ , 2) reconstructs only data within current segment
τ Xt − G t : t at the output for decoder log p θ .
Our method is naturally motivated by the piecewise-stable assumption on context dynamics.
Byremoving irrelevant data outside the current segment that are generated by past unassociated context,our method can estimate the current latent context more accurately.
3.1.2 Iterative Inference for the Segment Length
In this part, we focus on estimating p ( G t | τ 1: t ) , which is the posterior for current segment length,given p ( c h | τ t − G t : t ) , the posterior of context under given segment structure.
To compute this posteriordistribution, we first compute the joint distribution p ( G t , τ 1: t ) recursively based on p ( G t − 1 , τ 1: t − 1 )as follows 2 :
[Equation]
The three major components in Equation 4 in turn are the previous joint distribution, the evolutionprior, and the observation probability.
The previous joint distribution is iteratively provided at thebeginning of each time step.
The evolution prior p ( G t = i | G t − 1 = k ) measures the prior knowledgeon the segment length G t given the segment length of the previous time step G t − 1 =
k , where either
G t =
G t − 1 + 1 or G t = 1 holds.
As for the observation probability which is the third term of Equation 4, it measures how likely theobservations show up given the history of the segment.
To be specific, we have
[Equation]
In Equation 5, the term p ( s t , r t − 1 | s t − 1 , a t − 1 , c t −
i +1: t ) estimates the data likelihood for the nextstate-reward pair, where c is drawn from the posterior distribution given the data in the segment beforetimestep t .
To compute the RHS of Equation 5, we sample from q ϕ ( c t − i +1: t − 1 | G t
= i, τ t − i : t − 1 ) 4which is the belief context inferred from the whole history of the segment before timestep t , and usethe sampled c the decoder to compute the data likelihood.
Given the joint distribution, the posterior distribution of G t can be derived as
[Equation] [Equation]
We can also incorporate observable contexts x in the observation probability to improve accuracyduring training.
See Appendix A.3 for more details.
Belief Context as a Mixture Distribution
Given inferred posterior of G t in Section 3.1.2, the belief of the latent context at time step t can bederived as
[Equation]
This mixed probability of c t − g t +1: t which has taken all possible segment structures into consideration,can now represent the current belief b t of the latent context c .
3.2 Policy Optimization with Belief Context
Inspired by the belief MDP (Kaelbling et al., 1998), Bayes Adaptive MDP (BAMDP) (Duff, 2002)and
recent works on meta RL as task inference (Zintgraf et al., 2020), we incorporate the inferredbelief context into the augmented state.
At each time step t , the belief latent context is approximatedvia q ϕ using Equation 7.
Therefore, we define the augmented state as ( s, b ) ∈ S × B , where S isthe same state space as in LS-MDP and B is the set of belief latent contexts.
Accordingly, we havetransition T b ( s t +1 , r t , b t +1 | s t , a t , b t ) =
p ( s t +1 , r t | s t , a t , b t ) p ( b t +1 | s t , a t , r t , s t +1 , b t ) and reward
R b ( r t | s t , b t , a t )
This definition brings advantages in the sense that the information gathering andexploitation tradeoff is no longer a problem under such augmented states, since the transition andreward functions are no longer conditioned on exact c .
Now, the policy is defined as π ( a | s, b ) amapping from the augmented state space to the action space, and the agent’s objective is to maximize
[Equation]
3.3 Algorithm and Implementations of SeCBAD
In this section, we describe the overall algorithm and implementation details of SeCBAD.
See Figure3 for an overview of our framework.
As shown in Figure 3(a), we use a GRU (Cho et al., 2014)parameterized by ϕ as the recurrent encoder q ϕ , and distributions in latent context space is assumed tobe diagonal Gaussians with mean and variance parameterized by q ϕ .
The decoder includes transitionmodel p θ ( s t +1 | s t , a t , c ) , reward model p θ ( r t | s t , a t , c ) and observable context model p θ ( x t | c ) .
Theoutput of all the decoders are Gaussian distributions with mean parameterized by feed-forward neuralnetworks and fixed identity covariance.
Then, we estimate p ( G t | τ 1: t ) using q ϕ and p θ as described in
Section 3.1.2.
As shown in Figure 3(b), we combine the belief context based on different segmentaccording to p ( G t | τ 1: t ) to get the belief b t ( c ) =
q ϕ ( c t | τ 1: t ) , where one approach is to provide atotal of t mean, covariance and weights as policy input.
For simplicity, we choose G ∗ t with highestprobability in p ( G t | τ 1: t ) and use the corresponding q ϕ ( c t − G ∗ t +1: t | τ t − G ∗ t : t ) as the belief.
Empirically,we find this approximation leads to little performance loss.
We build our RL algorithm on the top of
PPO (Schulman et al., 2017b) to learn the policy π ψ ( a t | s t , b t ( c )) , where ψ denotes the parameters inthe actor and the critic.
We use the objective described in Section 3.2 to optimize the policy.
During deployment, q ϕ and p θ are fixed.
We first compute the segment posterior p ( G t | τ 1: t ) usingencoder q ϕ and the transition and reward model p θ .
Then, we estimate the belief b t ( c ) and feed itinto the policy as input.
For more implementation details, please refer to the Appendix A.6.
Experiments
In this section, we empirically evaluate SeCBAD on two tasks.
We first demonstrate our algorithmin a grid world environment in Section 4.1 to illustrate the significance of incorporating segmentstructure during inference.
For large-scale experiments, we test the proposed algorithm in Section4.2.
The results of multiple challenging tasks show that SeCBAD outperforms baselines in terms ofperformance and sample efficiency.
In Section 4.3, we further provide case studies of agent behaviorsand learned latent to gain more insights into the results.
[Figure] 4.1 Motivating Example
In this experiment, we use a motivating example to show that it is vital for the agent to consider thesegment structure to adapt to variations of unknown environment contexts.
The agent is in a 2 × 2 grid world containing 1 goal state s ∗ and 3 other states.
In s ∗ , the agent cangain +1 reward with probability p 0 , and 0 reward with probability 1 − p 0 .
For other states, the agentcan gain +1 reward with probability p 1 and 0 reward with probability 1 − p 1 , where p 0 > p 1 .
Theactions include up, down, left, right, and none.
The specific location of s ∗ is unobservable to theagent, and may change after a stochastic period.
Hence, the optimal behavior is to search for and stayin the goal state until the change happens.
The agent needs to rapidly detect the change in the goalstate search for it again.
We use the analytical belief update rule instead of variational inference in this experiment.
To bespecific, we assume the q ϕ and p θ is known and fixed and then infer p ( G t | τ 1: t ) and the belief b t ( c ) .b
t ( c ) here refers to the belief distribution of the location of s ∗ .
We concatenate ( s, b ) to train our
RL agent using this augmented state.
In this way, we aim to focus on the effectiveness of inferringsegment structure.
We test our algorithm against two benchmarks: (1) vanilla inference, where theagent updates belief without considering segment structure, like in Zintgraf et al.
(2020); Rakellyet al.
(2019), and (2) inference with oracle information, where the omniscient agent computes b t ( c )with the exact oracle segment structure.
Figure 4 illustrates how SeCBAD and the vanilla inference baseline behave.
It is shown that the beliefb t ( c ) of SeCBAD closely matches with the actual goal state s ∗ .
Also, compared with SeCBAD, thevanilla inference baseline needs more steps to detect the changes since it needs to correct the deviatedprior belief.
This result supports that incorporating segment structure into inference like SeCBAD issignificant 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 algorithmon several challenging control tasks with varying contexts. We conduct the experiments on modified
MuJoCo (Todorov et al., 2012) tasks.
These environments are commonly used in previous works(Zintgraf et al., 2020; Rakelly et al., 2019), and we further modify the contexts to provide challengescorresponding to LS-MDP as mentioned in Section 2. As for Ant Direction and Cheetah Direction,the policy needs to move towards the target direction as fast as possible.
As for Ant Velocity and Cheetah Velocity, the policy then needs to additionally move at the speed of the target velocity besidesthe target direction. For Cheetah Goal, the policy needs to navigate to varying goals and stay there.
For Ant environments, the direction and velocity are specified on a 2-dimensional plane, and forhalf cheetah environments, the direction and velocity are 1-dimensional. In these environments,the contexts refer to the target velocity, target direction, or location of the goal. To make it morechallenging, we further add noises to contexts within each segment.
Please refer to Appendix A.6for more details.
As for baselines, we select one representative algorithm for each type in Figure2 to compare since to the best of our knowledge, few existing works study the same setting as wedo.
We use VariBAD (Zintgraf et al., 2020) to represent those methods assuming one context inthe episode, and use FANS-RL (Feng et al., 2022) to represent those methods assuming Markoviancontexts.
Since LS-MDP is a special case of POMDP, we include a POMDP baseline PPO-RNN(Hausknecht and Stone, 2015) for comparison. The results are provided in Figure 5. [Figure]
The experiment results empirically show that SeCBAD achieves superior performance and sample efficiency on challenging control tasks. As illustrated in Figure 5, SeCBAD achieves higher scoresthan baselines. We provide some insights into the results.
For the method assuming that contexts staythe same within an episode, (Zintgraf et al., 2020) uses the learned latent contexts to decode the wholetrajectory including transitions and rewards in other segments. This reconstruction mismatch maylead to averaged latent contexts so that the policy cannot act correspondingly (see detailed analysisand the ablation study in Appendix A.7.3.) For Markovian contexts baseline (Feng et al., 2022) and
POMDP baseline (Hausknecht and Stone, 2015), they can perform relatively well on tasks where thecontexts 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 significantlyoutperforms these two baselines (Feng et al., 2022; Hausknecht and Stone, 2015), which proves thatthe specially designed joint inference component in SeCBAD is effective and can help improve theperformance. To better understand the proposed algorithm, we conduct a series of ablation studies. In Appendix
A.7.2, we study the effects of inaccurate prior on SeCBAD. In Appendix A.7.4, we study theimplementation choices on how to use p ( G t | τ 1: t ) . In Appendix A.7.5, we further test SeCBADagainst different levels of noises within each segment.
4.3 Case Studies on the Segment Structure
In order to provide more insights into the results, we present a case study on Ant Direction in thissection. We show the behavior and learned latent along with the inferred segment of a randomlyselected episode during testing of the SeCBAD algorithm on Ant Direction in Figure 6.
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 learnedlatent contexts in the second row and the inferred segment structure in the third row.
From the thirdrow, it can be seen that the agent can detect the segment in time as the inferred segment structureclosely matches the ground truth segment structure. As for the segment starting from the 430thenvironment step, the change in goal direction is negligible so that the agent automatically mergesthe two segments.
Since the latent contexts are calculated according to the inferred G ∗ t , the latentcontexts in the second-row change as 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 thechanges. We provide case studies of baselines and detailed analysis in Appendix A.7.1.
Related works
Our work is closely related to non-stationary RL , where the transition and reward functionsmay change over time.
Most existing works on non-stationary RL focus on inter-episode nonstationarity (Xie et al., 2021; Chandak et al., 2020a,b; Xie et al., 2022; Sodhani et al., 2021; Alegreet al., 2021; Poiani et al., 2021; Al-Shedivat et al., 2017), some of which adopt contextual MDPas formulation (Hallak et al., 2015).
Recently there are some works considering intra-episode nonstationarity
Ren et al Kamienny et al
Kumar et al Nagabandi al Feng al al the latent context to be finite and Markovian Fenget al.
latent environment can be modeled as a factored [Figure]
MDP.
In contrast to existing works on non-stationary MDP, we assume piecewise-stable contextwith abrupt changes within an episode, which is more realistic and can capture a wide range ofreal-world applications.
To rapidly adapt to dynamic context, the agent needs to continuously performinformation gathering behavior.
Bayesian RL (Duff, 2002; Zintgraf et al., 2020; Fellows et al., 2021)is an elegant framework to optimally tradeoff the exploration and exploitation in an unknown andstationary MDP.
As a special type of Belief MDP (Kaelbling et al., 1998), Bayes Adaptive MDP(BAMDP) (Duff, 2002) maintains a belief over the environment and uses this belief to augment thestate.
Our model can be viewed as a special case of belief MDP, where we only maintain belief overlatent context to trade off between information gathering and exploitation.
To accurately infer beliefcontext, 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., 2020; Humplik et al., 2019; Poiani et al.
,2021; Zintgraf et al., 2020) or context inference (Xie et al., 2021;
Feng et al., 2022; Ren et al., 2022).
However, none of these methods suit our setting.
We perform joint inference over latent context andsegment structure from observed data, so as to remove irrelevant data for more accurate belief contextinference.
LS-MDP can also be viewed as a special case of POMDP .
Recently, progress has beenmade in learning the latent dynamics model (Krishnan et al., 2015;
Karl et al., 2016; Doerr et al.,2018; Buesing et al., 2018;
Ha and Schmidhuber, 2018; Han et al., 2019; Hafner et al., 2019b,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, ithas been shown that (Hafner et al., 2019b) introducing more structured information will significantlyenhance the performance.
In this paper, we exploit the addtional assumption of LS-MDP to inferthe context based on segment structure.
The experiments show that SeCBAD can achieve betterperformance compared with existing methods.
Conclusion
In this paper, we propose SeCBAD, a Se gmented C ontext B elief A ugmented D eep RL method todeal with piecewise-stable context in non-stationary environments.
Piecewise-stable context is quitecommon in a wide range of real-world applications.
Compared with existing methods, our methodcan automatically detect the segment structure, which reflects when the context changes abruptly.
The detected segment structure can be further used to compute context belief with only relevantobserved data.
To the best of our knowledge, this is the first method that can model and leveragepiecewise-stable context in 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 agridwold environment and mujoco tasks with piecewise-stable context.
For future work, we plan toleverage various deep learning techniques to improve SeCBAD, which includes replacing the GRUencoder by Transformer to better capture the long-term dependency in the input trajectory segment.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions 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]
If you are 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]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(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 to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are 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 a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[No]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A AppendixA.1 Bound Derivation
The variational bound for our model can be derived using importance weighting and Jensen’sinequality.
We use ˜ c :=
c t − G t +1: t for simplicity:
[Equation]
A.2 Algorithm Framework
In this Section, we provide the detailed training and evaluation algorithms of SeCBAD.
The trainingalgorithm is detailed in Algorithm 1.
The evaluation algorithm is detailed in Algorithm 2.
Input: buffer B , imagination horizon H , interacting step T , batch size B , batch length L , number of trajetories N .
Initialize buffer B with S random seed episodes.
while not converged do ▷ Parameter Optimization for c = 1 , . .
C do
Draw B data sequences { ( s t , a t , r t , x t ) } k + L t = k from B
Calculate p ( G t | τ 1: t ) using Equation (6).
Sample g ∼ p ( G t | τ 1: t ) .
Infer belief state q ϕ ( c t − G t +1: t | G t =
g, τ t − G t : t ) .
Sample k ∈ N from [ t − g + 1 , t ] .
▷ Calculate the ELBO: sample 1 reconstruction step
Predict observable context: p θ ( x k | c, G t = g )
Predict reward: p θ ( r k − 1 | c, s k − 1 , a k − 1 , G t = g ) , and state p θ ( s k | c, s k − 1 , a k − 1 , G t = g ) .
Update ϕ, θ using Equation (2).
Update ψ using Equation (8).
end for
Reset environment and get s 1 , x 1 .
for t = 1 , . . .
, T do ▷ Data Collection
Calculate p ( G t | τ 1: t ) using Equation (6).
Calculate belief b t ( c ) using Equation (7).
Compute a t ∼ π ψ ( a t | s t , b t ) with action model.
Add exploration noise to action.
Execute a t and get x t +1 , s t +1 , r t .
end for
Add experience to buffer B = B ∪ { ( s t , a t , r t , x t )
Tt =1 } end while environment for t = 1 , . .
, T do
Calculate p ( G t | τ 1: t ) using Equation (6).
Choose g = arg max p ( G t | τ 1: t ) .
Compute b t ( c ) =
q ϕ ( c t − G t +1: t | G t =
g, τ t − G t : t ) .
Compute a t ∼ π ψ ( a t | s t , b t ) with action model.
Execute a t and get x t +1 , s t +1 , r t .
end for
A.3 Include x in observation probability
In Section 3.1.2, we recursively estimate the posterior of current segment length p ( G t | τ 1: t ) , whichconditions the segment length on the observation τ 1: t as conditions to keep consistent with the casein evaluation.
However, During training, we can use more information to help us get a more accurateestimation.
For instance, we can include the observable contexts x into τ 1: t , like τ X 1: t and update thecorresponding equations as follows:
[Equation]
The difference mainly lies in the observation probability term:
[Equation]
The observation probability term now estimates the data likelihood not only for the next state-rewardpair, but also for the next observable context x t .
There are many other choices here for how to choose the information during training.
It is alsopossible to only incorporate the observable contexts x instead of using τ X 1: t for simplicity.
We can stilluse the online inference method like we stated above, or we can use some more simple and heursticmethods to derive the posterior.
Since we can directly access x from the simulator, it is possible forus to calculate the posterior in advance, which can help to save a lot of computation during trainingtime.
A.4 Discussion with Related Works
Recently, some works (Ren et al., 2022; Feng et al., 2022) assume intra-episode context changes andmodel context c t at each time step, more specifically with Markovian transition on latent contexts.
Since Markovian assumptions are made in latent context space, these methods are theoreticallyable to capture the non-Markovian pattern in observable context space.
However, as discussedin our paper, one of the key challenges in deploying RL to real-world applications is to adapt tovariations of unknown environment contexts that stay stable for a stochastic period.
In contrastto general non-stationary environments, our problem setting assumes piece-wise stable context asspecial problem structure, which can be further exploited to improve performance against methods forgeneral non-stationary environments.
By contrast, we propose to jointly infer the segment structureas well as the segment latent context.
By doing so, we can avoid solving non-stationary environmentsin general case.
In our paper, we implement a one-step VariBAD to represent the methods thatmodel context c t at each time step.
For each time step, we use the whole trajectory as input just like
VariBAD.
However, instead of decoding the whole trajectory, we only make the agent to decode thecurrent time step.
A.5 Bayesian Update Rule
In Section 4.1, we use the Bayesian belief update rule to calculate the belief distribution.
Let Kdenotes the number of grids, then b
( θ ) ∈ R K is a simplex which refers to the probability location ofthe goal state and therefore a categorical distribution.
As for the belief update rule, we have
[Equation]
In above equation, we omit the transition term since the transition keeps stationary across theepisode.
As for the reward term p ( r | s, a, θ ) , it can be derived through the definition.
For the k -thgrid, suppose the observed state is the k ′ -th grid.
Then, if k = k ′ , we can get the reward term asp ( r = 1 | θ = k, s = k ′ , a ) =
p 1 and
A.6 Hyper parameters and Implementation Details
Network Architecture For Half-Cheetah environments, we use a GRU(Cho et al., 2014) with 128units as the dynamics model of the encoder q ϕ .
For each time step, q ϕ receives an encoded state,action and reward as input.
The encoded state, action and reward are the output three single layer fullyconnected networks of respective size [16 , 32 , 16] with ReLU as activation function.
We assume thelatent distribution are 5 -dimensional Gaussians with predicted mean and standard deviation.
As forthe transition model p θ ( s t +1 | s t , a t , c ) , reward model p θ ( r t | s t , a t , c ) and observable context modelp θ ( x t | c ) , the output of all the decoders are Gaussian distributions with mean parameterized by fullyconnected network (with hidden size [64 , 32] and ReLU activation) and fixed identity covariance.
Asfor the Ant-Cheetah environments, we use 10-dimensional Gaussians as the latent distribution.
Theother settings are kept the same as in Half-Cheetah environments.
As for the policy training part, we adopt PPO (Schulman et al., 2017b) to learn the policyπ ψ ( a t
| s t , b t ( c )) .
The actor and critic are parameterized by fully connected network of size [128 ,
128]with
Tanh as activation function.
Training Details
The inputs of both actor and critic is s t and b t .
To parameterize the distributionb t , we use its mean and standard deviation.
For simplicity, we only incorporate one possible segmentlength g into consideration and this approximation will not cause obvious performance drop accordingto our test.
During training, we randomly sample g from the posterior, while during evaluation, weuse the most possible g .
Then, we calculate the corresponding belief b t using g , and use the mean andstandard deviation as the inputs.
During training, we use the posterior p ( G t | x 1: t ) which conditionsonly on the observable contexts to simplify the computation.
During evaluation, we then use theposterior p ( G t | x 1: t ) like we stated in Section 3.1.2 since we can no longer access the training onlyinformation x .
Hyper parameters We train PPO with Adam optimizer with learning rate 7e-4, max gradient norm0.5, clip parameters 0.1, value loss coefficient 0.5, and entropy coefficient 0.01.
We use differenttraining schedule between the VAE part (encoder & decoder) and the policy part.
We use two Adamoptimizers with different learning rates.
The VAE optimizer uses a learning rate of 1e-3, and thepolicy optimizer uses a learning rate of 7e-4.
For the full details of hyper-parameter settings, pleaserefer to the code repository that we will publish soon.
Environment Details We give detailed descriptions of the environments we used.
Ant Directionenvironment is built upon the well-known Ant Mujoco environment.
For each environment step, weset a target direction in the 2D horizontal plane.
The reward function is given as v f − 0 .
2 v v , wherev f denotes the agent velocity along the target direction, v v denotes the agent velocity that is verticalto the target direction and is always non-negative.
Such reward function pushes the agent movesfaster along the target direction while penalizing the agent for the velocity component that is verticalto the target direction.
In experiment the target direction is chosen uniformly from [0 , 2 π ] .
For Ant Velocity , we set a target velocity v t =
( v t,x , v t,y ) in the 2D horizontal plane.
Denotethe agent velocity as v a =
( v a,x , v a,y ) .
We project the agent velocity v a to v t to get the forwardvelocity component v f = v t · v a / ∥ v t ∥ 2 .
The velocity component perpendicular to the targetvelocity is given by v v = (cid:113) ∥ v a ∥ 22 − v 2 f .
If forward velocity v f < ∥ v t ∥ 2 , we compute the reward asv f − 0 .
3 v v , meaning that we expect the agent moves faster along the direction of task velocity, butnot along the direction vertical to task velocity.
Otherwise when v f ≥ ∥ v t ∥ 2 , we compute reward as− v f + 2 ∥ v t ∥ 2 − 0 .
3 v v .
Observe that if ignore the term − 0 .
3 v v in the reward function, we assign thelargest reward when v f = ∥ v t ∥ 2 , and reward decreases as v f deviating from ∥ v t ∥ 2 , thus promotingthe agent follows the target velocity.
In experiment, v t x and v t y are independently chosen from
Uniform ( − 3 , 3) for each segment.
Cheetah Direction , Cheetah Goal and Cheetah Velocity are modified from Half Cheetah Mujocoenvironment.
The agent can only move along x -axis.
In the original version of Half Cheetah, thereward is given by the agent velocity along the positive x -axis.
Denote agent velocity as v a andagent location as x a .
In Cheetah Direction, we set the target direction to positive x -axis or negativex -axis, and the reward is v a if the target direction is positive x -axis or − v a otherwise.
We choose thepositive and negative directions with equal probability.
In Cheetah Goal, we set a target location x ton the x -axis in each environment step, and the reward function is −| x a − x t | , with x t sampled from
Uniform ( − 5 , 5) .
The agent location x a is also included into the observation space, which is not forthe original Half Cheetah, Cheetah Dir and Cheetah Velocity.
For Cheetah Velocity, we set a targetvelocity v t on the x -axis, the reward is −| v t − v a | , and v t is chosen from Uniform ( − 5 , 5) .
A.7.1 A Case Study on Ant Direction
[Figure]
P
[Figure]
s i t
i o
[Figure]
[Figure]
D i r e c t i o n
[Figure]
[Figure]
[Figure]
F o r w a r d v e l o c i t y
[Figure]
[Figure]
[Figure]
L a t e n t m e a n
[Figure]
To better understand the outperformance of SeCBAD, we show the behavior of SeCBAD and otherbaselines and the learned latent contexts in Figure 7.
The first row shows the agent’s location through time.
It can be seen that the agent of PPO RNN,
VariBAD, and Factor MDP choose to stay around a fixed position shortly after starting, while
SeCBAD can successfully guide the agent moving around.
In the second row, we show the agent’s direction along with the target (task) direction.
The orangecurve shows the time-varying task direction, and the blue curve shows the agent’s actual direction.
Since the goal is to maximize the velocity along the task direction, we plot the velocity along thetask direction in the third row.
It can be concluded that SeCBAD is able to detect and adapt to thevariations rapidly, and the agent is indeed moving along the task direction, while other methodsfail to detect and adapt to the non-stationary task in time.
After detecting the teak direction change,
SeCBAD’s velocity drops and then grows up as it tries to change the direction, while the velocities ofother baselines are close to zero and the learned behaviors are undesirable.
We provide some insights into the results.
We plot the mean of the latent in the fourth row of
Figure 7.
VariBAD (Zintgraf et al., 2020) tries to learn a single context for the whole episode andtrain the decoder by reconstructing the whole trajectory.
This may make the agent learn the samecontexts across the episode.
Therefore the learned latent mean is smooth and thus uninformative tothe changing environment.
Please refer to Appendix A.7.3 for detailed analysis.
FANS-RL (Fenget al., 2022) assumes that the contexts evolve in a Markovian pattern.
PPO-RNN (Hausknecht and
Stone, 2015) treats the problem in a general form as a POMDP and leverages little information aboutthe context pattern.
It is hard for these two methods to learn a meaning for contexts, especially incomplex environments like Ant Direction.
However, as for SeCBAD, thanks to the joint inferenceframework and segment decoder, the learned latent mean can change abruptly when the inferredsegment structure changes, which informs the policy that the contexts have changed.
The behavior in
Figure 7 shows that the agent can learn to adapt to the variations rapidly with this accurate enoughbelief.
[Figure]
Ant Direction
SeCBAD + Prior
SeCBAD + Prior
A.7.2 Prior Analysis
In this Section, we provide the analysis on the choice of prior p ( G t = k
| G t − 1 = i ) on Ant Direction.
The results are provided in Figure 8.
We choose two different prior functions.
For prior 1 , welet p ( G t = 1 | G t − 1 = i )
= 180 for any i and p (
G t
=
i + 1 | G t − 1 = i )
=
For prior 2 , weuse a more accurate prior by rolling out the generative process several times and approximate thep ( G t = k
| G t − 1 = i ) from the simulated data.
The performance using these two priors are shownabove.
[Figure]
It can be concluded that, two prior funciton have little performance gap.
To gain more insights, wedraw the probablity of p ( G t | τ 1: t ) in Figure 9.
The results show that the inferred posterior is verydeterministic, and the one with approximate prior is a little bit more blurry, but still deterministicenough.
This result proves that, the observation term dominates the probability term and our methodis robust with respect to the choice of prior.
A.7.3 Ablation study on the number of segments
In this section, we study how the number of segments affects the performance of SeCBAD and
VariBAD (Zintgraf et al., 2020).
As analyzed in Section 4.2, we hypothesize that methods assumingthat contexts stay the same within an episode like (Zintgraf et al., 2020) may end up learning anaveraged latent contexts.
We argue that this is not related to the length of the segment but is due tothe number of segments in an episode.
During reconstruction, (Zintgraf et al., 2020) uses the learnedlatent contexts at timestep t to decode the transitions and rewards in the whole trajectory.
Supposethere are n segments in the episode, then the probability of reconstructing the correct transition and
[Figure]
reward (i.e., those in the same segment) is only 1 n .
The noisy even erroneous training signal mayexacerbate as n grows and end up in learning averaged latent contexts for the whole episode.
To look further into this hypothesis, we conduct an ablation study on the number of segments.
We choose the Ant Direction environment and change the max episode length L ∈{ 100 , 200 , 300 , 400 , 500 } while letting the average segment length be fixed.
Then the expectednumber of segments n may vary.
We plot the averaged reward per step instead of the accumulatedrewards to cancel out the effect of L and study the effects on n .
As illustrated in Figure 10, for SeCBAD, n has little effect on performance.
However, for VariBAD,the performance deteriorates as the grows.
The results fit the above analysis and show the significanceof joint inference on the segment structure and context belief.
A.7.4 Ablation study on using p ( G t | τ 1: t )
[Figure]
use the inferred p ( G t | τ 1: t ) .
During the implementation, we may face a choice on how to use the inferred segment structurep ( G ∗ t | τ 1: t ) .
Since there are t possible choices for p ( G t | τ 1: t ) in timestep t , it’s hard to directly usethe full distribution.
One option is to use the G t with the highest probability.
Another option is tosample G t ∼ p to approximate the full distribution.
In this section, we perform an ablation studyon Ant Direction to show the performance of the two options.
We keep other parameters fixed andonly modify the way using G ∗ t .
As shown in Figure 11, the performances of the two options are veryclose.
Therefore, for stability, we choose the G ∗ t with the highest probability in experiments.
[Figure]
A.7.5 Ablation study on context robustness
In many real-world applications, the contexts are generally piecewise stable.
However, within eachsegment, the contexts might be slightly noisy.
Therefore, it is important to test the robustness of theproposed algorithm against the noise within each segment.
In this section, we conduct an ablationstudy on SeCBAD on Ant Direction by adjusting the standard deviation of the noises.
For each segment, we uniformly sample the mean of the contexts.
Then, for each timestep within eachsegment, we assume the contexts follow a Gaussian distribution.
We adjust the standard deviation ofthe Gaussian distribution to model different levels of robustness by setting σ = { 0 . 01 , 0 . 05 , 0 . 1 , 0 . 2 }and keep other parameters fixed.
As illustrated in Figure 12, SeCBAD is able to handle noises withineach segment since the performances of different σ are very close.
This further exhibit the ability of
SeCBAD to solve many real-world applications.