The Benefits of Model-Based Generalization in Reinforcement Learning
Authors: Kenny John Young
Aditya Ramesh
Louis Kirsch
Jürgen Schmidhuber
Venue: Submitted to ICLR 2023 (2023)
Keywords: model-based reinforcement learning
generalization

Abstract: 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.

Under review as a conference paper at ICLR 2023
THE BENEFITS OF MODEL-BASED GENERALIZATION IN REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Model- Based Reinforcement Learning (RL) is widely believed to have the poten- tial to improve sample efficiency by allowing agent to synthesize large amounts of imagined experience. Experience Replay (ER) can be considered a simple kind of model, which has proved extremely effective improving the stability and ef- ficiency 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 em- pirically 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 be useful. First, we provide general theorem motivating how learning a model as an inter- mediate step can narrow down the set of possible value functions more than learn- ing value function directly from data using the Bellman equation. Second, we provide an illustrative example showing empirically how a similar effect occurs 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- reinforcement learning (RL) refers to the class of RL algorithms which learn 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 's policy (Werbos 1987; Munro 1987 Jordan 1988 Sutton 1990, /Schmidhuber [990). Experience Replay (ER) can be seen 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 learned model compared to model- free 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 be rapidly propagated from states to their predecessors along previously observed transitions, without the need to actually revisit 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, we expect to obtain a sample efficiency benefit by using data generated by learned model compared to ER, we should look beyond these cases.
Let's now consider a situation where a model-free free agent using ER will likely fail to generalize from experience in way that an intelligent agent should be able to. Imagine agent has witnessed

---

Under review as a conference paper at ICLR 2023
tree lying across 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 decaying tree standing near 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 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 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 followed by 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" [ 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 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 value function policy on the resulting trajectories if they were real.
Having discussed how learned model generalization can provide a benefit over ER, it's worth noting that parametric value funetions 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 partial answer 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 cample-efficienee
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 (020) for an accessible discussion of MDPs and their usage in RI.
UNDERSTANDING THE BENEFIT OF MODEL-BASED GENERALIZATION
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 class of deterministic, episodic, MDPs M with fixed reward function, and transition function belonging to some known hypothesis class. Let HQ be the associated class of optimal action-value functions for MDPs in M. Now consider dataset D of transitions. Let H (D) be the subclass of action-value functions in HQ which obey the Bellman optimality equation for the transitions in D and let Has (D) be the subclass of optimal action-value function of MDPs in M which are consistent with D. Then the following are true:
1. H (D) C HD()).
2. For some choices of M and D, HD(() C. (D).).
3. For labular transition function class, that is one that includes every possible mapping from state-action pairs to next states, H (D)=
Intuitively Theorem 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 tabular model class, but rather only
Where the state consists of set of state- variables such that the distribution of values of each variable the next time-steps depends only subset of the variables in the current state.
2

---

Under review as a conference paper at ICLR 2023
for models which have some additional structure. In such cases, even if model-based and model- free approach begin with the same hypothesis class, the model-free approach can fail to leverage this structure and ultimately lose value-relevane information from the data in the process. While the assumptions of Theorem | 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 simple and intuitive case where similar intuition leads learned model be empirically beneficial with NN function approximation.
A SIMPLE CASE WHERE MODEL- -ASED GENERALIZATION IS USEFUL
In this section, we present an illustrative example where learning a parametric model from data, and then learning action-value function within that model, has clear advantage over learning an action-value function directly from the data. We also show that using multi-step model rollouts has advantage over single-step rollouts. To keep this example simple and intuitive possible, consider an offline RI setting with hand-selected datasets with varying coverage of the dynamics.
The environment in this section consists of 3x3 grid world with 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 reached, 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-encodee 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.
Ak-ivaluation
Part-to-Gool
SingleGoll
No-Evotaation
DON
Baai: configuratios
Exeluation
configaations
Modul
04
Step
Model
Figure 1: Left: Maze layouts included in the basic and evaluation sets. All datasets include every transition in each of the basic configurations, but different sets of transitions the evaluation figurations. Right: Frequency (over 30 random seeds) of trained agents' greedy policy selecting the correct action in cells of the displayed Length of green arrows indicate the frequency of greedy policies picking the corresponding action.
with different evaluation set coverage during training.
For explanatory purposes, break the data in the training datasets provided the agents into 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 walls one of the configurations illustrated in Figure [[eff)). 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-hased approach, we simple feedforward NN model, and DQN trained 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 capture the true dynamics. We train the model on transitions from the dataset and subsequently train the action-
3

---

Under review as a conference paper at ICLR 2023
value function model rollouts initialized from states in the dataset. We trained one model-basee based agent with single-step rollouts and another with 10-step rollouts. We trained each agent for million training steps and controlled the total number of (real or imagined) transitions used in each DQN update across all agents. See Appendix F 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.
We expect that all agents will succeed in the case where data is given for all transi- tions 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 model-based agent has no opportunity to generate useful novel transitions which do not already appear in the dataset.
All-Evaluation:
Path-to-Goal: We expect model-free DQN fail the Path-to-Govd to-Goal case. Consider the lower evaluation wall layout in Figure IIeft). For the bottom middle cell, there is no data in the Path- to-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 two walls. Note that this hypothesis implies a nontrivial prediction about how this simple model will generalize. Factored structure not hard-coded in the model, thus it is also plausible that the model predictions will be arbitrarily bad for the unobserved evaluation transitions.
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, can start in the far cell available in the dataset and rollout a trajectory which discovers the full path to the goal.
Single-cell:
No- No-Evaluation: Finally, all agents should fail in the case where there is no evaluation data available at all. The models will have opportunity to provide the learned value function with example tran- sitions 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 generative start-state model to produce plau- sible states for the start of rollouts which need not explicitly appear in the dataset 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 rightt) we observe that all the above predictions 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 model with 1-step rollouts can be helpful, by sampling counterfactual actions 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.
2 However, this raises questions about how the start-state model should generalize, if trained to maximize data likelihood, it could overfit and only generate states from the training set, rendering it unhelpful. We also verified that the model predictions in this case are near perfect with respect to predicting the agent position after each state-action pair across all transitions in the evaluation configurations.
4

---

Under review as a conference paper at ICLR 2023
3 FAVORABLE ENVIRONMENTS FOR ONLINE MODEL-BASHR LEARNING
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:
1. The state-space should have simple factored structure that we expect should be easy to learn for model with reasonable generalization properties.
2. The return should depend sharply on the policy, such that randomly behaving agent won't have much of 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 occa- sionally witness rewarding states
while behaving highly sub-optimally. Toward this, we include occasional random transitions to rewarding states in our subsequent experi- ments. 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 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 reconstruct rewarding circumstances using general knowledge of the transition dynamics. On practical note, there is significant work on exploration with sparse (or no) reward (Schmidhuber, 1991a by Thrun 1992; Schmidhuber 2010; Amin et al.] (02T) which 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 Appendix 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 different approaches. The environments are also Markov, use binary features, and are largely deterministie, simple models can work well, though also investigate more sophisticated latent-space models. Next, we describe the environments (see Appendix B for details).
ProcMaze
ButtonGrid
PanFlute
Figure 2: Left: Examples of states in the ProcMaze environment of size 4 with the agent shown in orange. Middle: An example state of the ButtonGrid environment with 3 buttons. The agent is shown in orange and the buttons in black (off) and white (on). Right: An instance of the PanFlute environment with pipes. The agent directly activates cells (a,b,c,d,e) through its actions after which the activation propagates up the associated pipe, one step at time, and dissipates at the end.
ProcMaze (Figure 2 left): Procedurally generated grid world mazes. The maze itself, along with the start state and goal state 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 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 randomized. Random behavior will tend to randomly perturb the buttons, precise policy is required to set them all to on. Complexity scales with the number of buttons.
PanFlute (Figure right): A minimal example of 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 the bottom of one pipe, after which the activation will propagate up the pipe, one step at a time, and dissipate after reaching
5

---

Under review as a conference paper at ICLR 2023
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" 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 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 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 wall will always block it. In ButtonGrri,, the connectivity of the grid 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 dynamies which are unaffected by the other pipes. In appendix C we present 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.
4 BENEFICIAL MODEL-BASED GENERALIZATION FOR ONLINE 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 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] (020), but with two major differences noted, to simplify the approach and reduce confounding factors in experiments. In particular, 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 memory. We experiment with Gaussian and Categorical latent variables. See Appendix [ for further details of these models. Finally, as 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 imag- ined) transitions used in each update. In particular, model-free DQN is updated on a batch of 320 transitions 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 | 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 1-step TD-error, with real 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 I 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.
For each combination of agent, environment, and data regime we performed an extensive grid search over the Q-network step-size and softmax cxploration temperature. We judged these hyperparameters to be the most likely to impact the relative performance of different methods. This grid search was performed for intermediately complex version of each environment (size 4 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
Experiment Design:
we report hyperparameter sensitivity results from this grid search.
6

---

Under review as a conference paper at ICLR 2023
ProcMaze
ButtonGrid
PanFlute
High Data
Experienco Rapley
-10
Parle Model
Reward
Retum/
Final
-20
30
0.15
0.10
Colagorical Leaat
Sinple Vevia
Gauslan Lalon:
0.20
0.15
Rate
440
0.10
0.05
-50
0.05
0.00
0.00
Low Data
-10
0.20
0.25
Retum/
Final
-20
30
0.15
0.20
Reward
0.15
Rate
40
0.10
0.10
-50
0.05
0.05
-60
1.00
3.0
3.5
4.0
4.5
5.0
6
6
10
Grid Size
Number of Buttons
Number of Pipes
Figure 3: Final performance of greedy policy for all environments in two different data regimes. High data means million environment interactions with I update per step, low data means 100 thousand interactions with 10 updates per step. Error bars show 95% confidence intervals.
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 Hafnnr et al.] (2020) that categorical latents tend to work better in the discrete control setting. Oddly, the simple model also performs much worse for I and 3 buttons in ButtonGrid in the high data regime. This may be because less buttons means less examples where button occupies each particular cell, which makes harder for the simple model to learn the underlying dynamics.
Results:
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 I and 3 buttons improves when mov- ing to the low data regime. Likely, this indicates underfitting in the high data regime which is helped by training more 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 heuris- tic. The true reward function nonzero, and thus the agent receives learning signal, only in the rare event that every pipe-end is active. The model might instead learn smoothed reward func- tion, 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 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 reward for each bin. Note that the ground truth reward is zero
look at the average model-predicee 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 probability under the model that all pipe-ends are active the next time-step.
7

---

Under review as a conference paper at ICLR 2023
Simple Model
Gaussian Latent
Prokcde Roward Addno Aip Sross
Predicted Revard
9.0
0.30
Predicted Probability
Prodcced ROASONNN
0.6
Rewarding Next State
0.4-
0.13
022
0.0
3.0
2.5
4.0
4.5
5.0
3.5
Number pipe lourenthroott
Number onds ondeeenteete
Log (buse 120 Modul training SlMgs
Figure 4: Left: Prediction, under simple model and Gaussian latent models, of reward and proba- bility that all pipe-ends active at the next time step function of true number of current and next pipe-ends active respectively on 9-pipe PanFlute. Right: In orange, the number of steps to reach near-optimal performance (95% of maximum possible reward rate) 9-pipe PanFlute when using a frozen simple model trained for variable number of steps. Numbered arrows indicate the number of seeds out of 30 which failed to reach near-optimal performance within million training steps (marked as black line). In green, the predicted reward for 6 active pipe-ends, displayed as a surrogate for the amount of model smoothing. Error bars show 95% confidence intervals.
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 TigureBleft) 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 addi- tionally 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, used simple models frozen 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 rright) rrggh)) 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 reasons. First, it may lead model-based algorithms to perform better than expected, acting as 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 way to drive exploration. However, as there are surely other situations where model smoothing is harmful, this would need to done with care.
RELATED WORK
There is a large body of empirical research showing the potential of learned models to improve sam- ple efficiency (Deisenroth & Rasmussen 2011 Buckman et al.; 2018 Kaiser al. 2019 Janner jet al. 2019) Curi et al., [020] Hafnnr et al. 2020)). J. A relatively small body of work directly com- pares ER with learned models. |aan 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 variety 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. (2013) and highlight the benefits of the model in particular when multi-step roll- outs 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 a learned para- metric model, and argue that if a model is used only to generate experience starting from observed states 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 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.
8

---

Under review as a conference paper at ICLR 2023
Dong al. (2020) share focus on highlighting situations where model- RL provides a sig- nificant 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 pol- icy 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 's Revenge without any sophisticated exploration mech- anism, as is usually required. [Holland et al. (2018) observe that their learned model sometimes outperforms the ground-trum -trtth 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. (2020)) 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 [ and the example Section[] 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 f Indeed, there is empirical evidence suggesting that MuZero's 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
Strchl et al. 2007 Diuk et al. 2009; Osband & Van Roy 2014 Sun et al. [019 Xu &
Tewari 2020p. Theorem in particular, is looslly related to Theorem 2 of Sun et al. (2019), which
shows that there exists a family of MDPs where 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 RI. An AIXI agent maximizes expected return over all computable world models with prior probability given by Solomonoff's s universal prior (Solomonofi 1978). which assigns higher probability to models with lower Kolmogorov complexity. While AIXI is itself not computable, computable approximations exist (Veness et al. T 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-basee 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 simple NN architecture.
In this work, we focus on the advantages of using models to generate experience for policy improve- ment. 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 as-
signment (Schmidhuber 1990 Heess et al. 2013; Buesing et al. 2018). ) exploration (Schmidhuber
| 1997 Pathak et al. (2017). representation learning (Lin & Mitchell 1992]] Gregor et al. 2019
Hessel et al. 2021). and answering learned queries (Schmidhuber 2015).
CONCLUSION
In Theorem T we highlighted how, for a model class with some known structure, we can narrow down the possible value functions more with 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 analogous model-free approach fails. As an aside we demonstrated how, by smoothing the reward and/or transition dynamics, expe- rience generated by learned model can provide useful signal for exploration that 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 better understanding of how learned models can im- prove 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 leam factored structure, and how more sophisticated architectures could improve on this.
'Though Gehring et al. 2021 suggest that the implicit model-based parameterization itself may yield fa-
vorable gradient dynamics.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Susan Amin, Maziar Gomrokchi, Harsh Satja, Herke van Hoof, and Doina Precup. A survey of exploration methods in reinforcement learning. arXiv preprint arXiv:2 2021.
Ankesh Anand, Jacob C. Walker, Yazhe Li, Eszter Vértes, Julian Schrittwiese,, Sherjil Ozair, Theophane Weber, and Jessica B Hamrick. Procedural generalization by planning with self- supervised world models. In International Conference on Learning Representations, 2022. URL
net/ forum? FineeiiddFFFBBees.s
Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning envi- ronment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations Python+NumPy programs, 2018. URL http:
//github. ctt.uub.cccomoooggee/e/eoo
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems, 31, 2018.
Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Cean-Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactuully-guiee. policy search. arXiv preprint arXiv: 181 1.06272, 2018.
Arunkumar Byravan, Leonard Hasenclever, Piotr Trochim, Mehdi Mirza, Alessandro Davide Ia- longo, Yuval Tassa, Jost Tobias Springenherg, Abbas Ahdolmaleki, Nicolas Heess, Josh Merel, et al. Evaluating model-based planning and planner amortization for continuous control. arXiv preprint arXiv:2110.0336 2021.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in a handful of trials using probabilistic dynamics models. Advances neural information processing systems, 31, 2018.
Sebastian Curi, Felix Berkenkamp, and Andreas Krause. Efficient model-based reinforcement learn- ing through optimistic policy search and planning. Advances in Neural Information Processing Systems, 33:14156-14111 2020.
Mare Deisenroth and Carl E. Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-I1), pp- 465-472. Citeseer, 2011.
Carlos Diuk, Lihong Li, and Bethany R. Leffler. The adaptive k-meteorologists problem and its application to structure learning and feature selection in reinforcement learning. In Proceedings of the 26th Annual International Conference on Machine Learning, pp. 249-256, 2009.
Kefan Dong, Yuping Luo, Tianhe Yu, Chelsea Finn, and Tengyu Ma. On the expressivity of neural networks for deep reinforcement learning. In International Conference on Machine Learning, pp. 2627-2637. PMLR, 2020.
Clement Gehring, Kenji Kawaguchi, Jiaoyang Huang, and Leslie Kaelbling. Understanding end- to-end model-based reinforcement learning methods implicit parameterization. Advances in Neural Information Processing Systems, 34:703-714 2021.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. Advances in Neural Information Processing Systems, 32, 2019.
David Ha and Jürgen Schmidhuber. World models. Advances in Neural Information Processing Systems, 31, 2018.
10

---

Under review as a conference paper at ICLR 2023
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.022 193, 2020.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap. Tom Erez, and Yuval Tassa. Learn- ing continuous control policies by stochastic value gradients. Advances in neural information processing systems, 28, 2015.
Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre. Theophane Weber, David Silver, and Hado van Hasselt. Muesli: Combining improvements in policy opti- mization. In International Conference on Machine Learning, pp. 4214 4226. PMLR, 2021.
Zacharias Holland, Erin J. Talvitie, and Michael Bowling. The effect of planning shape on dyna- style planning in high-dimensional state spaces. arXiv preprint arXiv: 1806.0 1825, 2018.
Marcus Hutter. Universal Artificial Intelligence: Sequential Decisions based on Algorithmic Prob- ability. Springer, Berlin, 2004. (On J. Schmidhuber's SNF grant 20-61847).
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in Neural Information Processing Systems, 32, 2019.
Michael I. Jordan, Supervised learning and systems with excess degrees of freedom. Technical Report COINS TR 88-27, Massachusetts Institute of Technology. 1988.
Lukasz. Kaiser, Mohammad Bahaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model-based reinforcement learning for atari. arXiv preprint wwXiv:/903.00377 2019.
Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching- Machine learning, 8(3):293-33) 1992.
Long-Ji Lin and Tom M. Mitchell. Memory approaches to reinforcement learning non-Markovian domains. Citeseer, 1992.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7540):529-33 -33-2 2015.
Thomas M. Moerland, Joost Broekens, and Catholijn M. Jonker. Model-based reinforcement learn- ing: A survey. arXiv preprint arXiv:2006.1617 16712, 2020.
P. W. Munro. A dual back-propagation scheme for scalar reinforcement learning. Proceedings of the Ninth Annual Conference of the Cognitive Science Society, Seattle, WA, pp. 165-176, 1987.
Ian Osband and Benjamin Van Roy. Near-optimal reinforcement learning in factored mdps. Ad- vances in Neural Information Processing Systems, 27, 2014.
Yangehen Pan, Muhammad Zaheer, Adam White, Andrew Patterson, and Martha White. Organizing experience: deeper look at replay mechanisms for sample- planning in continuous state domains. urXiv preprint arXiv: 1806 04624, 2018.
Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 16-17, 2017.
Jacques Richalet, André Rault, J.L. Testud, and J Papon. Model predictive heuristic control: Appli- cations to industrial processes, Automatica, 14(5):413-42, 1978.
Brian Sallans and Geoffrey E. Hinton. Reinforcement learning with factored states and actions. The Journal of Machine Learning Research, 5:1063-1088. 2004.

---

Under review as a conference paper at ICLR 2023
J. Schmidhuber. Artificial curiosity based on discovering novel algorithmic predictability through coevolution. P. Angeline, Z. Michalewicz. M. Schoenauer, X. Yao, and Z. Zalzala (eds.), Congress on Evolutionary Computation. pp. 1612-1618. IEEE Press, 1999.
Juergen Schmidhuber. On learning to think: Algorithmic information theory for novel com- binations of reinforcement learning controllers and recurrent neural world models. Preprint arXiv:1511.09249 2015.
Jürgen Schmidhuber. Making the world differentiable: On using fully recurrent self-supervised neural networks for dynamic reinforcement learning and planning in non-stationary environments. Technical Report FKI-126-90, htttp://peqp:
issia ch/-juergent (revised) bw ocr pdd, Institut für Informatik, Technische Universität München, 1990.
Jürgen Schmidhuber. A possibility for implementing curiosity and boredom in model-building ral controllers, In J. A. Meyer and S. W. Wilson (eds.). Proc. of the International Conference on Simulation of Adaptive Behavior: From Animals Animats, pp. 222-227. MIT Press/Bradford Books, 1991a.
Jürgen Schmidhuber. Curious model-building control systems. In Proceedings of the International Joint Conference on Neural Networks, Singapore, volume 2, pp. 1458-1463 IEEE press, 1991b.
Jürgen Schmidhuber. What's interesting? Technical Report IDSIA-35-97, IDSIA, 1997. extended abstract in Proc. Snowbird'98, Utah, 1998;
ftp://ftp. idsia. haanaaachhhhbieeeggennntteetttst see also [Schmidhuber (1999] 2002).
Jürgen Schmidhuber. Exploring the predictable. In A. Ghosh and S. Tsuitsui (eds.), Advances in Evolutionary Computing, pp. 579-5179 Springer, 2002.
Jürgen Schmidhuber. Formal theory of creativity, fun, and intrinsic motivation (1990-2010). IEEE Transactions on Autonomous Mental Development, 2(3):230-22) 247, 2010. ISSN 1943-0604. doi:
Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-669) 2020.
Ray Solomonoff. Complexity- induction systems: comparisons and convergence theorems. IEEE transactions on Information Theory, 24(4):422-432, 1978.
Alexander L Strehl, Carlos Diuk, and Michael L. Littman. Efficient structure learning in factored- state mdps. AAAI, volume 7, pp. 645-650, 2007.
Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference on learning theory, pp.
PMLR, 2019.
Richard S. Sutton. Learning to predict by the methods temporal differences. Machine learning.
3(1):9 44, 1988.
Richard S. Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. In Machine learning proceedings 1990, pp. 216-224 2244. Elsevier, 1990.
Richard S. Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, second edition, 2020.
Richard S. Sutton, Csaba Szepesvári, Alborz Geramifard, and Michael P. Bowling. Dyna- style planning with linear function approximation and prioritized sweeping. arXiv preprint arX6:: 2012.
Sebastian Thrun. Efficient exploration in reinforcement learning. Technical Report CMU-CS-92. 102, Mellenie-Mellon University, January 1992.
12

---

Under review as a conference paper at ICLR 2023
Hado P. van Hasselt, Matteo Hessel, and John Aslanides. When to parametric models in rein- forcement learning? Advances in Neural Information Processing Systems, 32, 2019.
Harm van Seijen and Rich Sutton. A deeper look at planning as learning from replay. In International conference on machine learning, pp. 2314-2322. PMLR, 2015.
Joel Veness, Kee Siong Ng, Mareus Hutter, William Uther, and David Silver. A monte-carlo aixi approximation. Journal of Artificial Intelligence Research, 40:95-142. 142,-12,222 2011.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
Paul J. Werbos. Building and understanding adaptive systems: A statistical/mumeriaal approach factory automation and brain research. IEEE Transactions on Systems, Man, and Cybernetics, 17, 1987.
Ziping Xu and Ambuj Tewari. Reinforcement learning in factored mdps: Oracle-efficient algorithms and tighter regret bounds for the non-episodic setting. Advances in Neural Information Processing Systems,
2020.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34, 2021.
13

---

Under review as a conference paper at ICLR 2023
A THEOREM MOTIVATING THE BENEFIT OF MODEL GENERALIZATION
Here, we present and prove a simple theorem motivating the benefit of learning a parametric model over learning a value function directly from ER. Intuitively speaking, the theorem states that, when narrowing down the set of possible value functions based on observed data, we can rule out more we first rule out models directly, and demand the value function be consistent with the reduced model class, than if we only demand the value function obeys the Bellman optimality equation with respect to observed transitions.
We state the theorem within the formalism of finite MDPs. An MDP consists of state-space S, action-space A, reward function r x A + I and transition function p. In general, p maps state-actions pairs to probability distributions over possible next states. However, in Theorem [ consider deterministic MDPs, meaning each state-action pair maps to a distribution with probability particular next state s' and zero for all other next states. In this case, it will be convenient to write S x A - S as mapping from state-action pairs to the only possible next state. We assume S and A are finite sets for simplicity.
An agent interacts with an MDP in a series of time-steps, beginning in some state So E S. At each time-step, the agent observes the current state S, E S and selects an action A, E A in response. The environment then transitions to the next state S = p(S,, AA)) and the agent receives a reward Re41 r(,,, At). Interaction continues until some designated terminal state Sx = is reached at which point the episode is over. The goal of an agent is to find policy, mapping from states to actions π - which maximizes the expected return E, [G.|S, s). where G2 = >x=++1 Rg, from any given state until the terminal state Sy = 1 is reached when actions are selected according to We assume that all policies eventually reach 1. with probability one, such that the return is well defined. The action-value function of a policy is defined as 9x (s, a) = E [G,|S; = 8, A, = a). the expected return if action a selected in state s and policy is followed from that point forward.
5
The optimal action-value function q" (8, a) = max= (= a) is defined as the maximum action value over all policies æ. q (s. a) is known to be the unique solution the Bellman optimality equation ("(s,a) a) = r(s, a) + E maax. (S(41, a') d'){S, = s, A,, = a] where the value of all actions in the terminal state 1 are defined to be zero. In the case of a deterministic transition function, this reduces to simply 9" (s, a) = r(s. a) + max.' 9" (oss,)).a)
We state Theorem [ for deterministic MDPs, but analogous results likely hold for general MDPs, albeit significantly complicated by the fact that in the general case, models and value functions can only be ruled out with high probability based on observed data, as opposed to with certainty.
Theorem I Consider class of episodic MDPs, M, with fixed reward function r S x A + R and deterministic transition function belonging Assume, for all E H, all policies lead to eventual termination.
hypothesis class H S {p : S S A + S).
Since each MDP in the class deterministic, following a deterministic policy n, beginning in will give rise to specific state-action trajectory (T( (s0, da, .1, 01,1,
where 8p 8, 8,00 = a, is the terminal state, and p E H. Define also G(Tp(TT, s,))
r(St. 0t). the return associated with the trajectory.
Next, define the class of optimal action value functions associated with H: 8, + B E H : Vs, a g(s, a) max G(T,(T,,,a)),
or equivalently: HQ = bq :q::s00 8,
Rp E Vs, a q(s, r(s, a) + maxq(p((s
a'
Consider a dataset D = {Sn, an, 5/)]m E 10, N}} of transitions such that p(sn, on) = s', for some p H. We now define two different notions of hypothesis classes over action-value functions which are consistent with D:
HD(D) {g HQlVn q(3n, an) r(sn. an) + max a')}
a"
Hy(D) = {q E HolEp E H: (Yn p(sn, an) = s4,) A (Vs, a q(s.)(ss))) = r(s, a) + max q(p(s a), a'))
a'
'Again, policies can more generally map states to distributions over actions, but focus here on the deterministic case, and note that there is always a deterministic optimal policy.
14

---

Under review as a conference paper at ICLR 2023
Where B stands for Bellman consistency and M stands for model consistency. In words, these the hypothesis classes consisting of value functions which obey the Bellman optimality equation with respect the observed transitions, and the hypothesis class consisting of true optimal value functions for transition dynamics which are consistent with the observed transitions respectively. Then the following true:
1. (D) _ HD)))
2. For some choices M and D, (D) C (D).).
3. For tabular transition function class, that is one that includes every possible mapping from state-action pairs next states, HM (D)==H()D))
Proof. We begin by proving part 1. Towards this, assume q E HM(D)). Then by definition we have the following:
3p H (Vn p( an) s/)) (Vs.aq(s,a) = a) maxq(p(s,
Ep E H (Vn p(5n: an) = s/) (vn an) an) = r(sn; an) + max max (p((s, on),
Vn (( an) = (Sn, an) + max g (
q (D()),
which proves part
To prove part 2, it suffices construct specific M and D for which (D) Hn(D). Towards this, consider a class of deterministic MDPs M with state-space S N UL where the dynamics of each component of (nonterminal) S; E S are factored such that S. does not influence S++1ij] for i * j. Furthermore Sy[0] {0, 1, 2). [ {0, 1), S.[2] f0, 1}. The reward function is defined as follows:
( if x11] x and s(0]
rss, a)
0 otherwise.
The transition dynamics for S, are fixed within the class such that
S [0] = S - if 5, >
S 1L. otherwise.
Thus [.] acts counter for the number of steps until termination, and the agent gets positive reward only both bits ST-1 [1] and ( are set to at the step prior to termination. The action space is A {0, 1,2). This model class is illustrated in Figure
Now let the dataset D consist of the following transitions:
(s: (2:00,))) (1,1,0)
(2: (1,0, (1.00)
(s:(22,,0),::1.::): (1,0,,)
(s: (2,1,,1),:::11::(111
(s : (2,00)) (1,0,0)
(1) Given the factored dynamics and that the dynamics of S,(0) are fixed within the model class, this data suffices uniquely determine p H to be the transition function identified by the following equations (when S,41 is nonterminal), in addition to the known dynamics for Sy [)]
(s: (2,1,1), 2,5' 1,).
S [1] [n))() if A, [s.[1] otherwise
[not(S,[2)) if A, = [ otherwise.
S
15

---

Under review as a conference paper at ICLR 2023
R
1
1
2
A
Figure 5: Illustraiion of the model class used as an example where selecting hypothesis based on model consistency more selective than Bellman consistency. The state consists of three compo- nents from left to right. The left most component is fixed within the class and simply acts as countdown to termination. The other two components have unknown internal transition dynamics but are known to not be influenced by other state components. The reward function is known to be zero except for when the two rightmost components are set to one at termination.
That is A, toggles component 1 of the state and A, toggles component 2, while A, leaves both components unchanged. In particular, the data is such that it contains an example of the effect of each action each possible bit value of S,[1] and S,(2]. meaning the above transition function is the only possibility. Thus HM consists of the singleton of the true optimal value function q (s, a) for this MDP, which is 1 if and only if the agent has sufficient steps remaining after taking action a before termination to toggle bits [11] and S S both to 1.
On the other hand, we can show that Hp (D) is a larger set. particular, consider the following action-value function:
((8 a) 1 if s|1] s[2] = and s{0] = 0
(2)
otherwise.
Note that ((),a) E HD))), as it satisfies the Bellman equation,
((sn, an) = r(sn, a,) + max a'),
a'
for all transitions D, although it is not the true optimal value function and thus is not in H (D). The value function ğ(s, is also in HQ. since it is the value function of the MDP with altemna- tive factored transition dynamics where (for example) [1]][1] S = regardless of their previous value the action choice. This suffices prove part 2 of the Theorem.
Finally, we prove part 3. Since we have already shown in part that E HM((D) - q HD)) in the general case, it suffices show that, with the additional restriction of tabular class H of transition matrices, we have E (D) = q E (D). Recall that by tabular we mean one which includes every possible mapping from state-action pairs next states. Thus, in this case we are free to choose p(s, a) to be an independently selected s' for every (s, a) pair, and know that the resulting p E H. Assume q E HB (D), then by definition we know q E HQ. meaning
3p E Н Vs, a q(s, a) r(s, a) + max (p((s a),
(3)
d'
and, by the definition of HD(D), we know
Vn (Sn, an) r(Sns an) max q(s a') a').
(4)
16

---

Under review as a conference paper at ICLR 2023
Now, given tabular H set, for all p(5n, an) 8/ which we know from Equation 4 gives us:
((6n, an) = r(sn; an) + max g(p( an), a').
Now for 8, where An (s,a) a) (8n: an), given Equation know that for some choice of p(s, a) it holds that:
q(s, a) = a) + max maxx(p(s,a))a') a'
Thus, given the choice of tabular H. we can choose p to simultancously satisfy p(an, an) s for all n and the Bellman optimality equation for all (s, a), and indeed:
q E (D(D) = Eee E H (An p(8n, 5')) ((s,)) = r(s, a) + max (()x
o
E (D))),
which completes the proof of part 3.
Note that, for simplicity, the dataset D in the example for part 2 does not consist of full episodic trajectories. However, it is straightforward which the same outcome still holds, for example by appending each transition in Equation 1 with any sequence of transitions leading termination with reward.
come up with dataset of episodic trajectories for
Avoiding inclusion of the rewarding transition in the dataset is also not necessary to obtain an anal- ogous result. For example, consider adding another action a 3 to the above example with known dynamics that always switch [1+] and S.+1 2) to 1, but always gives immediate reward of 1. Now consider adding the following episodic sequence (here including rewards for clarity) to the dataset of Equation]
(s: a:: (1,0,0)
(s: ::: 1,)1,
1).
Though the data still uniquely specifies the model, (D(D) will contain the incorrect action-value function
if 3 and x0] + 1
((s,a)
1 if s[1]= s/2] = and s(0] = 0 and a * 3
0 otherwise,
which is minor modification of Equation include the case where 3.
It is also straightforward to come up with similar examples where the data uniquely specifies model. but H (D) includes value functions for which the associated greedy policy is unique and suboptimal in the true model.
Practical learning algorithms like stochastic gradient descent applied with NN function approxima- tion don't work by directly ruling out hypotheses based on the data. Rather than considering models within an explicitly defined class, NN would have some inductive bias towards particular models and away from others depending on the architecture and hyperparameters. Nonetheless, the idea behind Theorem D helps to give insight into why we should expect learning parametric model to provide sample efficiency benefit.
Intuitively, we can think of the process performing many updates to a sufficiently high capacity neural network, with data from dataset, as incrementally constraining the possible functions rep- resented by the network. The specific function converged to will depend on the initialization and random batches selected for updates. Theorem I suggests that learning model as intermediate step can impose more constraints on the possible value functions. This in turn should increase the chance of converging to value function that generalizes well from limited data.
17

---

Under review as a conference paper at ICLR 2023
B FURTHER ENVIRONMENT DETAILS
The environments we investigate are all Markov and use flat binary observation vectors. To allow the model-based agent to learn about the reward or termination function, we include occasional random transitions to rewarding or terminal states. The probability of these transitions is chosen to result in much lower return than is achievable with optimal performance in each environment. Except for these random transitions, all environments are deterministic (with the exception of rare random transitions to rewarding or terminal states) so simple models can be expected to work well, though we also investigate more sophisticated latent-space models. Here, we will describe each of the environments in some detail.
The first environment, ProcMaze, is illustrated in FlgfreZleft). ProcMaze is episodic environ- ment. ProcMaze consists of procedurally generated gridworld mazes, where an agent has to navigate from a start state to goal state. The maze itself, along with the start state and goal state are random- ized at the start of each episode. A reward of 11 is given for each step until the goal is reached, at which point the episode terminates. Also the agent rarely randomly teleported to the goal (prob- ability 0.1/T where T' is the time required to complete the worst case problem instance for the grid size under the optimal policy) such that it can obtain knowledge of the reward function even with a poor behavior policy. Difficulty could be scaled by increasing the grid size. The observations consist of a flat binary vector including: one hot vectors for the goal location and agent location, a vector which is one if and only if a cell contains a wall, and vector which is one if and only if cell does not contain a wall. The action space includes attempting to move in each cardinal direction, and no-op. An attempted move will fail if it would lead the agent into wall or the edge of the grid. In each episode a new maze is generated using randomized depth first search, which produces reasonable mazes and guarantees the goal is reachable.
The second environment, ButtonGrid, is illustrated in Figure Mmiddle). ButtonGrid is a continuing environment, with no termination. ButtonGrid consists of grid world with a set of randomly placed buttons. An agent (orange in the figure) can move around on the grid and if it hits a button it will toggle it either on (black in the figure) or off (white in the figure). Reward is given whenever all the buttons are set to on, at which point and the button locations are randomized, but the number of buttons held fixed, and all buttons set to off. Occasionally all the buttons will spontaneously switch to on (probability 0.1/(grid size)?), meaning the agent can receive examples of the reward function even while behaving suboptimally. Importantly, once to solve this environment, as they are toggled on and off by repeated contact, an agent must also carefully avoid hitting them again after the first time they are pressed. Difficulty can be scaled by increasing the number of buttons on the grid, as well as the grid size, but we focus on the former. The observations consist of a flat binary vector including: one hot vectors for the agent location, a vector which is one if and only if a cell contains button which is turned on, and a vector which is one if and only if a cell contains a button which is turned off. The action space includes attempting to move in each cardinal direction, and no-op. An attempted move will fail only if it would lead the agent into the edge of the grid.
does not suffice touch each button
The third environment, PanFlute is illustrated in Figure right). PanFlute is continuing environ- ment with no termination. PanFlute is intended as a minimal instantiation of an environment with combinatorial complexity in terms of optimal behavior, but a simple factored transition structure. The observations consists of binary value for each square in the figure. An agent has actions available to it, (a,b,c,d,e) in the figure. Each action will activate the associated cell at the bottom of a specific pipe. The pipe associated with the last action (alphabetically) has a length of one cell, every other pipe is one cell longer than its (alphabetical) successor. a cell is activated at a given time-step, it will deactivate and activate the cell above it in the same pipe at the next time-step. A reward of 1 is received if the cells at the end of each pipe are simultancously active, otherwise the reward is always zero. An active pipe-end will always deactivate at the next step regardless of whether reward is obtained. Occasionally, the cells at the end of all pipes will activate spontancously (probability 1 1/n²), thus allowing the agent to observe rewarding situation without having to create it through its own actions. Otherwise, due to the arrangement of pipe lengths, the only way for the agent obtain reward is to choose each of the n actions in sequence. We can scale the difficulty of the environment by changing the number of actions n. The probability of a random sequence of n actions reaching the rewarding state (aside from spontaneous activation) is 1/n".. Observations consist of flat binary vector which includes the active/inactive state of each cell in each pipe.
18

---

Under review as a conference paper at ICLR 2023
C EXPERIMENTS IN AN ENVIRONMENT WITHOUT STRUCTURED TRANSITIONS
This work primarily focuses on highlighting environments in which model-based learning is ex- pected to be beneficial. Nevertheless, it is worthwhile to contrast this with what happens in environ- ments which do not have such favorable characteristics. To that end, we ran an additional experiment on an environment without factored structure, which we will now present.
The environment for this experiment, which we refer to as OpenGrid, was simply an open grid with goal in the bottom right corner and reward of -1 for every step until the goal is reached at which point termination occurs. The agent starts in a random location each episode. As in our other experiments, include occasional spontaneous transitions the goal (probability 0.1/(grid size)): The agent location is simply represented by a one-hot vector (effectively tabular) so there is really structure to exploit. The learned model must essentially memorize every individual transition to learn the dynamics.
Our experimental design was the same in Section 4. We tune the Q-network step size and softmax exploration temperature from the same set of values on a grid size of size 12 and then used the best hyperparameters for each agent the other grid sizes. In this experiment, we focused on the low- data regime where the simple model tended to have the biggest advantage over ER in our other experiments.
Return Grid Size
Grid Size 12 Hyperparameter Sensilivity
-10-
-15
20 Ratum
-20
Fini
-30
Copelie Roubly
-40
-30
Pactor Mooe
-40
-35
7.5
10.0
12.5
15.0 17.5
20.0
-10 -6
Size
(Base Tomperature
Figure 6: Left: Final performance of greedy policy v.s grid size for OpenGrid in the low data regimes, that is 100 thousand interactions with 10 updates per step. Right: Softmax temperature and step-size sensitivity curves for each approach resulting from the grid-sean grid-search on size 12 OpenGrid. In these plots, the other hyperparameter fixed to its best value from the gric-search while varying the temperature step-size. Error bars show 95% confidence intervals.
The results are shown in Figure 6 with performance v.s. grid-size on the left and hyperparameter sensitivity curves from the initial tuning the right. In contrast to other experiments, here see that the simple model becomes worse relative to ER as the environment complexity increases. the model has no ability extrapolate beyond the data. The best it can do
This is reasonable is memorize what is already in the replay buffer and the limitations of finite model capacity and imperfect optimization prevent from doing so perfectly. This result helps contextualize our main results for euvironments with factored structure by showing how the performance of the model-base approach suffers in simple environment without such structure.
D MODEL DETAILS
Our latent-space model consists of the following components:
Representation Model: 6, (p((lo))
Observation Reconstructor: â ~ po( (cb)))
Transition Predictor: &t ~ po a1-1)
Reward Predictor: f ~ po (fel0t-1,, a
Termination Predictor:
Po po 11.
All components are implemented as NNs with 0 representing the combined parameter vector. or the observation from the environment at time i, is the action, is the reward, is the continu-
19

---

Gr-1) ) uses a sigmoid activation to output Under review as a conference paper at ICLR 2023
ation probability[ is the latent-state constructed by the model from the observation from which transitions, rewards and terminations are all predicted. The associated versions of each variable with hats are predictions made by the model. The Q-netwrrk associated with the latent-space mod- els are always trained to predict action-values directly from Φi as opposed to first reconstructing the observation.
For training the model, we use a loss very similar to that employed by Hafner et al. (2020)
Le( log(po(o:lo=)) log(po(r -,
(og .,-1)
+
where de ~ go (ox We also employ KI.-balancing, as described by Hafner et al. (2020), with a = 0.8. We experiment with both Gaussian and Categorical latent variables for Φt. For the categorical case we use the straight-through estimator to propagate gradients through the discrete latent variables where necessary.
all cases, we train the model on randomly sampled transition from replay- buffer. Note that it is not necessary to train on sequences in our case, as the lack of recurrence means that the loss at each stme-step can be independently evaluated. The observation reconstructor po(Orlor, uses sigmoid activation to output the means of a vector of Bernoulli distributions since all tested environments use binary observations. The reward predictor pe (fr| (t-1)) uses a linear activation and outputs the mean of univariate Gaussian, in which case the above loss is effectively mean-squared error. The termination predictor pe single Bernoulli termination probability.
For the simple model, the reward and termination predictors are the same except that they take raw observations input instead of latent variables. Likewise the transition predictor po a--1) works directly in observation space, and is trained with a negative log-likelihood loss relative to the true observations.
E ILLUSTRATIVE EXPERIMENT DETAILS
Here we give some additional detail on the setup the illustrative experiment in Section [ For the most part, we used the same hyperparameters as main experiments, detailed in Appendix F The only exception is that for this simple experiment we did not tune any hyperparameters, but rather fixed the Q-learning step-size 2c 4 and softmax exploration temperature to 0.1. The softmax exploration temperature only applies to the model in this case, since the model-free agent was trained fixed data and thus never selects actions during learning.
We train all agents for 1,000,000 training steps, to insure convergence, which was excessive given the small fixed dataset used for training. To control for the total number of value function updates, and the total number of (real or imagined) transitions used in each update, we made the following choices: for model-free DQN and the 1-step model-based agent the value function is update on batches of 320 transitions, from the dataset or the learned model; the 10-step model-based agent uses a batch of 32 sequences of length 10 generated by the model to equate the number of transitions per update. The models are always trained on 32 transitions from the dataset in each update.
'Two out of three of the environments we experiment with are continuing, thus termination will never occur and this prediction could be omitted, however, 7/ should be learned casily and thus it should make little difference whether it is included or not.
20

---

Under review as a conference paper at ICLR 2023
F HYPERPARAMETERS
Shared Hyperparameters
Value
Number of Hidden Layers Number of Hidden Units
3
200
Hidden Activation
ELU
Optimizer
Adam
Adam 31
0.9
Adam $2
0.99
Adam
le-5
Adam weight decay
1e-6
Q-learning Step-Size
Tuned (See Appendix G
Discount Factor
0.9
Batch Size
32 for model-based, 320 for model-free
Exploration Strategy
Softmax
Sofumax Temperature
Tuned (See Appendix G
Target Network Update Frequency
100
Buffer Size
100,000
Training Start Time
1000
Model Hyperparameters
Model Learning Step-Size
2e-4
Rollout Length
10
Categorical Latent Hyperparameters
Number of Features
32
Width of Features
32
KL Balancing
0.8
KI. Loss Scale
1.0
Gaussian Latent Hyperparameters
Number Features
32
KL Balancing
0.8
KL Loss Scale
1.0
Minimum std
0.1
std activation
(x/2)
Table 1: Table of hyperparameters used experiments in Section
G HYPERPARAMETER TUNING AND SENSITIVITY EXPERIMENTS
Here, we present hyperparameter sensitivity plots for step-size and softmax exploration temperature which resulted from our initial grid-search to select hyperparameters for the main experiments in Section[4 The initial grid-search tried the set (0.0125, 0.025, 0.05, 0.1, 0.2, 0.4, 0.8, 1.6, 3.2) for the softmax temperature and the set 21.25e- 05, 2.50e 05, 5.00e - 05, 1.00e 04, 2.00e - 04, 4.00e 04, 8.00e 04, 1.60e 03, 3.20e 03} for the step-size. This range was extended in some cases where there was significant positive trend at the boundary of the range for either parameter, this only effected results for ER and the Gaussian latent model and the impact was negligible compared to using the best parameters in the initial range. In each case, the seeds was used to evaluate each hyperparameter setting in terms of final performances of the greedy policy. The hyperparameters with the best final performance were selected in each case for in our main experiments. Sensitivity curves for step-size and temperature are shown in Figure 7 and Figure respectively.
performance of 30 random
21

---

Under review as a conference paper at ICLR 2023
ProcMaze
ButtonGrid
PanFlute
0.10
Pxpeeween Repdey
High Data
-10
0.150
Parkic Doca
0.08
Final Retum -20
3.100
Conpurtel Loken
Reward
0.00
0.073
-50
0.02
0.023
20.0
117.5
-150 -125 -10.5
-1
-14
-32
-10
-10
-14
-12
-16
Data
0.150
0.00
-20
Ratum
0.100
Reward
-40
-50
0.02
0.050
-00
Log Step
Step-Sow
Figure 7: Step-size sensitivity curves for each approach resulting from the grid-search termediate level difficulty for each environment (size 4 ProcMaze, button ButtonGrid, 7 pipe PanFlute). In these plots, the softmax temperature is fixed its best value from the grid-search while varying step-size. The search was extended in few cases when there was a significant posi- tive trend the boundary of the initial search grid. Error bars show 95% confidence intervals.
in-
ProcMaze
PanFiule
High Data
Doedee Pepply
0.08
-20
0.255
Sinpa Vodal
Ralam/
Rhat
-30
0.000
Calegrtoil Loard
50
0.02
0.000
-60
0.025
Data
-10
1150
-20
1125
Final
0.05
Retum/
-30
0075
-10
0.00
BBoos Theppeatiune
(Baose Temparation
(Baan 2) Terpanitre
Figure 8: Softmax temperature sensitivity curves for each approach resulting from the grid-search on an intermediate level of difficulty for each environment (size 4 ProcMaze, 4 button ButtonGrid, pipe PanFlute). In these plots, the step-size is fixed to its best value from the grid- rearhh while varying softmax termperature. The search was extended in few cases when there was significant positive trend at the boundary of the search grid. Error bars show 95% confidence intervals.
22

---

Under review as a conference paper at ICLR 2023
H LEARNING CURVES
High Data
ButtonGrid
Button
Buttons
3 Butions
Buttons
5 Buttons
Buttons
Bulloos
Buttons
0.
Dime
10:0
140
6PA
160
1.90
100
PanFlute
Pipes
5 Pipes
Pipes
Pipes
Pipes
10
$o.00
Baus
2.10
20.09
gass
10.00
0.0%
ProcMaze
Experience Replay
Grid
Grid
Perfect Model
Simple Model
Categorical Latent
Low Data
Gaussian Latent
ButtonGrid
Button
Buttons
Buttons
Buttons
Buttons
Buttons
7 Buttons
Buttons
suce
33000
1000C0
103000
50000
100069
100000
10c
PanFlute Pipes
5 Plpes
6 Ploes
7 Pipas
B Plpes
Pipes
Taus
MED3
MED3
10C
LOCOA
ProcMaze
Grid
Grldd
50000 100000
Figure 9: Full learning curves for experiments in Section Performance of the greedy policy is plotted every 5000 updates and smoothed with a moving average over the last 10 values. Error bars show 95% confidence intervals.
I ABLATION EXPERIMENTS
In this section, we present some additional ablation studies to better understand the impact of some of our experimental design decisions made the main paper. Figure T0 highlights the impact of removing spontaneous transitions to rewarding states in each of the environments. Figure []] compares the performance of the simple model with 1-step and 10-step model rollouts, where the total number of simulated transitions used in each batch used for DQN updates is controlled.
23

---

Under review as a conference paper at ICLR 2023
ProcMaze
ButtonGrid
PanFlute
0.30
High Data
-10
0.20
Vodal
0.25
-20
0.20
Returi-
-30
0.15
Reward
Rate
40
0.10
50
0.05
-60
0.00
Low Data
-10
0.20
-20
Retur
Reward
-0
Rate
40
-50
-60
0.00
3.0
3.5
4.0
4.5
5.0
5
10
Grid Size
Numbar Buttons
Number Pipes
Figure 10: Comparing the performance of ER and the simple model in variants each environment class with spontancous rewards disabled to the original environment. Dotted lines show the curve with spontancous rewards disabled. The effect of this is highly variable, but is generally detrimental to the model- approach. Perhaps surprisingly, ER appears perform better without the random transitions to rewarding states in ButtonGrid, perhaps due to elimination of the resulting noise from the learning signal. The impact was largest in PanFlute, where the absence of spontaneous rewards negatively effected both model-free and model-based approaches, but in the high-data regime lead model-free outperform model-based. Error bars show 95% % confidence intervals.
ProcMaze
ButtonGrid
PanFlute
High Data
Parted Model
0.20
Simoble Vedil
Final -20
0.15
Retum/
Reward
30
0.05
50
Low Data
0.20
0.25
Final
0.20
Retum/
0.15
Reward
0.15
Rate
0.10
0.10
-60
0.05
0.00
4.0
4.5
5.0
Grid Size
Number Buttons
Number of Pipes
Comparing simple model performance with 1-step and 10-step rollouts, with perfect model included for reference. Dotted line is 1-step rollouts, solid line is 10-step. In ProcMaze and PanFlute, 10-step rollouts consistently perform much better, however is ButtonGrii 1-step rollouts perform better for lower button counts in the high data regime, while 10-step rollouts generally perform better in the low data regime. Error bars show 95% confidence intervals.
Figure 11:
24


Reviews:
[{"review_id": "vm9WO0rQa-", "reviewer": "Reviewer_JiYh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper provides a high-level characterization of the properties of RL tasks in which model-based algorithms can be expected to improve generalization to new observations. Two primary quantities arise from this analysis: first, the state space should exhibit a simple underlying factored structure which the inductive bias of the model is well-suited to learning; second, the distribution of rewarding trajectories that an agent sees should be sufficiently sparse as to make learning from Bellman backups difficult, but still provide non-zero reward density so that there is some signal for a learned model to leverage.", "summary_of_the_review": "Overall, this paper provides a nice set of interpretable examples that highlight the different behaviour of model-free and different model-learning algorithms. One of these examples in particular illustrates how the smoothness bias in DNN function approximators can facilitate learning and generalization in model-based algorithms. The paper suffers two main limitations: first, its main theoretical result is a relatively weak statement, and doesn’t directly offer a concrete set of implications for the effect of model-learning methods on generalization. Second, it is not clear whether the empirical findings concerning the behaviour of model-based and model-free methods will generalize outside of the small environments presented in the paper.\n", "strength_and_weaknesses": "### Strengths\n\n- The paper provides nice intuitive examples of generalization in maze environments that highlight the limitations of model-based and value-based RL algorithms.\n- The paper is well-written and does a good job at summarizing and emphasizing the key take-aways from each section.\n- The experiment setup is clear.\nThe analysis of different learned models was enlightening as it highlighted the importance of the inductive bias of the world model, and because it provided a striking but also simple setting where a learning algorithm benefits more from an imperfect learned model than from the perfect world model. I appreciated that the authors went further to validate their hypothesis that the learned models were benefitting from reward smoothing in these contexts\n- Theorem 1 provides a nice contrast to the findings of Parr et al. in the linear setting. \n\n### Weaknesses\n\n- The paper focuses on a handful of small toy environments. While this is useful to guide the reader’s intuition and to make the experimental findings interpretable, it also limits the generality of the paper’s conclusions.\n- The degree to which the findings in this paper also apply to larger scale models is an open question. While the toy settings considered in the paper do a good job of illustrating some environment qualities can than influence the efficacy of simple model-based algorithms, it’s not clear to me that these are the only parameters that can influence the relative performance of model-learning methods in more challenging tasks. For example, the influence of model learning on exploration, or on the stability of the underlying learning method, could also play a role in an agent’s final performance in these settings. It’s not clear to me whether the environment features and model smoothing properties identified in this paper would be a dominant factor in performance compared to these other factors. \n\n- I suspect that the value-based agents might be worse at generalization than policy-based agents based on prior literature on e.g. the ProcGen suite (Raileanu et al., 2020), and would have liked to see the inclusion of a policy-gradient method in the set of baselines. \n- The presentation of the proof of theorem 1 could be improved: the definition of the transition dynamics for $S_t[0]$ for example, seem to be defining different quantities in the top and bottom line of the bracket, and the hypothesis class $H$ from which the q-functions are selected is not defined. I assume it is the set of all real-valued functions over the state space, but this should be stated explicitly. \n- While Theorem 1 is a nice observation, the statement itself is relatively weak. It does not, for example, provide a means of quantifying the degree to which the hypothesis space can be reduced by model classes with certain properties. Nor does it address settings where the algorithm may need to trade off between a limited set of environment steps and the approximation error of a model. This limits the significance of the theoretical contribution of the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "mtAUORtz_FQ", "reviewer": "Reviewer_8CUu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper investigates the question of when and why model generalization is better than value function generalization in reinforcement learning. The paper starts with a theoretical result that gives the intuition on this question: reinforcement learning with a learned model of the transition dynamics can reduce the optimal action value function space than simply solving bellman constraints with experience replay. Then the authors provide some hypothesis/conditions that tends to favor model-based RL, and then justify them by testing on some designed toy experiments.", "summary_of_the_review": "Interesting question, nice examples and discussion of empirical observation. Needs more thorough and rigorous theoretical study on some key questions.", "strength_and_weaknesses": "Strength:\n- An interesting and important question to study, as when and why model-based RL is guaranteed to be better than experience replay can significantly influence us to design practical RL algorithms for real-world applications\n- Well-written and easy to follow, with many concrete examples and interesting discussions of empirical observations\n- Good introduction of the background\n\nWeaknesses/Questions:\n- The first motivating example in the introduction (an agent trying to cross a stream to reach the food) is not very clear and convincing. It is hard to understand the formulation of the state (is it a bandit problem given your state definition?), and what is the data (some state+reward pairs?) in your experience replay, and why model-free learning clearly/definitely fails (it can also generalize via value function)?\n- Perhaps the main rigorous answer is Theorem 1, where I was expecting to see when and why model-based RL is better. However, I found the result is not surprising/contains enough theoretical insights. It is basically saying, if we are given a set of hypothesis class for the transition dynamics, and we can successfully identify one that perfectly matches the data, then using that transition function (probably optimal if data is enough) to solve Bellman optimal equation is better than directly using the data to solve. This might be trivial and in practice, the key question is actually what theoretical guarantee we can achieve when the approximate model makes an error (upper bounded by some constant). If the data D is very limited and the best model consistent with D is still far from the true dynamics, then is model-based RL always better (since it incurs compounding error)? Another essential question to study is how will the model error propagates to solving Bellman optimality equations (the data is always from true dynamics but generated ones can be wrong)?\n- It seems most above important questions are not rigorously answered. The following sections provide nice examples and hypothesis on the conditions. But I feel they should be an empirical evidence to some more general theorems. I was expecting a formal sufficient (maybe not necessary) and hopefully general enough condition for when model generalization + compounding error is better than ER.\nTheoretically this should depend on specific algorithm, data collecting policy and sample size, etc. Otherwise, it may be hard to summarize some general principles that guides the design of algorithm for more complicated and realistic environments.\n- The observation that sometimes learned model is even better than the true model is very interesting. In principle, even using more smooth model can help exploration or optimization, this should eventually lead to bias. Will this be an issue and can we alleviate it using annealing (gradually reduce to true model)? In case of no bias, it seems to be something like a reward shaping or simply a coincidence that two MDPs have same solution. This may deserve more thorough study.\n- \"factored structure\" appears many times but seems fuzzy. Can you give a formal definition?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "IPVF3CHduO", "reviewer": "Reviewer_4bHT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the benefits of model-based generalization by comparing the performance of DQN on the rollouts of a learned model, and DQN on a replay buffer. Theoretically, this paper proves that learning a transition model can be more efficient (in terms of pruning invalid solutions) than learning a Q-function using Bellman equations. Empirically, this paper proposes three environment characteristics where the model-based methods could be better, and demonstrates them using constructed toy environments. \n\n**Summary of the experiments**\n\nSetup: \n\nExperience replay (ER): Run DQN with a replay buffer of the past trajectories (called $\\mathcal{D}$)\n\nModel-based: use $\\mathcal{D}$ to train a model. And then train a DQN on fake trajectories generated by the learned model starting from a random state from $\\mathcal{D}$ (similar to MBPO).\n\nThe goal is to make the training of DQN exactly the same except for the data distribution (ER vs. model rollouts).\n\nResults: \n\nThis paper compares ER and model-based methods on toy environments in the online RL setting. The number of updates to the Q-network is the same for model-based and ER. In every update, the Q-network is trained using 320 samples from the replay buffer for ER, and 32 model rollouts with length 10 for model-based methods (so the batch size is also the same). This paper shows that for the toy environments model-based is in general better than ER (Figure 3).", "summary_of_the_review": "On the plus side, the problem setup of this paper is neat and interesting, the results are clean and intuitive. On the minus side, more discussions about the practical relevance and ablation studies are needed since the results are mainly empirical. As a result, I recommend a weak rejection for the current version of this paper.", "strength_and_weaknesses": "Strengths:\n\n- This paper studies the benefit of model-based generalization by comparing model-based rollouts with experience replay. To some extent, this comparison is much fairer because the only difference is the distribution of training data (replay buffer vs model rollouts). This approach is neat and results in some interesting theoretical results (Theorem 1).\n- The illustrative example in Section 2 is very easy-to-follow and already well demonstrates the intuitions of this paper.\n\nWeaknesses:\n\n- This paper only discusses one side of the story. In fact, model-based method outperforms experience replay in all the environments studied in this paper. Together with Theorem 1, is it possible to make the claim that model-based methods are always better than experience replay? In other words, discussions and ablation studies about the necessity of the three characteristics in Section 3 are missing. Hence, the scope of the conclusions of this paper is unclear.\n- It’s unclear to me when and how the characteristics proposed in Section 3 can be applied to benchmarking environments such as OpenAI Gym. For example, if we predict whether model-based methods are better using the intuition in this paper, does the prediction match reality? Since this is mostly an empirical paper, I would expect more discussions on the practical impact of the conclusions.\n- the experiments provide one possible explanation for why MB is better than ER --- learning the transitions can be easier because of the implicit bias of the function class (e.g., the factored structure of the transition) and therefore MB has better sample efficiency. This explanation is somewhat expected and whether it is the main factor in benchmarking experiments such as Mujoco is unclear.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "BkIKylM7Fq", "reviewer": "Reviewer_5uCS", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work formalizes the generalization benefits of model-based RL methods in the context of deterministic episodic MDPs who's transitions probability belong to a known hypothesis class. The authors then discuss the intuition gained from their theory in various illustrative environments and take a closer look at the case when transitions can be assumed to be factored.", "summary_of_the_review": "Assuming the novelty isn't brought in question, this paper seems like a clear accept. Ideas are well motivated and of interest to RL.", "strength_and_weaknesses": "This work examines an important open question about the benefits of model-based RL. Although these methods are commonly believed to be more sample efficient, formal results on the subject are sparse. This work will likely be of interest to the RL community.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject