T HE B ENEFITS OF M ODEL -B ASED G ENERALIZATION IN R EINFORCEMENT L EARNING
A BSTRACT
Model-Based Reinforcement Learning (RL) is widely believed to have the potential to improve sample efficiency by allowing an agent to synthesize large amounts of imagined experience.
Experience Replay (ER) can be considered a simple kind of model, which has proved extremely effective at improving the stability and efficiency of deep RL.
In principle, a learned parametric model could improve on
ER by generalizing from real experience to augment the dataset with additional plausible experience.
However, owing to the many design choices involved in empirically successful algorithms, it can be very hard to establish where the benefits are actually coming from.
Here, we provide theoretical and empirical insight into when, and how, we can expect data generated by a learned model to be useful.
First, we provide a general theorem motivating how learning a model as an intermediate step can narrow down the set of possible value functions more than learning a value function directly from data using the Bellman equation.
Second, we provide an illustrative example showing empirically how a similar effect occurs in a more concrete setting with neural network function approximation.
Finally, we provide extensive experiments showing the benefit of model-based learning for online RL in environments with combinatorial complexity, but factored structure that allows a learned model to generalize.
In these experiments, we take care to control for other factors in order to isolate, insofar as possible, the benefit of using experience generated by a learned model relative to ER alone.
Model-based reinforcement learning (RL) refers to the class of RL algorithms which learn a model of the world as an intermediate step to policy optimization.
One important way such models can be used, which will be the focus of this work, is to generate imagined experience for training an agent’s policy (Werbos, 1987; Munro, 1987; Jordan, 1988; Sutton, 1990; Schmidhuber, 1990).
Experience
Replay (ER) can be seen as a simple, nonparameteric, model (Lin, 1992; van Hasselt et al., 2019) where experienced interactions are directly stored, and later replayed, for use in a learning update.
ER already captures many of the benefits associated with a learned model as compared to modelfree incremental online algorithms (i.e. model-free algorithms which perform a learning update using each transition only at the time it is experienced).
In particular, ER allows value to be rapidly propagated from states to their predecessors along previously observed transitions, without the need to actually revisit a particular transition for each step of value propagation.
Propagating value only at the time a transition is visited can make model-free incremental online algorithms extremely wasteful of data, particularly in environments where the reward signal is sparse.
As Lin (1992) and van Hasselt et al.
(2019) have discussed, it is often not obvious why we’d expect experience generated by a learned model to improve upon ER, as a replay buffer is essentially a perfect model of the world insofar as the agent has observed it.
This is especially true in the tabular case, where a model does not generalize from the observed transitions.
It is also true for policy evaluation in the case where the value function and model are linear (Sutton et al., 2012).
In this case, learning the least-squares linear model from the data, and then finding the TD(0) solution (Sutton, 1988) in the resulting linear MDP is identical to finding the TD(0) solution for the empirical MDP induced by the observed data.
Hence, if we expect to obtain a sample efficiency benefit by using data generated by a learned model compared to ER, we should look beyond these cases.
Let’s now consider a situation where a model-free agent using ER will likely fail to generalize from experience in a way that an intelligent agent should be able to.
Imagine an agent has witnessed a tree lying across a stream which allows it to cross in order to reach food on the other side.
Imagine the same agent has previously pushed against a decaying tree and knocked it over, but in that case gained nothing by doing so.
We may then expect that upon seeing a decaying tree standing near a stream with food on the other side, the agent would be able to synthesize these two experiences to decide to intentionally push the tree over to provide a bridge over the stream.
We argue that ER alone is unlikely to achieve the kind of generalization necessary to perform such extrapolation.
To understand why, let’s express the state as a combination of two abstract binary factors (tree fallen, food across stream).
The agent has observed state (1,1), and also observed separately that from state (0,0) it can achieve state (1,0) through a certain action (call that action push-tree ).
Now, the agent observes state (0,1).
An agent using ER can easily learn that (1,1) is valuable since it is followed by a food reward, but since it has never executed push-tree from (0,1) to reach that valuable state, there is no reason to expect it will assign that action an elevated value.
On the other hand, a learned model that generalizes appropriately could guess, perhaps based on some inductive bias toward factored dynamics 1 , that taking action push-tree in (0,1) would lead to (1,1) in the next step.
Having done so, we could then use the model in various ways to incorporate this information into the value function and/or policy.
One common way of using the model, that we will focus on in this work, is to generate (potentially multi-step) rollouts of simulated experience and train a value function or policy on the resulting trajectories as if they were real.
Having discussed how learned model generalization can provide a benefit over ER, it’s worth noting that parametric value functions also generalize.
Why should model generalization be inherently better than value function generalization?
This question was already raised in the work of Lin (1992), which first introduced ER for RL.
The next section gives a partial answer as a theorem which shows how learning a model as an intermediate step can narrow the space of possible value functions more than learning a value function directly from the data using the Bellman equation.
After motivating the benefit of model-based generalization theoretically, we will present an intuitive case where learning a parametric model is empirically beneficial with NN function approximation.
Subsequently, we will present extensive experiments, in environments with factored structure, which highlight the sample-efficiency benefits of model-based learning for online RL.
We will also analyze an interesting instance we came across during these experiments where an agent using a learned model outperforms one using the perfect model due to smoothed reward and transition dynamics.
We use the MDP formalism throughout this paper.
We refer the unfamiliar reader to Sutton & Barto (2020) for an accessible discussion of MDPs and their usage in RL.
U NDERSTANDING THE B ENEFIT OF M ODEL -B ASED G ENERALIZATION
We now present a simple theorem that provides at least part of the answer to the question of how model generalization can be considered more useful than value function generalization.
We state the theorem informally here, and formally with proof in Appendix A.
Theorem 1.
Consider a class of deterministic, episodic, MDPs M with fixed reward function, and transition function belonging to some known hypothesis class.
Let H Q be the associated class of optimal action-value functions for MDPs in M .
Now consider a dataset D of transitions.
Let
H B ( D ) be the subclass of action-value functions in H Q which obey the Bellman optimality equation for the transitions in D and let H M ( D ) be the subclass of optimal action-value function of MDPs in M which are consistent with D .
Then the following are true:
H M ( D ) ⊆
H B ( D ) .
For some choices of M and D , H M ( D ) ⊂
H B ( D ) .
For a tabular transition function class, that is one that includes every possible mapping from state-action pairs to next states, H M ( D ) =
H B ( D ) .
Intuitively Theorem 1 states that, if we want to narrow down the possible optimal action-value functions from data, we can in general prune more if we narrow down the possible models first than if we only demand that the value functions obey the Bellman optimality equation with respect to the observed data.
However, part 3 states that this is not the case for a tabular model class, but rather only for models which have some additional structure.
In such cases, even if a model-based and modelfree approach begin with the same hypothesis class, the model-free approach can fail to leverage this structure and ultimately lose value-relevant information from the data in the process.
While the assumptions of Theorem 1 differ from the reality of training neural networks (NNs) with ER, this theorem provides intuition for why learning a parametric model can improve sample efficiency.
In the next section, we will highlight a simple and intuitive case where similar intuition leads a learned model to be empirically beneficial with NN function approximation.
A S IMPLE C ASE WHERE M ODEL -B ASED G ENERALIZATION IS U SEFUL
In this section, we present an illustrative example where learning a parametric model from data, and then learning an action-value function within that model, has a clear advantage over learning an action-value function directly from the data.
We also show that using multi-step model rollouts has an advantage over single-step rollouts.
To keep this example as simple and intuitive as possible, we consider an offline RL setting with hand-selected datasets with varying coverage of the dynamics.
The environment in this section consists of a 3x3 grid world with a goal in the top left corner and with arbitrary walls (the location of which is fixed within an episode).
Reward is -1 at each step until the goal is reached, at which point the episode terminates.
Actions consist of standing still, or moving in a cardinal direction.
Observations are flat binary vectors consisting of one-hot-encodings of the agent’s position and the goal position (though goal position is fixed here), and two binary vectors indicating, respectively, the presence and absence of walls in each cell.
[Figure]
the
For explanatory purposes, we break the data in the training datasets provided to the agents into a basic set and an evaluation set .
The basic set consists of all possible transitions with wall layouts having a single wall within one of the 8 non-goal cells.
Every training dataset contains this entire basic set, in addition to some limited data from the evaluation set.
The evaluation set consists of transitions with wall layouts of 2 walls in one of the configurations illustrated in Figure 1(left).
We consider training datasets with 4 different levels of evaluation set data coverage:
• All-Evaluation : All possible transitions within the evaluation wall layouts.
• Path-to-Goal : Only evaluation layout transitions which follow the path to the goal.
• Single-Cell : Only the transitions starting from the cell furthest from the goal (with respect to the only open path).
• No-Evaluation : No transitions from the evaluation set.
We compare model-based and model-free algorithms in this setting.
For the model-free approach, we use a Deep Q-Network (DQN; Mnih et al.
(2015)) with ER.
For the model-based approach, we use a simple feedforward NN model, and DQN trained on model generated transitions.
The model takes an observation as input and outputs a predicted reward, Bernoulli termination probability and vector of Bernoulli probabilities that each feature is active in the next state.
Note that this model cannot capture correlation between features in the next observation given the current observation.
However, as the proposed environment has deterministic transitions, it is sufficient to capture the true dynamics.
We train the model on transitions from the dataset and subsequently train the actionvalue function on model rollouts initialized from states in the dataset.
We trained one model-based agent with single-step rollouts and another with 10-step rollouts.
We trained each agent for 1 million training steps and controlled the total number of (real or imagined) transitions used in each DQN update across all agents.
See Appendix E for further detail on the experiment setup.
Each agent is evaluated by checking the frequency out of 30 independent runs with which the greedy action under the learned value function is optimal in each cell of an evaluation layout.
We say the agent has failed if there exists any cell in which the majority of runs select the wrong greedy action.
We next discuss how we expect each agent to perform with each level of evaluation set coverage.
All-Evaluation : We expect that all agents will succeed in the case where data is given for all transitions in the evaluation layouts.
Here, a model-free agent has access to all the data needed to backup value from the goal and determine the value of each state-action pair, and a model-based agent has no opportunity to generate useful novel transitions which do not already appear in the dataset.
Path-to-Goal : We expect model-free DQN to fail in the Path-to-Goal case.
Consider the lower evaluation wall layout in Figure 1(left).
For the bottom middle cell, there is no data in the Pathto-Goal dataset for actions besides moving right.
For every basic-set layout, moving right is worse than any action which remain in the same cell.
We predict the model-free agent will incorrectly generalize from the basic set to conclude that the right action is also worse in this new configuration.
We expect both model-based agents to succeed with Path-to-Goal data.
All missing transitions are one step away from the available data, but require a different action selection.
We predict that the outcome of these missing actions can be learned by generalization from the basic set.
As in the simple situation described in the introduction, this environment has significant factored structure.
To determine the effect of an action, it suffices to look at the agent’s current location and whether there is a wall in the cell it is attempting to enter.
We predict the agent will be able to learn this basic structure from the training data, even in the absence of specific data about the case where there are two walls.
Note that this hypothesis implies a nontrivial prediction about how this simple model will generalize.
Factored structure is not hard-coded in the model, thus it is also plausible that the model predictions will be arbitrarily bad for the unobserved evaluation transitions.
Single-cell : We expect the single-step model to fail when only transitions from a single evaluation cell are included in the dataset.
In this case, the model rollouts have no chance to reconstruct anything not already available explicitly in the dataset given all single step transitions from the far cell are included and this is likely to be insufficient for reasons already explained.
However, the 10-step model has the potential to succeed.
If the model generalizes as expected, it can start in the far cell available in the dataset and rollout a trajectory which discovers the full path to the goal.
No-Evaluation : Finally, all agents should fail in the case where there is no evaluation data available at all.
The models will have no opportunity to provide the learned value function with example transitions from the evaluation layout given model rollouts are initialized with states from the dataset.
Note that the situation may be different if we had used a generative start-state model to produce plausible states for the start of rollouts which need not explicitly appear in the dataset.
2 Using the model to plan for immediate action selection at evaluation time, as in model predictive control, would also help here, as the agent could directly plan a response for the previously unseen state.
In Figure 1(right), we observe that all the above predictions are confirmed.
We reiterate that this is a nontrivial empirical result.
It relies on the simple model, a feedforward NN, with no explicit bias toward factored solutions, generalizing in a particular way to state-action pairs that do not explicitly appear in the dataset.
At least in this case, the model indeed seems to generalize in a way that provides a significant advantage over ER alone.
This experiment also straightforwardly illustrates how even a model with 1-step rollouts can be helpful, by sampling counterfactual actions or, in the case of stochastic environments, counterfactual chance outcomes.
However, multi-step rollouts can succeed in situations where there is insufficient data for one-step rollouts to be helpful.
F AVORABLE E NVIRONMENTS FOR O NLINE M ODEL -B ASED L EARNING
We next describe three environments which exemplify properties that should make online learning with a parametric model particularly useful.
We aim for the following environment characteristics:
The state-space should have a simple factored structure that we expect should be easy to learn for a model with reasonable generalization properties.
The return should depend sharply on the policy, such that a randomly behaving agent won’t have much of a learning signal for policy improvement.
This should make model generated experience more useful, as Bellman backups alone will tend to be mostly uninformative. 3.
To allow the model-based agent to learn about the reward function, the agent should occasionally witness rewarding states even while behaving highly sub-optimally.
Toward this, we include occasional random transitions to rewarding states in our subsequent experiments.
The same is true for terminal states in environments where termination is desirable.
This does not contradict the previous point, as an agent may be able to occasionally obtain some reward while behaving sub-optimally, but have difficulty obtaining more.
The last characteristic may seem contrived, but we argue that it is quite natural.
For example, one can imagine an agent gathering edible plants for a long time before working out how to grow their own.
It is often much harder to discover reward, without ever observing it, than to work out how to reconstruct rewarding circumstances using general knowledge of the transition dynamics.
On a practical note, there is significant work on exploration with sparse (or no) reward (Schmidhuber, 1991a;b; Thrun, 1992; Schmidhuber, 2010; Amin et al., 2021) which is orthogonal to our focus.
Hence, we mitigate the issue by making it easier to learn the reward function.
We ablate these spontaneous transitions to rewarding states in Appendix I to test the impact of this choice.
In addition to the above, the environments we investigate allow scaling of problem complexity to test the limitations of different approaches.
The environments are also Markov, use binary features, and are largely deterministic, so simple models can work well, though we also investigate more sophisticated latent-space models.
Next, we describe the environments (see Appendix B for details).
[Figure]
a b c d e
R
[Figure]
ProcMaze (Figure 2, left): Procedurally generated grid world mazes.
The maze itself, along with the start state and goal state are randomized in each episode.
Complexity is scaled by increasing the grid size.
Negative reward is given for each step until the goal is reached.
ButtonGrid (Figure 2, middle): A 5 by 5 grid with randomly placed buttons.
An agent can move around and, if it hits a button, will toggle it on or off.
If all buttons are on, a reward is given and button locations are randomized.
Random behavior will tend to randomly perturb the buttons, a precise policy is required to set them all to on.
Complexity scales with the number of buttons.
PanFlute (Figure 2, right): A minimal example of an environment with combinatorial complexity of optimal behavior, but simple factored transition structure.
PanFlute consists of n pipes of cells where each pipe evolves independently.
Each action directly activates the cell at the bottom of one pipe, after which the activation will propagate up the pipe, one step at a time, and dissipate after reaching the end cell of the pipe.
A reward is received if the cells at the end of all pipes are simultaneously active, which can only be achieved by choosing each of the n actions in a certain order, a probability of 1 /n n under random behavior.
We scale the complexity by changing the number of pipes n .
We expect ER alone to be of limited utility in each of these environments as each of them require precise control to obtain significantly more reward than random behavior, especially as the problem complexity is scaled up in each case.
Since we always include random transitions to rewarding states, an agent can easily learn that these states are good, but until it reaches the rewarding state by it’s own actions, it wont be able to learn much from the states which proceed rewarding states.
On the other hand, each environment has significant factored structure which a model-based agent can learn, and subsequently use to imagine many novel, plausible, states in its rollouts.
In ProcMaze, an agent moving into a specific empty space will have the same effect regardless of the rest of the maze layout, and attempting to move into a wall will always block it.
In ButtonGrid, the connectivity of the grid is independent of the button layout, and stepping on a specific button always has the same effect regardless of the layout of the rest of the buttons.
In PanFlute, each action always has the same effect, and each pipe evolves according to dynamics which are unaffected by the other pipes.
In appendix C, we present an experiment in an environment without factored structure and find that
ER outperforms a learned model in this case, which helps to contextualize our other results.
B ENEFICIAL M ODEL -B ASED G ENERALIZATION FOR O NLINE RL
We now empirically evaluate the performance of model-based and model-free learning algorithms on the environments described in Section 3.
We experiment with variants of each environment with a range of different complexities to test how different approaches scale to more complex environments.
All tested approaches use DQN for behavior learning, but vary in how training examples for DQN are generated.
Our model-free approach draws transitions randomly from an ER buffer for training.
We test several types of learned model.
The first is the simple feedforward NN model introduced in Section 2.
The second is a latent-space model (Schmidhuber, 1997; Watter et al., 2015; Ha &
Schmidhuber, 2018) inspired by Dreamer (Hafner et al., 2019; 2020), but with two major differences to simplify the approach and reduce confounding factors in our experiments.
In particular, as noted, we use DQN instead of actor-critic and, since our environments are Markov, we forgo the recurrent network of Dreamer and model single-step stochastic transitions with no memory.
We experiment with Gaussian and Categorical latent variables.
See Appendix D for further details of these models.
Finally, as a strong baseline, we include a perfect model, which uses the ground truth environment dynamics instead of a learned model, but is otherwise the same as the simple-model agent.
As in Section 2, we control for the total number of updates, and the total number of (real or imagined) transitions used in each update.
In particular, model-free DQN is updated on a batch oftransitions from the replay buffer while all model-based approaches use 32 model rollouts of length 10, beginning in a state from the replay buffer.
In Appendix I, we perform an ablation in which 1-step rollouts are used instead and find that longer rollouts are generally helpful.
In each update, the model is trained using the same batch of transitions which initialize the rollouts.
All agents use a softmax behavior policy, and are evaluated under the greedy policy.
All Q-networks are trained on discounted 1-step TD-error, with real or imagined transitions, using a discount factor of 0 . 9 .
We experiment with 2 different data regimes to get a more complete picture of how each approach scales with available data.
In the high data regime, we use one update per real environment step and train for a total of 1 million steps.
In the low data regime, we use 10 updates per step and train for 100 thousand steps.
Note that the total number of updates is the same in each case.
Experiment Design: For each combination of agent, environment, and data regime we performed an extensive grid search over the Q-network step-size and softmax exploration temperature.
We judged these hyperparameters to be the most likely to impact the relative performance of different methods.
This grid search was performed for an intermediately complex version of each environment (size
ProcMaze, 4 button ButtonGrid, and 7 pipe PanFlute) and the same hyperparameters used for the other complexity levels.
We evaluated each hyperparameter setting based on mean final performance of the greedy policy over 30 random seeds.
We were able to run 30 seeds efficiently in parallel on a single GPU using automatic batching in JAX (Bradbury et al., 2018).
Other hyperparameters, including model step-size, were fixed to reasonable defaults (see Appendix F), not tuned for any specific approach.
In Appendix G, we report hyperparameter sensitivity results from this grid search.
[Figure]
Results: We present results for all environments in Figure 3 (see Appendix H for learning curves).
As the results for PanFlute differ substantially from the results for ButtonGrid and ProcMaze, we will discuss them separately.
For ButtonGrid and ProcMaze in the high data regime, the simple model and categorical latent model both significantly outperform ER for sufficiently complex environment instances.
The Gaussian latent model generally performs quite poorly, which corroborates the results of Hafner et al.
(2020) that categorical latents tend to work better in the discrete control setting.
Oddly, the simple model also performs much worse for 1 and 3 buttons in ButtonGrid in the high data regime.
This may be because less buttons means less examples where a button occupies each particular cell, which makes it harder for the simple model to learn the underlying dynamics.
In the low data regime, the simple model outperforms ER to a greater extent, while the performance of the categorical latent model degrades significantly.
This can perhaps be understood by noting that the hypothesis class of the simple model is simpler (the latent-space model can model correlated features, while the simple model cannot) and thus able to generalize well from less data when the simple class is sufficient.
Performance of the simple model for 1 and 3 buttons improves when moving to the low data regime.
Likely, this indicates underfitting in the high data regime which is helped by training more on each example.
Overall, the results for the simple model and categorical latent model in the high data regime, and simple model in the low data regime, show a clear indication of the sample efficiency benefit that can be obtained by using a learned model in these environments.
Results for PanFlute are qualitatively different.
Most surprisingly, the Gaussian-latent model and the simple model outperform the perfect model in some of the harder problem instances.
This is intuitively strange, and seems to indicate that model errors somehow improve performance.
Why do Some Learned Models Outperform the Perfect Model on PanFlute?
We hypothesize that the smoother reward and dynamics learned by the model serves as a powerful exploration heuristic.
The true reward function is nonzero, and thus the agent receives a learning signal, only in the rare event that every pipe-end is active.
The model might instead learn a smoothed reward function, proportional to the number of pipe-ends activated.
With this smoothed reward the agent could incrementally learn to activate more pipe-ends, gradually improving toward the correct sequence.
To test this hypothesis, we looked at the models learned at 10,000 time-steps, in 9-pipe PanFlute (high data).
We generated a large amount of random trajectory data with a policy that selects actions in alphabetical order with 80% probability and uniformly randomly otherwise to get a good mix of different numbers of active pipe-ends.
We bin this data by the number of active pipe-ends and then look at the average model-predicted reward for each bin.
Note that the ground truth reward is zero for all expect the 9 active pipe-end bin.
To test for favorable smoothing in the transition dynamics, we used the same data.
This time, we bin the data by the number of pipe-ends active at the next step and look at probability under the model that all pipe-ends are active at the next time-step.
[Figure]
The results, for both predicted reward and predicted probability that all pipe-ends will be active at the next time-step (“predicted probability of rewarding next state” in the figure) are displayed in
Figure 4(left).
We observe that the models indeed tend to learn smoother rewards than the ground truth in a way that might provide a useful exploration heuristic.
The Gaussian latent model additionally learns smoothed transition dynamics, predicting all pipe-ends will activate with appreciable probability when in reality only most will, this may provide additional benefit in this problem.
As an additional test of the benefit of model errors, we used simple models frozen at various points in training to train a value function from scratch for 1,000,000 time-steps.
We plot when near-optimal performance of the greedy policy is first reached.
The results, shown in Figure 4(right) clearly show that a model trained for an intermediate amount of time is most useful for reaching good performance quickly.
We also plot the mean predicted reward for states with 6 active pipe-ends for each model as an indication of smoothing, as expected, this decreases with more model training.
The smoothing effect highlighted in these results is interesting for two reasons.
First, it may lead model-based algorithms to perform better than expected, acting as a confounding variable when interpreting results.
Second, it may be genuinely useful.
One could even design algorithms which learn policies within relaxed versions of a model, as a way to drive exploration.
However, as there are surely other situations where model smoothing is harmful, this would need to be done with care.
R ELATED W ORK
There is a large body of empirical research showing the potential of learned models to improve sample efficiency (Deisenroth & Rasmussen, 2011; Buckman et al., 2018; Kaiser et al., 2019; Janner et al., 2019; Curi et al., 2020; Hafner et al., 2020).
A relatively small body of work directly compares ER with learned models.
Van Seijen & Sutton (2015) show an exact functional equivalence between a variant of replay with linear TD(0), and linear Dyna (Sutton, 1990).
Pan et al.
(2018) provide an empirical study comparing ER with learned models under a variety of search control strategies, i.e. different methods for choosing which state-action pairs to update.
Holland et al.
(2018) empirically compare ER with a parameteric model in the arcade learning environment (ALE;
Bellemare et al.
and highlight the benefits of the model in particular when multi-step rollouts are used.
Relative to these works, we focus on understanding how a learned model provides a benefit, and highlighting properties of environments where this benefit is most prominent.
Van Hasselt et al.
(2019) make a strong case that ER provides many of the benefits of a learned parametric model, and argue that if a model is used only to generate experience starting from observed states it is unlikely to provide additional benefit.
We investigate the comparison between ER and learned models further, and argue that there is good reason to believe a learned model can improve sample efficiency over ER in environments with structure, such as factored dynamics.
This is true even if the model is used only to augment training data with rollouts starting from observed states.
Dong et al.
(2020) share our focus on highlighting situations where model-based RL provides a significant benefit.
While we focus on the generalization benefits, they motivate the expressivity benefit of model-based RL in continuous state spaces by showing there exist MDPs where the optimal policy is exponentially more complex to represent than the dynamics.
The benefit of model smoothing, observed in PanFlute, may shed light on some observations in the literature.
Hafner et al.
(2020) suggest model smoothing as an explanation for why DreamerV2 can achieve good performance on Montezuma’s Revenge without any sophisticated exploration mechanism, as is usually required.
Holland et al.
(2018) observe that their learned model sometimes outperforms the ground-truth model in Seaquest, though they do not suggest a specific explanation.
Our work has implications for implicit model-based algorithms (as defined in the survey paper of Moerland et al.
such as MuZero (Schrittwieser et al., 2020).
In such algorithms, the model is not trained to predict future observations, but only task relevant aspects of the future such as policy, value and reward.
In light of Theorem 1, and the example in Section 2, it is unclear whether such techniques fully exploit the benefits of model-based learning due to the limited training signal used to constrain the model.
Indeed, there is empirical evidence suggesting that MuZero’s sample efficiency can be improved by using additional training signals (Ye et al., 2021;
Anand et al., 2022).
Our work is loosely inspired by the literature on exploiting factored structure in MDPs (Sallans &
Hinton, 2004; Strehl et al., 2007;
Diuk et al., 2009; Osband & Van Roy, 2014; Sun et al., 2019; Xu &
Tewari, 2020).
Theorem 1, in particular, is loosely related to Theorem 2 of Sun et al.
(2019), which shows that there exists a family of MDPs where a model-based approach can be exponentially more efficient than any model-free approach.
We focus on factored structure to motivate the benefit of model generalization, but do not use algorithms explicitly designed to exploit this factored structure.
AIXI (Hutter, 2004) provides a particularly general approach to model-based RL.
An AIXI agent maximizes expected return over all computable world models with prior probability given by
Solomonoff’s universal prior (Solomonoff, 1978), which assigns higher probability to models with lower Kolmogorov complexity.
While AIXI is itself not computable, computable approximations exist (Veness et al., 2011).
Here, we take a step back from considering all computable models, and highlight the benefit of model learning in general, relative to model-free learning with ER.
We focus on understanding how model-based learning can facilitate generalization, and thus improve sample efficiency, in a way that model-free learning does not.
Our empirical results suggest this applies even for the practical subset of computable models representable by a simple NN architecture.
In this work, we focus on the advantages of using models to generate experience for policy improvement.
Learned models can also be useful in other ways, for example, local planning for immediate action selection (Richalet et al., 1978; Chua et al., 2018; Byravan et al., 2021), improving credit assignment (Schmidhuber, 1990; Heess et al., 2015; Buesing et al., 2018)
, exploration (Schmidhuber, 1990; 1997; Pathak et al., 2017), representation learning (Lin & Mitchell, 1992; Gregor et al., 2019;
Hessel et al., 2021), and in answering learned queries (Schmidhuber, 2015).
C ONCLUSION
In Theorem 1 we highlighted how, for a model class with some known structure, we can narrow down the possible value functions more with a model-based approach than a model-free approach.
We provided empirical evidence that the intuition behind this theorem holds in practice when we use NN function approximation in domains with factored structure.
In such cases, we have verified through extensive experiments that a model-based method can maintain strong performance as the complexity of the environment increases beyond the point where an analogous model-free approach fails.
As an aside we demonstrated how, by smoothing the reward and/or transition dynamics, experience generated by a learned model can provide a useful signal for exploration that can sometimes lead to better performance than even a perfect model.
Overall, we believe that our work can help to ground future work in model-based RL in a better understanding of how learned models can improve sample efficiency.
An interesting direction for future work lies in better understanding the inductive biases that allow simple NN models to generalize in a way that allows them to efficiently learn factored structure, and how more sophisticated architectures could improve on this.
A CKNOWLEDGMENTS
We thank Richard Sutton, Michael Bowling, Yi Wan, Arsalan Sharifnassab and Tian Tian for useful conversations.
Funding for this work was provided by NSERC, Alberta Innovates, DeepMind,
CIFAR, Amii, ERC Advanced Grant (742870) and the Swiss National Science Foundation grant (200021 192356).
Computational resources for this work were provided by the Digital Research
Alliance of Canada as well as the Swiss National Supercomputing Centre (CSCS project s1127)