Variance-Seeking Meta-Exploration to Handle Out-of-Distribution Tasks
Abstract
Meta-Reinforcement Learning (meta-RL) yields the potential to improve the sample efﬁciency of reinforcement learning algorithms.
Through training an agent on multiple meta-RL tasks, the agent is able to learn a policy based on past experience, and leverage this to solve new, unseen tasks.
Accordingly, meta-RL promises to solve real-world problems, such as real-time heating, ventilation and air-conditioning (HVAC) control without accurate simulators of the target building.
In this paper, we propose a meta-RL method which trains an agent on ﬁrst order models to efﬁciently learn and adapt to the internal dynamics of a real-world building.
We recognise that meta-agents trained on ﬁrst order simulator models do not perform well on second order models, owing to the meta-RL assumption that the test tasks should be from within the same distribution as the training tasks.
In response, we propose a novel exploration method called variance seeking meta-exploration which enables a meta-RL agent to perform well on complex tasks outside of its training distribution.
Our method programs the agent to prefer exploring task dependent state-action pairs, and in turn, allows it to adapt efﬁciently to challenging second order models which bear greater semblance to real-world problems.
Introduction
Reinforcement learning (RL) algorithms aim to learn a policy that maximises the reward in one speciﬁc environment.
The coalescence of deep learning with RL has signiﬁcantly improved RL’s ability to solve diverse and challenging problems [21, 14, 2].
However, in the learning phase, an
RL agent still requires extensive exploratory interactions to learn a good policy.
These exploratory actions can be dangerous and can have high associated costs in the real-world, which means it is currently impractical to train RL agents directly in a real-world environment.
Meta-RL methods aim to improve the sample efﬁciency of RL by training with an objective to quickly perform well on many related environments [24, 6, 19].
During the training phase, the agent is trained on a number of meta-training tasks sampled from a training task distribution.
At meta-test time, the agent’s learning performance is measured on an unseen task.
Critically, meta-learning methods assume that the training and testing tasks are always sampled from the same distribution.
Therefore, meta-RL may fail to generalise to tasks that lie far outside the training distribution [12, 13].
As the real-world lies outside the training distribution by construction, this assumption means Meta-RL agents may fail to adapt to real-world problems.
We show that this problem appears in practice, when we try to apply Meta-RL to price-adaptive control of heating, ventilation and cooling (HVAC) systems.
Agents trained on simple ﬁrst-order temperature dynamics fail to generalise to more complex and out-of-distribution (OOD) second-order
Deep Reinforcement Learning Workshop (NeurIPS 2021) simulation models.
We hypothesize that one of the reasons that meta-RL fails in OOD tasks is because the exploration strategies learnt by the agent overﬁt on the training task distribution.
In response to this problem, we propose an exploration method called variance-seeking meta-exploration .
This method learns to identify states and actions with high variance in their returns across tasks, and then focuses on exploring these states to reduce uncertainty when faced with a new task.
We demonstrate empirically that this exploration method improves performance of Meta-RL on OOD tasks.
Preliminaries
We deﬁne a probability distribution p ( M ) over a range of possible tasks modeled as Markov Decision
Processes [18, MDPs] M i = (cid:104) S, A, R, T i (cid:105) .
Each MDP within this distribution has the same state space s ∈ S and action set a ∈ A , as well as the same reward function R ( s t , a t ) → R .
The state transition function T i ( s t , a t , s t +1 )
→ Pr( s t +1 | s t , a t ) varies across MDPs according to a latent variable ψ identifying the task, however, ψ i the MDP identiﬁer for M
i is kept hidden from the agent at all times.
We use the terms MDP and task interchangeably throughout the paper.
In standard RL, an MDP-speciﬁc policy π ψ i ( s t ) → a t is learned which maximises the discounted reward objective for a given discount factor γ and episode length H ,
[Equation]
for ﬁxed
[Equation]
[Equation]
Policy π ψ i induces Q π ψi ( s t , a t ) , the expected discounted sum of rewards an agent receives in state s t by taking action a t and then following policy π ψ i = arg max a Q π ψi ( s t +1 , a ) .
One way to learn an optimal control policy is Q -learning [26], which incrementally updates the current value estimate Q ω with step size α to minimise the magnitude of the 1-step temporal difference (TD) error
δ [23],
[Equation]
In practice, state space S is often too large to track values of Q ω exactly.
Deep Q-Networks [14,
DQN] solves this problem by using a deep neural network to approximate values of Q ω , minimising the magnitude of the TD error in eq.
(2) using stochastic gradient descent.
2.1 Memory-based Meta-RL
In meta-RL the objective is to learn a global policy π which is able to maximise the reward for all
MDPs sampled from p ( M ) , under the same objective (1) for each sampled M i ,
[Equation]
Maximising (3) forces the agent to learn a policy π which can adapt within H time steps to maximise the reward.
In doing so, the agent must learn the underlying task structure for MDPs in p ( M ) , to efﬁciently handle new tasks sampled from p ( M ) .
In other words, the agent learns how to learn.
Memory based meta-RL [24, 3] uses a recurrent neural network (RNN) based policy to maximise eq.
(3) The reward and action at the previous step ( r t − 1 ) and ( a t − 1 ) respectively are appended to the state space of the agent.
The RNN policy selects action a t , conditioned on the hidden state h t of the RNN, in addition to the state s t .
Therefore, the policy becomes π ( a t | h t , s t ) .
Intuitively, the recurrent update of the the hidden state can be understood as an approximation of the hidden MDP identiﬁer ψ i .
Instead of explicitly inferring ψ i , the agent learns to infer an internal representation of ψ i as h t , given agent’s experience up till t − 1 in the current MDP.
In addition to identifying the MDP, the agent also learns a policy π ( a t | h t , s t ) that acts optimally according to it’s belief ( h t ) about the current MDP at time t .
Both the MDP inference procedure and the policy are learnt simultaneously by optimising the RNN during meta-training.
We use the terms RL 2 and memory based meta-RL interchangeably in this paper.
2.2 Meta-RL for HVAC control
Price-adaptive control of heating, ventilation and cooling (HVAC) systems aims to reduce the energy consumption of HVAC systems during periods of peak demand.
By doing so, the reliance on nonrenewable sources, which is maximum during peak periods, can be reduced.
We cast this HVAC control problem as an RL problem, where the agent controls a central cooling and heating system of different sized buildings via thermostat setpoints (Figure 1).
As part of this task, the agent must pay for its energy usage according to real-time energy prices.
Because of the high costs of sub-optimal actions, training an RL agent directly in the real world is infeasible.
One solution to this problem is to train the agent on a carefully designed accurate simulator of the target building so as to avoid real-world exploration [1, 27, 10].
This, however, is highly time-consuming and expensive, making it infeasible to deploy on a wider scale.
In response, we propose a meta-RL solution which can efﬁciently and cost-effectively train an agent to adapt to the dynamics of any real-world building.
More speciﬁcally, we use meta-RL to train an agent on a number of easy to develop ﬁrst order dynamics models [15].
As a result, the agent learns to adapt to the new and unseen dynamics of a building.
In doing so, our proposed method dispenses with the need to construct accurate simulators of the target building.
We test the agent’s ability to adapt to real world buildings by measuring its performance on much more realistic and expensive second order dynamics models [7].
In this case, the distribution of all possible ﬁrst order model corresponds to the meta-training task distribution p ( M ) .
The latent parameter ψ that generates MDPs in p ( M ) corresponds to the parameters: time to cool and time to heat of the ﬁrst order model.
At test time, the test MDP is not sampled from p ( M ) and is instead sampled from an unseen distribution of more complex second order models of real world buildings.
This allows us to test the agent’s ability to adapt to out-of-distribution tasks.
The state-space of the agent consists of: the room temperature, the outside temperature, the current energy price, the energy price forecast for next ﬁve time periods, and the action and reward for previous step (for meta-RL).
At each step, the agent decides the target temperature setpoint for the heating system.
So the discrete action-space consists of natural numbers between 0 and 40.
For the price data, we use publicly available 15-minute dispatch prices from the Australian electricity market operator [16].
Exploration using task-dependent state-action pairs
To address the transfer problem identiﬁed in the previous section, we propose a novel exploration objective term to bias exploration.
The key idea behind our exploration objective is that to identify the
MDP; the agent focus on exploring upon MDP identiﬁer ψ i First, we motivate our objective using a simple grid navigation task.
Then, we show how to use our exploration objective to improve meta-RL exploration.
3.1 Motivation behind exploration objective
In theory, Q-learning learns a policy for a single MDP.
However, in this case, we train the agent using data from many randomly sampled MDPs from p ( M ) .
Because for the same state and action,
Q ( s t , a t ) can have different value across MDPs in p ( M ) , the agent cannot learn a ﬁxed Q ( s t , a t ) that holds true for all MDPs.
However, we show via an example that after sufﬁcient training, the variance of TD error differs across the state-action pairs.
We use a simple grid navigation task to demonstrate how the variance of TD error across MDPs differs in state-action pairs.
In this task, the agent’s goal is to navigate to a goal in a 5 × 5 grid within 15-time steps.
However, the goal is randomly assigned to one of the four possible goal locations in the top right corner.
This task is a modiﬁcation of the navigation task proposed by [28].
p ( M ) in this example includes only four MDPs where each possible goal location gives a new MDP.
The MDP identiﬁer ψ i here represents the goal location.
We train tabular Q-learning on MDPs sampled according to p ( M ) .
ψ i remains hidden for the agent at all times.
Upon convergence, we measure the variance of TD error across all tasks.
The results are shown in ﬁgure 2, TD-error variance of the state-action pairs near the possible goal locations is highest.
This experiment demonstrates that the variance of TD error differs across state-action pairs.
As demonstrated above in our grid navigation example, the TD error varies only for certain stateaction pairs.
Because the Q ( s t , a t ) is learnt using the TD error in Eq (3), the Q ( s t , a t ) of only a few state-action pairs will be updated when in a new MDP.
This means that the Q ( s t , a t ) of most state-action pairs will stay constant across MDPs.
We call the state-action pairs whose Q ( s t , a t ) are independent of MDP identiﬁer ψ i , task invariant state-action pairs.
The state-action pairs whose
Q ( s t , a t ) do depend on ψ i and therefore have a high variance of TD errors across tasks are called task dependent state-action pairs.
During the exploration steps, when a meta-RL agent infers the underlying MDP, the task dependent state-action pairs can provide useful data.
On the other hand, task invariant state-action pairs tell the agent little about the underlying MDP.
For example, in the grid navigation task, the high TD error variance state-actions are close to the possible goal location.
The agent must visit these high variance state-action pairs near the possible goal locations to infer the underlying MDP.
On the other hand, any state-actions far off from the goal locations cannot be used to infer the MDP.
This shows that the high TD error variance of a state-action pair is a good signal for MDP inference.
Duan et al.
[3] show that memory based meta-RL correctly learns which states to explore for MDP inference in a maze navigation problem similar to our grid example.
However, in real-world tasks such as the HVAC problem, the state-space can be large and continuous, meaning it is impossible to say which state-action pairs the meta-RL agent will learn to use for MDP inference.
Because meta-RL agents learn to maximise the discounted reward over MDPs in p ( M ) , they learn to use the state-action pairs that allow MDP inference only for the tasks within p ( M ) .
The state-actions that the agent learns to use for MDP inference may not necessarily be task dependent.
Moreover, as we show in our experiments, the learnt state-actions struggle to help in MDP inference in OOD tasks.
To avoid learning state-actions that only allow MDP inference for tasks within p ( M ) , we propose biasing the agent’s inference steps towards the task-dependent state-action pairs.
In doing so, we hope that the agent will use task-dependent state-action pairs for MDP inference instead of learning the state-actions that only allow MDP inference for within distribution tasks.
The simple grid navigation example above has only a small number of state-action pairs, and the variance of each state-action pair can be explicitly calculated and stored in a table.
However, in most problems like our HVAC problem, the state space is continuous, and the variance of each ( s t , a t ) just like the Q ( s t , a t ) needs to be approximated.
To approximate the TD error variance for continuous states, we use a separate a function approximator f θ , which learns to approximate f θ ( s t , a t )
≈ δ ( s t , a t ) , where δ ( s t , a t ) is the TD error of a ﬁxed DQN for state-action pair ( s t , a t ) .
f θ ( s t , a t ) minimises the squared error of its predictions.
The objective it minimises is
[Equation]
If the the variance of TD error across MDPs is high for ( s t , a t ) , approximator f θ ( s t , a t ) will struggle to learn a stable mapping and therefore have a high error for ( s t , a t ) .
On the other hand, if TD error variance is low for ( s u , a u ) , f θ ( s u , a u ) can learn to predict the ﬁxed TD error correctly (given enough data) and therefore minimise its error.
We use the mean squared error of f θ ( s t , a t ) to approximate the TD error variance for continuous states.
Following the work of [22], we assume that memory based meta-RL or RL 2 performs MDP inference in its ﬁrst n steps in a MDP. Therefore, we formulate a training objective that encourages the agent to collect data from task dependent state-action pairs in it’s initial interactions with the MDP. The individual components involved are summarised as follows:
• A DQN d ω ( s t , a t ) , parameterised by ω ; • A function approximator f θ ( s t , a t ) , predicting the TD error of d ω ( s t , a t ) .
f θ is implemented as a deep neural network; and • A memory-based meta-RL policy π φ ( a t | h t , s t ) .
To encourage the agent to visit task dependent state-action pairs in its ﬁrst n steps, we transform the agent’s reward during ﬁrst n steps.
We reward the agent for visiting those state-actions where the squared error of f θ ( s t , a t ) is maximum, as shown below
[Equation]
where R is the environment reward and κ and ν are tunable hyperparameters.
To incentivize exploration in task dependent state-action pairs, the agent is trained to maximise the augmented reward (cid:101) R in the ﬁrst n steps.
The agent then maximises environment reward R for the rest of the episode.
Next we describe the steps involved in training the agent.
Algorithm 1 shows the complete pseudo-code for training.
Firstly, the DQN d ω ( s t , a t ) is trained using data from MDPs sampled from p ( M ) .
Once the training loss of the DQN stabilises and stops decreasing, the weights of the DQN are frozen for all the next steps.
The purpose of DQN is only to identify the high TD error variance states.
Input: Training distribution of MDPs p ( m )
Initialise: DQN d ω ( s t , a t ) , Two replay buffers β d and β f , TD error approximating network f θ ( s t , a t ) , meta-RL policy π φ ( a t | h t , s t ) 1: while not done do (cid:46)
Training the DQN 2:
Sample a random MDP M i from p ( M ) 3:
Rollout the DQN policy d ω ( s t , a t ) in M i 4:
Store the resulting trajectory in replay buffer β d 5:
Sample a random trajectory τ i from β d 6:
Update DQN weights ω using τ i 7: end while 8: return ω (cid:46)
The DQN weights after training 9: while not done do (cid:46)
Training the meta-RL policy 10:
Sample a random MDP M i from p ( M ) 11: for each step t in 0 , 1 , ..., H − 1 do 12:
Select action: a t = π φ ( a t | h t , s t ) 13:
Advance MDP: s t +1 , r t +
= M i ( a t ) 14:
Compute TD error of DQN: δ ω ( s t , a t ) using Eq (4) 15:
Compute the variance of TD error: r error using Eq (8) 16: if t <= n then : 17:
Augment reward: ˜ r t +1 =
κr t +1 + νr error using Eq (9) 18: else 19:
Set reward to environment reward ˜ r t +1 =
r t + 20: end if 21:
Store ( s t , a t , δ ω ( s t , a t ) ) in β f 22: end for 23:
Update meta-RL policy weights φ to maximise ˜ r using Eq (1) 24:
Update f θ ( s t , a t ) weights θ using entire β f by minimising L ( θ ) in Eq (7) 25: end while 26: return φ
In the next phase, we train the function approximator f θ ( s t , a t ) simultaneously with the meta-RL policy π φ ( a t | h t , s t ) .
During early training steps, the error of f θ ( s t , a t ) is high because of small training data.
Because the policy π φ ( a t | h t , s t ) gets rewarded for visiting the states where the error of f θ ( s t , a t ) is maximum, it will maximise the data collected from these state-actions.
As more data becomes available for training f θ ( s t , a t ) , it’s error begins to reduce.
The error continues to decrease until only the high TD-error variance state-actions are left to be learnt.
Because the learning target for f θ ( s t , a t ) has a high variance, the error continues to stay high even when π φ ( a t | h t , s t ) focuses on collecting data only from these state-action pairs.
Towards the end of training we expect an equilibrium, where f θ ( s t , a t ) cannot reduce it’s error further because of high variance in it’s target values (ﬁgure 3b).
The meta-RL policy on the other hand cannot get a higher reward because the error of f θ ( s t , a t ) is lower for all state-action pairs except those with a high TD error variance.
This concurrent training regime pits f θ ( s t , a t ) and π φ ( a t | h t , s t ) against each other.
This is similar to a minimax two player game.
Experiments
To evaluate our method, we compare our exploration strategy against state of the art Meta-RL algorithms on its ability to generalize to out-of-distribution tasks in the HVAC domain.
Algorithms We compare our method against the two memory based meta-RL methods in literature.
The ﬁrst method is the original memory-based meta-RL method proposed by [24, 3].
By maximising the RL reward in Eq (1), this method learns a Bayes optimal trade off for exploration and exploitation [17].
If we do not augment the environment reward in the ﬁrst n steps, our method reduces to this original method.
The second method we compare against is E-RL 2 [22].
This method has shown strong improvements over RL 2 for tasks requiring complex exploration.
By setting the environment reward in the ﬁrst n steps to 0 during the gradient update, the gradient computation of the RL objective [Figure] in Eq (1) explicitly assigns credit for the ﬁrst n exploratory steps.
Our method becomes equivalent to
E-RL 2 if we augment the reward to 0 for the ﬁrst n steps.
Similar to E-RL 2 , our method continues to receive the original environment reward at each step including the the ﬁrst n steps of the policy roll out.
We also compare the three agents against an Oracle, which is a RL agent trained speciﬁcally on the test MDP until convergence.
The oracle was trained using the IMPALA algorithm
[4] because of its strong demonstrated performance in challenging tasks using a single machine.
We closely followed the training method and hyperparameters of [25] and trained all our agents using
IMPALA.
Similar to [24], all our agents used LSTM [8] instead of a RNN.Importantly, in our method, because we want to emphasise on the error signal from f θ ( s t , a t ) we set κ in in Eq (7) to 0.05 (close to 0).
While ν in set to 0.55, which we found to be the good in terms of ﬁnal reward.
Tasks To test the ability to generalise to out of distribution tasks, we compare the performance of the three methods on the HVAC problem.
The three methods were trained on the MDPs sampled from a distribution of 1st order models.
The ﬁrst order models in the training distribution were generated by randomly varying the time to heat ( α heat ) and time to cool ( α cool ) parameter of the model.
At the start of each episode, a new MDP was randomly sampled from this 1st order training distribution.
All three methods converged around the same number of steps.
To allow a fair comparison between methods, we trained all agents for exactly 410 episodes.
We chose 410 episodes, because beyond this, the three agents showed no improvement in average reward.
The number of exploration steps in each episode was set to n = 40, which is around 3.5 hours of exploration in the real-world.
All hyper-parameters were ﬁxed across the three agents.
Evaluation Currently the meta-RL literature lacks a consistent evaluation procedure [5].
A number of studies report the meta-training performance [19, 5, 22].
However, to report the performance on out-of distribution tasks we followed the evaluation procedure of [25].
At test time, the weights of all agents were ﬁxed and the reward after n steps in each MDP was recorded.
This was necessary because both E-RL 2 and our method are trained to maximise the reward after n steps.
We report the mean and median rewards of all agents across 30 random seeds in Table 1, where each random seed was evaluated for 100 episodes in a ﬁxed MDP.
Both the 1st order and the 2nd order evaluation tasks were held ﬁxed for all random seeds.
4.1 Results and discussion
On the ﬁrst order (within distribution) MDPs, E-RL 2 achieves the highest reward, slightly higher than RL 2 (Table 1).
We hypothesise that the strong performance of E-RL 2 on within distribution tasks can be attributed to a better exploration strategy learnt in comparison to RL 2 .
Our method performs worse than both E-RL 2 and RL 2 .
This is expected because our method unlike RL 2 does not learn a Bayes optimal exploration strategy.
Instead our augmented reward biases the agent to visit high TD error variance states, which may not be the best for MDP inference within this speciﬁc task [Table] distribution.
On the other hand RL 2 and E-RL 2 optimise only the RL reward.
Therefore they are capable of learning state-actions that allow the best MDP inference for the speciﬁc task distribution.
However, despite showing strong results on within distribution tasks, both RL 2 and E-RL 2 fail to replicate their successes in OOD second order tasks.
We hypothesise that this is because the state-actions that they rely upon to make MDP inference are learnt speciﬁcally for the training task distribution.
Therefore, when faced with an OOD task, these learnt state-actions may fail to aid the MDP inference process.
On the other hand our exploration method does not directly learn the state-action pairs for MDP inference.
Instead, our augmented reward biases the agent to use high TD error variance state-actions to make MDP inference, allowing for better MDP inference (close to the
Oracle) in the OOD tasks than the state-actions learnt directly by meta-RL.
Related work
Hochreiter et al.
[9] use RNN’s for meta-learning in supervised learning problems.
Building on the work of [9], [24, 3] use RNNs for meta-learning in reinforcement learning problems.
Ortega et al.
[17] cast memory based meta-RL in a Bayesian framework and show that the meta-RL strategies learnt are approximately Bayes optimal.
To maintain an explicit belief over the possible MDPs,
Varibad [28] learns a variational auto-encoder (VAE) to infer the MDP.
The policy is conditioned on the VAEs output and learns an approximate Bayes optimal behaviour.
Training the VAE makes this method signiﬁcantly slower in training than memory-based methods [19], while the performance at test time is similar.
Therefore we used memory-based meta-RL in our work for faster training time.
Moreover, in comparison to other meta-RL methods such as MAML [6] and PEARL [19], we chose memory based meta-RL for this work because of the low computational requirements when deployed in the real-world.
A number of studies have applied RL to control HVAC systems.
Ruelens et al.
[20] use ofﬂine-RL on the predicted energy prices and demand for residential buildings.
However, the policy learnt ofﬂine could be sub-optimal due to any errors in the forecasting model.
This is in contrast to a meta-RL solution where the agent learns to adapt its policy to deal with any errors in the forecast.
To minimise expensive exploration in the real-world, [10] ﬁrst learn a simple linear dynamics model of a data centre from the data collected during exploration, and then use this learnt model to train the agent.
This model based RL approach can be quite expensive in particular if the underlying dynamics model is non-linear.
In comparison our model-free meta-RL approach makes no assumptions about the underlying dynamics model of the target building.
Other studies [1, 27, 11] use RL for HVAC control in accurate simulation models of the target building.
However, developing accurate simulators is expensive and time-consuming.
Our method instead uses simple and easy to develop ﬁrst order models to learn RL policies that can generalise to more complex environments including the real-world.
Conclusions and future work
In this paper, we propose training a meta-RL agent on cheaper ﬁrst order models to overcome the signiﬁcant problems associated with developing expensive simulators of target buildings for HVAC control.
Our experiments show, however, that meta-RL fails to deliver strong performances in more complex second order models.
This is problematic because the real-world is much more complex than the ﬁrst order models.
To allow meta-RL to accomplish more complex tasks than those within its training distribution , we propose a novel exploration method called variance seeking meta-exploration.
This method biases the agent to explore task dependent state-action pairs.
We show in our experiments that the agent trained using our method can generalise to second order tasks much better as compared to existing methods.
This shows the beneﬁt of adding an inductive bias towards certain state-action pairs during the exploration phase for OOD tasks.
Our method also sheds light on the importance of
TD error variance for meta-RL exploration.
Further research into the properties of states with high
TD error variance across environments, and the resulting behaviour of exploring in such states, is necessary to provide a sound theoretical underpinning to our method.
In addition to this, we are also working to test our meta-RL solution to control large real-world buildings to minimise their carbon footprint.