Planning with Expectation Models for Control

In model-based reinforcement learning (MBRL), Wan et al. (2019) showed conditions under which the environment model could produce the expectation of the next feature vector rather than the full distribution, or a sample thereof, with no loss in planning performance. Such expectation models are of interest when the environment is stochastic and non-stationary, and the model is approximate, such as when it is learned using function approximation. In these cases a full distribution model may be impractical and a sample model may be either more expensive computationally or of high variance. Wan et al. considered only planning for prediction to evaluate a fixed policy. In this paper, we treat the control case—planning to improve and find a good approximate policy. We prove that planning with an expectation model must update a state-value function, not an action-value function as previously suggested (e.g., Sorg & Singh, 2010). This opens the question of how planning influences action selections. We consider three strategies for this and present general MBRL algorithms for each. We identify the strengths and weaknesses of these algorithms in computational experiments. Our algorithms and experiments are the first to treat MBRL with expectation models in a general setting. Methods that scale with computation are most likely to stand the test of time. We refer to scaling with computation in the context of artificial intelligence (AI) as using more computation to provide a better approximate answer. This is in contrast to with the common notion of scaling in computer science as using more computation for solving a bigger problem exactly. The recent success of modern machine learning techniques, in particular that of deep learning, is Department of Engineering, University of Guelph, Guelph, Canada Department of Computing Science, University of Alberta, Edmonton, Canada Alberta Machine Intelligence Institute, Edmonton, Canada. Correspondence to: Katya Kudashkina <ekudashk@uoguelph.ca>. Pre-print. Copyright 2021 by the author(s). primarily due to their ability to leverage the ever-increasing computational power (driven by Moore’s Law), as well as their generality in terms of dependence on data rather than hand-crafted features or rule-based techniques. Key to building general-purpose AI systems would be methods that scale with computation (Sutton, 2019). RL is on its way to fully embrace scaling with computation— it already embraces scaling with computation in many aspects, for example by leveraging search techniques (e.g., Monte-Carlo Tree Search, see Browne et al., 2012; Finnsson & Björnsson, 2008), and modern deep learning techniques such as artificial neural networks. Methods that resort to approximating functions rather than learning them exactly have not been fully investigated. Extending the techniques used in the simpler tabular regime to this function approximation regime is an obvious first step, but some of the ideas that have served us well in the past might actually be impeding progress on the new problem of interest. For example, Sutton & Barto (2018) showed that when dealing with feature vectors rather than underlying states, the common Bellman error objective is not learnable with any amount of experiential data. Recently, Naik et al. (2019) showed that discounting is incompatible with function approximation in the case of continuing control tasks. Understanding function approximation in RL is key to building general-purpose intelligent systems that can learn to solve many tasks of arbitrary complexity in the real world. Many of the methods in RL focus on approximating value functions for a given fixed policy—referred to as the prediction problem (e.g., Sutton et al., 1988; Singh et al., 1995; Wan et al., 2019). The more challenging problem of learning the best behavior is known as the control problem—that is approximating optimal policies and optimal value functions. In the control problem, an RL agent learns within one of two broad frameworks: model-free RL and model-based RL (MBRL). In model-free RL, the agent relies solely on its observations to make decisions (Sutton & Barto, 2018). In model-based RL, the agent has a model of the world, which it uses in conjunction with its observations to plan its decisions. The process of taking a model as input and producing or improving a policy for interacting with a modeled environment is referred to as planning. Models and planning are helpful. One advantage of models and planning is that they are useful when the agent faces unfamiliar or novel ar X iv :2 10 4. 08 54 3v 1 [ cs .A I] 1 7 A pr 2 02 1 Planning with Expectation Models for Control situations—when the agent may have to consider states and actions that it has not experienced or seen before. Planning can help the agent evaluate possible actions by rolling out hypothetical scenarios according to the model and then computing their expected future outcomes (Doll et al., 2012; Ha & Schmidhuber, 2018; Sutton & Barto, 2018). Planning with function approximation remains a challenge in reinforcement learning today (Shariff and Szepesvári, 2020). Planning can be performed with various kinds of models: distribution, sample, and expectation. Wan et al. (2019) considered planning with an expectation model for the prediction problem within the function approximation setting. In this paper we extend Wan et al.’s (2019) work on the prediction problem to the more challenging control problem, in the context of stochastic and non-stationary environments. This will involve several important definitions. We start off with discussing important choices in MBRL (Section 1). This is followed by fundamentals on planning with expectation models in the general context of function approximation (Section 2). We then show (in Section 3) that planning with an expectation model for control in stochastic non-stationary environments must update a state-value function and not an action-value function as previously suggested (e.g., Sorg & Singh, 2010). Finally, we consider three ways in which actions can be selected when planning with state-value functions, and identify their relative strengths and weaknesses (Sections 4 & 5). 1. Choices in Model-Based RL Model-based methods are an important part of reinforcement learning’s claim to provide a full account of intelligence. An intelligent agent should be able to model its environment and use that model flexibly and efficiently to plan its behavior. In MBRL, models add knowledge to the agent in a way that policies and value functions do not (van Hasselt et al., 2019). Typically, a model receives a state and an action as inputs, and computes the next state and reward (Kuvayev & Sutton, 1996; Sutton et al., 2008; Hester & Stone, 2011). This output is used in planning to further improve policies and value functions. In this section we discuss three important choices one needs to make in MBRL. Learned models vs. Experience Replay. One choice to make is where planning improvements could come from: learned models or experience replay (ER) (Lin, 1992). In replay-based methods the agent plans using experience stored in the agent’s memory. Some replay-based examples include deep Q-networks (DQN) (Mnih et al. 2013; 2015) and its variations: double-DQN (van Hasselt et al., 2016); extension of DQN to include prioritized sweeping (Schaul et al. 2016); deep deterministic policy gradient (Lillicrap et al., 2016); rainbow DQN (Hessel et al., 2018). MBRL methods in which 1) the model is parameterized by some learnable weights, 2) the agent learns the parameters of the model, and 3) the agent then uses the model to plan an improved policy, are referred to as planning with learned parametric models. Learned parametric models are used in the Dyna architecture (Sutton, 1991), in normalized advantage functions method that incorporates the learned model into the Q-learning algorithm based on imagination rollouts (Gu et al., 2016), and in MuZero (Schrittwieser et al., 2019). The latter is a combination of a replay-based method and a learned model: the model is trained by using trajectories that are sampled from the replay buffer. If the environment is non-stationary, then the transitions stored in the replay buffer might be stale and can slow down or even hinder the learning progress. In this work, we focus on the learned parametric models. Types of Learned Models. Another important choice to make in MBRL is a choice of a model type. A model enables an agent to predict what would happen if actions are executed from states prior to actually executing them and without necessarily being in those states. Given a state and action, the model can predict a sample, an expectation, or a distribution of outcomes, which results in three model-type possibilities. The first possibility is when a model computes a probability p of the next state as a result of the action taken by the agent. We refer to such model as a distribution model. Such models have been used typically with an assumption of a particular kind of distribution such as Gaussian (e.g., Chua et al., 2018). For example, Deisenroth & Rasmussen (2011) proposed a model-based policy search method based on probabilistic inference for learning control where a distribution model is learned using Gaussian processes. Learning a distribution can be challenging: 1) distributions are potentially large objects (Wan et al., 2019); and 2) distribution models may require an efficient way of representing and computing them, and both of the tasks can be challenging (Kurutach et al., 2018). The second possibility is for the model to compute a sample of the next state rather than computing the full distribution. We refer to such model as a sample model. The output of sample models is more compact than the output of distribution models and thus, more computationally feasible. In this sense sample models are similar to experience replay. Sample models have been a good solution to when