Trust, but verify: model-based exploration in sparse reward environments
Abstract
We propose trust-but-verify (TBV) mechanism, a new method which uses model uncertainty estimates to guide exploration.
The mechanism augments graph search planning algorithms with the capacity to deal with learned model’s imperfections.
We identify certain type of frequent model errors, which we dub false loops , and which are particularly dangerous for graph search algorithms in discrete environments.
These errors impose falsely pessimistic expectations and thus hinder exploration.
We conﬁrm this experimentally and show that TBV can effectively alleviate them.
TBV combined with MCTS or Best First Search forms an effective model-based reinforcement learning solution, which is able to robustly solve sparse reward problems.
TBV framework
This work attempts to pave the way in harnessing the power of graph search methods to reinforce model-based reinforcement learning.
Majority of the planners were designed be supplied with a perfect model , otherwise they often yield an unexpected, or outright wrong, result.
Although in some environments, like games or simulators, a perfect model is available, in most real-world scenarios it is a luxury we cannot afford.
We are thus led to two seemingly conﬂicting desiderata for the choice of a method: we would like to use models which are robust to model errors, and planners which leverage the structure of the underlying problem, and hence are more susceptible to model errors.
We advocate developing methods, which will facilitate adoption of search algorithms into setting with imperfect models.
In this work we make a step towards implementing this plan.
We identify and study certain type of model errors, which we dub false loops , and we propose Trust, but verify (TBV) method, which can successfully immunize planners with respect to said errors.
1.1 Model errors
Sutton & Barto (2018, Section 8.3) conveniently categorize model errors into two categories: optimistic and pessimistic.
A situation when a model predicts higher reward or better state transitions (optimistic error) is easier to handle since while the agent enters an erroneous region, it collects corrective data.
More dangerously, predictions more pessimistic then reality may prevent the agent ever to visit certain regions.
In this work, we concentrate on errors related to the graph structure of the problem.
By deﬁnition, the edges in this graph correspond to transitions induced by actions, and thus the errors are “incorrect edges”.
Perhaps the simplest prediction task (for a model) is to determine if an edge is one-step loop (i.e. if an action will change or not the current state).
In this work, we focus on the loop errors, as they turn out to be critical to performance (see Figure 1 for examples).
In Figure 2 we present numerical results that show rather frequent occurrence of false loop .
This is a pessimistic error, which, unless corrected, e.g. by TBV, prevents exploration.
To read broad context of existing work, see Section A. in Appendix.
1.2 Trust, but verify method
Trust, but verify (TBV) is a method which can be integrated with graph search planners.
The key idea of TBV is to prioritize visits in states for which the model is suspected to be pessimistic.
In each state the planner is asked for a recommended action.
If the model uncertainty of the current state is high, we may suspect that the planner’s results are misleading.
In such a case, a different action is chosen to verify the model’s accuracy.
This promotes exploration and thus reduces errors in the model.
To measure the uncertainty, we utilize ensembles, which is a popular approach in the literature.
To avoid potential problems with a priori uncontrollable scale of uncertainty, we quantify its magnitude as quantiles of the uncertainty distribution induced by the planner’s search graph.
The details follow the pseudo-code for TBV, listed in Algorithm 1.
For completeness, in Appendix, we also present a typical model-based training loop in Algorithm 2 .
As shown in Algorithm 1, TBV interacts with a planner, which uses an imperfect learned model.
During its execution, the planner expands a search graph, g p , in order to propose an action a p .Each state-action pair in g p is given a score reﬂecting uncertainty assigned to it by the model, and which is quantiﬁed by a function D ISAGREEMENT _ MEASURE .
Based on that, TBV decides whether to keep the planner’s decision a p (trust the model) or override it (and explore to verify model predictions).
More precisely, a state-action pair with the maximal score is considered if it reaches above the threshold given by quantile _ score .
We found using quantiles instrumental in avoiding tuning thresholds which are problem-speciﬁc and change during the training.
Furthermore, QR is relatively easy to tune (see Section 2) and is the only hyperparameter introduced by TBV.
In the algorithm, we use additional randomization to prevent multiple revisits of the same state-action pair, while waiting for the model (and scores) update.
This could be alternatively realized by more principled approaches, which we leave for future work.
Using TBV yields little computational overhead beyond the necessity of using ensembles (in efﬁcient implementations STATE _ SCORE and TRANSITION _ SCORE calls can be inlined into planner.
CHOOSE _ ACTION ).
D ISAGREEMENT _ MEASURE depends on the state space representation; in our experiments, we found standard deviation computed on states work well.
Other methods, like Bayesian inference, could be used as well.
1.3 Planners
Online Best First Search Best First Search (BestFS) is a family of search algorithms that builds a graph by expanding most promising nodes among already visited.
A classical BestFS expands nodes until it ﬁnds the goal state.
In order to use it in on-line planning regime, we choose to extend the
Algorithm
Require: model ensemble of models used in Algorithmplanner planner that uses model
QR quantile rank ∈ (0 , 1)
Use:
QUANTILE ( l , q ) computes the q -quantile of list l
RANDOM () random number from U (0 , 1) function CHOOSE _ ACTION ( state ) a p ← planner.
CHOOSE _ ACTION ( state ) g p ← planner.
GET _ GRAPH _ OF _ PLANNING scores ← [] for state ∈
g p do scores.
APPEND ( STATE _ SCORE ( state )) quantile _ score ← QUANTILE ( scores, QR ) one _ step _ scores ← [] for a ∈ actions do one _ step _ scores.
APPEND ( TRANSITION _ SCORE ( state , a )) if max( one _ step _ scores ) > quantile _ score and RANDOM () > 0 .
5 then return arg max action one _ step _ score elsereturn a p function TRANSITION _ SCORE ( state , action ) predictions ← [] for network ∈ model do next _ state, next _ reward ← network.
PREDICT _ NEXT _ STATE ( state , action ) predictions.
APPEND (( next _ state, next _ reward )) return D ISAGREEMENT _ MEASURE ( predictions ) function STATE _ SCORE ( state ) return max action ( TRANSITION _ SCORE ( state , action )) graph by ﬁxed amount of vertices in any planning phase (typically 10 ). The proposed action is the ﬁrst edge on the shortest path to the best node in the subgraph searched so far. An important design decision is choosing heuristics for assigning nodes numerical values. In our experiments, we concentrated of exploration and used Disagreement_Measure (standard deviation) between ensemble predictions.
By design BestFS ignores one step loops and thus is especially prone to false loops described in Section 1.
Monte Carlo tree search Monte Carlo Tree
Search (MCTS) is a well-established planning algorithm which was successfully applied to complex problems.
Due to its simplicity and effectiveness, it has numerous extensions (see
Browne et al.
(2012) for a survey) including famous AlphaZero Silver et al.
We use implementation of Milos et al.
(2019) taking adventage of graph structure and using ensembles of values to guide search.
For more details see
Appendix A.7 and Milos et al.
Experiments
The experiments below were chosen to highlight the aforementioned false-loops model errors.
We have selected ToyMontezumaRevenge environment and the Tower of Hanoi puzzle.
These environments have sparse rewards and distance to the goal spans at least a few hundred steps.
As a performance metrics for the methods, we use the minimal distance to the goal state, measured in steps, and treated as a function of total environment steps used in training.
We apply this measure to two planners: BestFS and MCTS (each in the vanilla version, with ε -greedy exploration, and TBV).
For comparison, we also use RND Burda et al.
(2018), a strong model-free exploration baseline.
In all experiments, we found that planners equipped with TBV mechanism signiﬁcantly outperform other baselines.
ToyMontezumaRevenge ToyMontezumaRevenge is navigation, maze-like, environment introduced by Roderick et al.
(2018) as a testing ground for long-horizon planning and exploration (see
Figure 7 in Appendix).
For both planners using epsilon-greedy improves over vanilla version but not as much as TBV (see
Figure 3).
For BestFS, 4 out of 10 experiments were unable to leave the ﬁrst room.
This stemmed from the fact that the agent could not make progress by getting stuck in the early stages of exploration due to the inability to pass through false-loops .
This is visualized in Figure 2, where it is also shown that TBV helps.
Tower of Hanoi The Tower of Hanoi is a classical puzzle consisting of 3 pegs and n disks (see section A.8.2 and Figure 8 in Appendix).
As can be seen in Figure 4, TBV method signiﬁcantly improves both MCTS and BestFS planners in this domain.
This is due to the nature of the Tower of Hanoi, which can create an illusion that the biggest disk stays in the same position.
This can result in the false belief of the model that the biggest disk does not move, hence causing the false-loop errors.
A Appendix A. Training setup
Code for all our experiments can be accessed at https://github.com/ComradeMisha/
TrustButVerify .
Our experiments adhere to the general model-based training loop logic, described in Algorithm 2.
We use a distributed system with 32 workers solving distinct episodes, where data gathered across a batch of workers is collected in two common experience buffers: the replay buffer for trainable model of size 50000 and the replay buffer for value function of size 30000.
Before the solving of actual episode, we collect 1000 random trajectories, that we use for initial training of the model of environment.
This initial training has critical importance for the performance of our agents.
Episodes are limited to: 600 steps for ToyMontezumaRevenge and 1000 steps for the Tower of Hanoi.
To update value network for MCTS we follow approach of (Milos et al., 2019) (similar to Hamrick et al.
for BFS experiments we simply calculate future discounted returns for each transition encountered by agent.
A. Model-based training loop
In Algorithm 2 we present general training loop used in our experiments.
Algorithm 2 Model-based training loop # Initialize parameters of ensemble of models # Initialize parameters planner specific networks # Initialize buffer repeat episode ← COLLECT _ EPISODE buffer.
ADD ( episode )
B ← buffer.
BATCH
Train ensembles of models
Update planner speciﬁc networks until convergence function COLLECT _ EPISODE s ← env.
RESET episode ← [] repeat a ← TBV_planner.
C HOOSE
A CTION (s) (cid:46) TBV-augmented graph planner s (cid:48) , r ← env.
STEP ( a ) (cid:46) using models, see Algorithmepisode .
APPEND (( s, a, r, s (cid:48) )) s ← s (cid:48) until episode is done return episode
A. Network architectures for model and value A.3.1 Single networks
In every experiments we use two neural network architectures: one to estimate the value function of a given state and the other to predict the outcome of taking a step in the environment.
In
ToyMR we represent state as a vector of length 17 and in the Tower of Hanoi as a vector of length 3 × number of discs (it is due to one-hot encoding of the peg number for each disc).
In both environments, for value estimation we use an MLP architecture with two hidden layers of 50 neurons and ReLu non-linearity and for model we use an MLP with four hidden layers ofneurons and ReLU activations.
The model output consists of: the difference between next and current observations (delta target), reward and theepisode termination ﬂag (0 or 1).
A.3.2 Ensemble
For both value estimation and model we use ensemble of networks with architectures described in
Section A.3.1.
The ﬁnal estimation of value is an average of predictions across ensemble.
The The standard deviation of this predictions multiplied by κ is used by MCTS as an auxiliary score added to value.
In Algorithm 3 we present the procedure of transforming the output of ensemble of model networks to a valid prediction of environment signal.
In our experiments, it turned out to be beneﬁcial to train the model to predict change between observations rather then the ready observations.
The transformation performs the following step: collects predictions of all networks in ensemble, then averages next predicted change of observations, predicted rewards and predicted end of episode ﬂags.
Averaged observation change is added to current state, then clipped to the range of possible values (some coordinates of the state vector are binary variables), rounded to integers and returned as the next predicted state.
Reward and end of episode ﬂag are predicted to be true if their average value across networks in ensemble is larger than 0 .
Algorithm 3 Transforming model ensemble predictions function PREDICT _ STEP ( state , action ) next _ observations ← [] rewards ← [] is _ done _ flags ← [] for network ∈ ensemble do : ( next _ obs _ delta, reward, is _ done ) ← network.
PREDICT ( state, action ) next _ obs ← state + next _ obs _ delta next _ observations.
APPEND ( next _ obs ) rewards.
APPEND ( reward ) is _ done _ flags.
APPEND ( is _ done ) predicted _ reward ← AVERAGE ( next _ reward ) if predicted _ reward > 0 .
5 then predicted _ reward ←else predicted _ reward ←predicted _ is _ done ← AVERAGE ( is _ done _ flags ) if predicted _ is _ done > 0 .
5 then predicted _ is _ done
← True else predicted _ is _ done ← False predicted _ next _ obs ← AVERAGE ( next _ observations ) predicted _ next _ obs ← CLIP ( predicted _ next _ obs ) predicted _ next _ obs ← ROUND ( predicted _ next _ obs ) return
To stabilize performance of value and model ensembles we used additional masks mechanism: we create a number of networks (given by ensemble size parameter), but each worker uses only a random subset of given size (given by number of masks parameter).
A. Quantile Rank value analysis
Figure 5 presents performance of agent for different values of Quantile Rank.
A. Hyper-parameters
In table 3 we present hyper-parameters used in our experiments.
[Table]
For MCTS we took hyperparameters from Milos et al.
As they used setup without learned model we needed to tune model architecture and model training parameters on ToyMontezumaRevenge domain.
Then, we separately tuned epsilon and Quantile Rank for each combination of domain/planner.
[Table]
A. RND and PPO
Interestingly, it was quite hard to tune RND on ToyMontezumaRevenge (the agent struggled to ﬁnd the second key and explored at most 14 out of 24 rooms).
The reason for this could be related to the fact that RND needs more data to work, and our data regime was low (originally RND was trained with billions of transitions).
Additionally, RND was designed and tested for visual observations, so our setup could not lend to its strengths.
Hyperparameters choice for RND can be found in Tableand Table 4.
For Hanoi Tower RND continued to explore the environment when left for longer training (up to 25M transitions, see 6).
Its progress is, however, very slow when compared to other methods.
This is in line with our observation from the previous section that RND is designed to work on larger scale experiments.
A. MCTS
A vanilla MCTS constructs a search tree in four stages: node selection, leaf expansion, rollout, and backpropagation (see Browne et al.
(2012, Section 3.1) for a comprehensive review).
In modern approaches, the rollout phase is often replaced by a neural network evaluation step.
MCTS is often equipped with auxiliary mechanisms exploiting the graph structure of the problem.
In our work, we use an MCTS implementation with four such mechanisms: hard loop-avoidance inside the search tree, soft loop-avoidance within to the whole episode, the transposition tables and amortized value estimates, see Milos et al.
for details.
Transposition tables (Childs et al., 2008; Gelly et al., 2012) enable sharing information between different tree nodes, which correspond to the same state (i.e. have been reached using distinct trajectories).
Similar in spirit are mechanisms, which attempt to avoid visiting a node multiple times.
These are akin to classical graph search methods in which a node it typically visited only once.
This may take place on the level of the whole episode, within one planning phase (in the search tree).
A “soft” way of implementing such mechanisms is to use “virtual loss”, which assigns a temporary negative value when a graph node is ﬁrst encountered.
This encourages the exploration and tree building mechanism to avoid this node.
Such solutions have been proposed in (Segal, 2010) to enhance parallelism and later used in (McAleer et al., 2019) to tackle Rubik’s cube search.
A hard version (equivalent to setting virtual loss to −∞ ) strictly prohibits entering the same node twice.
In (Milos et al., 2019) it is reported to be a signiﬁcant efﬁciency boost for modest planning budgets
[Table]
[Figure]
[Table]
FCN
[Table]
We use a code from https://github.com/chrisgrimm/deep_abstract_q_network on the classical Sokoban puzzle.
Recent works (Milos et al., 2019; Hamrick et al., 2020), propose to utilize “amortized value estimates” (which are calculated using the internal MCTS statistics) as targets for value function for training.
Our version uses a learned value function (similarly to AlphaZero) to evaluate new nodes.
Similarly to Lowrey et al.
and Milos et al.
we use also value function ensemble to ensure exploration during planning.
A. Environments A.8.1 ToyMontezumaRevenge
ToyMontezumaRevenge has a greatly simpliﬁed visual layer compared to the original Montezuma’s
Revenge Atari game, but it retains much of its exploration difﬁculty.
In our experiments, we consider the biggest map containing 24 rooms and sparse rewards: the agent gets a reward 1 only if it reaches the treasure room, otherwise the episode is terminated after 600 steps.
Observation is represented as a tuple containing current room location, agent position within the room, and status of all keys and doors on the board.
It turns out that traps (see caption below Figure 1) pose a challenge for BestFS algorithm.
In order to mitigate this issue, we to modiﬁed the novelty mechanism to explicitly avoid traps (performing actions resulting in episode termination without positive reward, according to the learned model).
For a fair comparison, we also added the same mechanism to epsilon-greedy: when sampling from action space, only actions not leading to traps were considered.
Similar improvements were unnecessary for MCTS.
A.8.2 Tower of Hanoi
The objective of the game is to move the entire stack of disks from the starting peg to the goal peg.
The rules are that only one disk can be moved at a time and it is disallowed to put a larger disk on top of a smaller one.
This makes it a challenging combinatorial problem, and an optimal solution requires 2 n − 1 moves.
This is an interesting domain for planning methods as model-free algorithms struggle to deal with larger instances (Troussard et al.
(2020) used n = 4 , and Edwards et al.
(2018) used n = 3 ).
In this paper, we use a considerably harder version, with n = 7 .
Furthermore, we consider a more difﬁcult version of the puzzle, where the agent does not have access to the disk sizes.
This does not allow to generalize the dynamics of the environment from partial data, and it makes it challenging to learn the rule that "the smaller disk can always be placed on larger".
Instead, the agent needs to learn the relation between sizes for each pair of disks separately, as the exploration progresses.
A. Related Work
Many model-based reinforcement learning algorithms follow the framework laid out in Dyna, see
Sutton (1991).
Kaiser et al.
(2020) uses a model to collect ﬁctitious playouts and obtains impressive results for low data regime on Atari.
In the continuous domains, a similar approach is adopted by
Kurutach et al.
Simultaneous training of a single policy on model ensemble is proposed to tackle model errors.
Interestingly, this approach is reported to perform better than back-propagation through time.
Clavera et al.
(2018) indicates that polices trained on model ensembles might be over-conservative and proposed to use policy meta-learning.
Another approach to dealing with model errors works by choosing the unroll horizon for a model.
Janner et al.
(2019), proposed short model-generated branched rollouts starting from data collected on “real environment”.
Perhaps surprisingly, they ﬁnd that one-step rollouts provide competitive results.
Similarly, short rollouts are used in Feinberg et al.
and Buckman et al.
(2018) to improve value estimates.
Janner et al.
; Buckman et al.
(2018) uses ensembles to reduce model bias.
The former generates diverse data similarly to Chua et al.
(2018), and the latter automatically adapt the planning horizon based on uncertainty estimates.
Xiao et al.
(2019) proposed another adaptive horizon mechanism based on measurement of model errors using principled Temporal Difference methods.
There is a huge body of work about planning; for classical results, the interested reader is referred to Cormen et al.
(2009), Russell & Norvig (2003) and the references therein.
Traditional heuristic algorithms such as A ∗ (Hart et al.
or GBFS (Doran & Michie (1966)) are widely used in practice.
In Agostinelli et al.
the authors utilise the value-function to improve upon the A ∗ algorithm and solve Rubik’s cube.
Similarly, Orseau et al.
(2018) bases on the classical BFS to build a heuristic search mechanism with theoretical guarantees.
The Monte Carlo Tree Search (MCTS) algorithm, which combines heuristic search with learning, led to breakthroughs in the ﬁeld, see
Browne et al.
(2012) for an extensive survey.
Famously, Silver et al.
(2018) develop MCTS-based technique to master the ancient game of go.
The POLO algorithm presented in Lowrey et al.
(2019) proposes to enrich MPC planning with ensemble-based risk measures to augment exploration.
This work is extended by NEEDLE (Milos et al.
to tree-based MCTS planners and by Lu et al. (2019) to successfully tackle life-long learning scenarios of environmental change.
Hamrick et al. (2020) explores how to use better statistics collected while searching to calculate signal for value function training.
The above works plan using a perfect model.
Recently, a lot of research has focused on planning with learned models.
Nagabandi et al.
(2018) successfully blends the strengths of model-based and model-free approaches.
PlaNet (Hafner et al. (2019a)) and Dreamer Hafner et al.
(2020) train latent models, which are used for planning.
Conceptually, a similar route was explored in MuZero (Schrittwieser et al.
and Universal Planning Networks (Srinivas et al.
Farquhar et al.
(2017) and Oh et al.
(2017) investigate the possibility of creating neural network architectures inspired by the planning algorithms.
When a learned model is unrolled during planning, errors typically accumulate dramatically.
Racanière et al.
(2017) and Guez et al.
(2018) are two approaches to learn the planning mechanism and make it robust to model errors.
Sekar et al.
(2020) uses ensembles to measure model uncertainty and deal with model inaccuracies.
In a somewhat similar spirit, Eysenbach et al.
(2019) treats the replay buffer as a non-parametric model forming a graph and uses ensembles of learned distances as a risk-aware mechanism to avoid certain types of model errors.
There is an emerging body of work on the intersection of model learning and exploration.
Pathak et al. (2017) uses a self-supervised inverse dynamics model to provide intrinsic reward driving exploration.
Sekar et al.
(2020) proposes a method of learning task agnostic global world model, which can be utilised to quickly solve new incoming tasks.
Somewhat similarly, Henaff (2019) draws inspiration from E 3 algorithm (Kearns & Singh, 2002), actively searching for data useful in model training.
Shyam et al.
(2019) has a similar objective implemented using novelty mechanism derived from
Bayesian perspective on exploration.
Our work also falls into this area.
We construct a model-based approach capable of exploring sparse reward environments.
Compared to (Pathak et al., 2017) we use uncertainty measures and planning.
Sekar et al.
(2020) uses ensembles similarly to us but instead of an on-line planner utilises Dreamer (Hafner et al., 2019b).
Henaff (2019) is perhaps the closest to our work, however it puts emphasis on continuous environments and consequently leaves aside using graph structure and related model errors.
Similarly, Shyam et al.
(2019) largely ignores graph structure and online planning, though brieﬂy mentions a loop problem stemming from this fact.