High-Accuracy Model-Based Reinforcement Learning, a Survey

One of the shortcomings of conventional reinforcement learning methods is that they only focus on expected value, ignoring the variance of values. This is problematic when few samples are taken for each trajectory $\tau$. Uncertainty modeling methods have been developed to counter this problem. Gausian processes can learn simple processes with good sample efficiency, although for high-dimensional problems they need many samples. They have been used for probabilistic inference to learn control (Deisenroth and Rasmussen, 2011) in the PILCO system. This system was effective on Cartpole and Mountain car (Figure 11), but does not scale to larger problems.

A related method uses nonlinear least-squares optimization (Tassa et al., 2012). Here the model learner uses quadratic approximation on the reward function, which is then used with linear approximation of the transition function. With further enhancements this method was able to teach a humanoid robot how to stand up (see Figure 9).

We can also sample from a trajectory distribution optimized for cost, and to use that to train the policy, with a policy-based method (Levine and Koltun, 2013). Then we can optimize policies with the aid of locally-linear models and a stochastic trajectory optimizer. This approach, called Guided policy search (GPS), has been shown to train complex policies with thousands of parameters learning tasks in MuJoCo such as swimming, hopping and walking. Alternatively, we can compute value gradients along the real environment trajectories, instead of planned ones, and re-parameterize the trajectory through sampling, to mitigate learned model inaccuracy (Heess et al., 2015). This was done by Stochastic value gradients (SVG) with global neural network value function approximators.

Learning arm and hand manipulation directly from video camera input is a challenging problem in robotics. The camera image provides a high dimensional input and increases problem size and complexity of the subsequent manipulation task substantially. Both Finn and Levine (2017); Ebert et al. (2018) introduce a method called Visual foresight. This system uses a training procedure where data is sampled according to a probability distribution. Concurrently, a video prediction model is trained. This model generates a sequence of future frames based on an image and a sequence of actions, as in GPS. At test time, the least-cost sequence of actions is selected in a model-predictive control planning framework (see Section 3.2.2). This approach is able to perform multi-object manipulation, pushing, picking and placing, and cloth-folding tasks (which adds the difficulty of material that changes shape as it is being manipulated).

Ensemble methods, such as a random forest of decision trees (Ho, 1995), are widely used in machine learning (Bishop, 2006), and they are also used in controlling uncertainty in high dimensional modeling. Ensemble methods mitigate variance and improve performance by running algorithms multiple times. They are used with success in model-based deep reinforcement learning as well. Chua et al. (2018) combine uncertainty-aware modeling with sampling-based uncertainty propagation, creating a method called Probabilistic ensembles with trajectory sampling, PETS. (This approach is described in the next section, see Algorithm 6). An ensemble of probabilistic neural network models is used by Nagabandi et al. (2018). Ensembles perform well; performance on pusher, reacher, and half-cheetah (see Figure 9) is reported to approach asymptotic model-free baselines such as PPO (Schulman et al., 2017). Ensembles of probabilistic networks (Chua et al., 2018) are also used with short rollouts, where the model horizon is shorter than the task horizon (Janner et al., 2019). Results have been reported for hopper, walker, and half-cheetah, again matching the performamce of model-free approaches.

The ensemble approach is related to meta learning, where we try to speed up learning a new task by learning from previous, related, tasks (Brazdil et al., 2008; Hospedales et al., 2020; Huisman et al., 2021). MAML is a popular meta learning approach (Finn et al., 2017), that attempts to learn a network initialization $x$ such that for any task $M_{k}$ the policy attains maximum performance after one policy gradient step. The MAML approach can be used to improve model accuracy by learning an ensemble of dynamics models and by then meta-optimizing the policy for adaptation in each of the learned models (Clavera et al., 2018). Results indicate that such meta-learning of a policy over an ensemble of learned models indeed approaches the level of performance of model-free methods with substantially better sample complexity.

The next group of methods that we describe are the latent models. Central to all our approaches is the need for improvement of model accuracy in complex, high-dimensional, problems. The main challenge to achieve high accuracy is to overcome the size of the high-dimensional state space. The idea behind latent models is that in most high-dimensional environments there are elements that are less important, such as background trees that never move, that have little or no relation with the reward of the agent’s actions. The goal of latent models is to abstract away these unimportant elements of the input space, reducing the effective dimensionality of the space. They do so by learning the relation between the elements of the input and the reward. When we focus our learning mechanism on the changes in observations that are correlated with changes in these values, then we can improve the efficiency of learning high-dimensional problems greatly. Latent models thus learn a smaller representation, smaller than the observation space. Planning takes place in this smaller representation space.

The value prediction network (VPN) was introduced by Oh et al. (2015, 2017) to achieve this goal. They ask the question in their paper: “What if we could predict future rewards and values directly without predicting future observations?” and describe a nework architecture and learning method for such focused value prediction models. The core idea is not to learn directly in actual observation space, but first to transform the actual state respresentation to a smaller latent representation model, also known as abstract model. The other functions, such as value, reward, and next-state, then work with the smaller latent representations, instead of the actual high-dimensional states. By training all functions based on the values (Grimm et al., 2020), planning and learning occur in a space where states are encouraged only to contain the elements that influence value changes. In VPN the latent model consists of four networks: an encoding function, a reward function, a value function, and a transition function. All functions are parameterized with their own set of parameters (Figure 8). Latent space is lower-dimensional, and training and planning become more efficient. The figure shows a single step rollout, planning one step ahead, as in Dyna-Q (Algorithm 4).

The training of the networks can in principle be performed with any value-based reinforcement learning algorithm. Oh et al. (2017) report results with $n$-step Q-learning and temporal difference search (Silver et al., 2012).

VPN (Oh et al., 2017) showed impressive results on Atari games such as Pacman and Seaquest, outperforming model-free DQN (Mnih et al., 2015), and outperforming observation-based planning in stochastic domains. Subsequently, many other works have been published that further improved results (Kaiser et al., 2019; Hafner et al., 2018, 2019, 2020; Sekar et al., 2020; Silver et al., 2017b; Ha and Schmidhuber, 2018b). Many of these latent-model approaches are complicated designs, with multiple neural networks, and different learning and planning algorithms.

The latent-model approach is related to world models, a term used by Ha and Schmidhuber (2018a, b). World models are inspired by the manner in which humans are thought to contruct a mental model of the world in which we live. World models are often generative recurrent neural networks that are trained unsupervised using a variational autoencoder (Kingma and Welling, 2013, 2019; Goodfellow et al., 2014) and a recurrent network. They learn a compressed spatial and temporal representation of the environment. In world models multiple neural networks are used, for a vision model, a memory model, and a controller (Ha and Schmidhuber, 2018b). By using features extracted from the world model as inputs to the agent, a compact and simple policy can be trained to solve a task, and planning occurs in the compressed or simplified world. World models have been applied by Ha and Schmidhuber (2018a, b) on a car racing game (Kempka et al., 2016). The term world model actually goes back to 1990, where it was used by Schmidhuber (1990b). Latent models are also related to dimensionality reduction (Van Der Maaten et al., 2009).

The architecture of latent models, or world models, is elaborate. The dynamics model typically includes an observation model, a representation model, a transition model, and a value or reward model (Karl et al., 2016; Buesing et al., 2018; Doerr et al., 2018). The task of the observation model is to reduce the high-dimensional world into a lower-dimensional world, to allow more efficient planning. Often a variational autoencoder or LSTM is used.

The Arcade learning environment is one of the main benchmarks in reinforcement learning. The high-dimensionality of Atari video input has long been problematic for model-based reinforcement learning. Latent models were instrumental in reducing the dimensionality of Atari, producing the first successes for model-based approaches on this major benchmark.

Related to the VPN approach (Oh et al., 2015, 2017) is other work that uses latent models on Atari, such as Kaiser et al. (2019), that is aimed at video prediction, outperforming model-free baselines (Hessel et al., 2017), reaching comparable accuracy with up to an order of magnitude better sample efficiency. The approach by Kaiser et al. (2019) uses a variational autoencoder (VAE) to process input frames, conditioned on the actions of the agent, to learn the world model, using PPO (Schulman et al., 2017). The policy $\pi$ is then improved by planning inside the reduced world model, with short rollouts. This behavior policy $\pi$ then determines the actions $a$ to be used for learning from the environment.

Latent models are also used on continuous MuJoCo problems. Here the work by Hafner et al. (2018, 2019) on the PlaNet and Dreamer systems is noteworthy including the application of their work back to Atari (Hafner et al., 2020), which achieved human-level performance. PlaNet uses a Recurrent state space model (RSSM) that consists of a transition model, an observation model, a variational encoder and a reward model (Karl et al., 2016; Buesing et al., 2018; Doerr et al., 2018). Based on these models a Model-predictive control agent is used to adapt its plan, replanning each step (Richards, 2005). The RSSM is used by a Cross entropy method search (Botev et al. (2013), CEM) for the best action sequence. In contrast to model-free approaches, no explicit policy or value function network is used; the policy is implemented as MPC planning with the best sequence of future actions. PlaNet is tested on continuous tasks and reaches performance that is close to strong model-free algorithms. A further system, called Dreamer (Hafner et al., 2019), builds on PlaNet. Using an actor critic approach (Mnih et al., 2016) and backpropagating value gradients through predicted sequences of compact model states, the improved system solves a diverse collection of continuous problems from the Deepmind control suite (Tassa et al., 2018), see Figure 9. Dreamer is also applied to discrete problems from the Arcade learning environment, and to few-shot learning (Sekar et al., 2020). A further improvement achieved human-level performance on 55 Atari games, a first for a model-based approach (Hafner et al., 2020), showing that the latent model approach is well-suited for high-dimensional problems.

As we saw in Section 2.1, methods for continuous action spaces typically sample full trajectory rollouts to get stable actions. At each planning step, the transition model $T_{a}(s)\rightarrow s^{\prime}$ computes the new state, using the reward to update the policy. Due to the inaccuracies of the internal model, planning algorithms that perform many steps will quickly accumulate model errors (Gu et al., 2016). Full rollouts of long and inaccurate trajectories are therefore problematic. We can reduce the impact of accumulated model errors by not planning too far ahead. For example, Gu et al. (2016) perform experiments with locally linear models that roll out planning trajectories of length 5 to 10. This reportedly works well for MuJoCo tasks gripper and reacher.

In their work on model-based value expansion (MVE), Feinberg et al. (2018) also allow imagination to fixed depth, value estimates are split into a near-future model-based component and a distant future model-free component. They experiment with model horizons of 1, 2, and 10. They find that 10 generally performs best on typical MuJoCo tasks such as swimmer, walker, and cheetah. The sample complexity in their experiments is better than model-free methods such as DDPG (Silver et al., 2014). Similarly good results are reported by Janner et al. (2019); Kalweit and Boedecker (2017), both approaches use a model horizon that is much shorter than the task horizon.

Taking the idea of shorter trajectories for planning than for learning further, we arrive at Model-predictive control (MPC) (Kwon et al., 1983; Garcia et al., 1989). Model-predictive control is a well-known approach in process engineering, to control complex processes with frequent re-planning of a limited time horizon. Model-predictive control uses the fact that while many real-world processes are not linear, they are approximately linear over a small operating range. Applications are found in the automotive industry and in aerospace, for example for terrain-following and obstacle-avoidance algorithms (Kamyar and Taheri, 2014). In optimal control, four MPC approaches are identified: linear model MPC, nonlinear prediction model, explicit control law MPC, and robust MPC to deal with disturbances (Garcia et al., 1989). In this survey, we focus on how the principle of continuous replannning with a rolling planning horizon performs in nonlinear model-based reinforcement learning.

It is instructive to compare the MPC and linear quadratic regulators (LQR) approach, since both methods come from the field of optimal control. MPC computes the target function with a small time window that rolls forward as new information comes in; it is dynamic. LQR computes the target function in a single episode, using all available information; it is static. We observe that in model-based reinforcement learning MPC is used in the planning part with the behavior policy $\pi$ being the target and the transition function $T_{a}(\cdot)$ the input; for LQR the transition function $T_{a}(\cdot)$ is the target, and the environment samples $(s_{t},r_{t})$ are the input. Thus, one could conceivably use both MPC and LQR, the first as planning and the second as learning algorithm, in a model-based approach.

An iterative form of LQG has indeed been used together with MPC on a smaller MuJoCo problem (Tassa et al., 2012), achieving good results. MPC used step-by-step real-time local optimization; Tassa et al. (2012) used many further improvements to the trajectory optimization, physics engine, and cost function to achieve good performance.

MPC has also been used in other model learning approaches. Both Finn and Levine (2017); Ebert et al. (2018) use a form of MPC in the planning for their Visual foresight robotic manipulation system (that we have seen in a previous section). The MPC part uses a model that generates the corresponding sequence of future frames based on an image to select the least-cost sequence of actions.

Another approach uses ensemble models for learning the transition model, while using MPC for planning. PETS (Chua et al., 2018) uses probabilistic ensembles (Lakshminarayanan et al., 2017) for learning. In MPC fashion only the first action from the CEM-optimized sequence is used, re-planning at every time-step (see Algorithm 6).

MPC is a simple and effective planning method that is well-suited for model inaccuracy, by restricting the planning horizon. MPC is used with success in model-based reinforcement learning, with high-variance or complex transition models. MPC has also been used with success in combination with latent models (Hafner et al., 2018; Kaiser et al., 2019).

Our third planning approach is different, it integrates planning with learning in a fully differentiable algorithm. Let us see how this works.

Model-based reinforcement learning consists of two distinct functions: learning the transition model and planning with the model to improve the behavior policy (Figure 5). In classical, tabular, reinforcement learning, both learning and planning functions are designed by hand (Sutton and Barto, 2018). In deep reinforcement learning, one of these functions is approximated by deep learning—the model learning—while the planner is still hand-written. End-to-end learning and planning breaks this classical planning barrier. End-to-end approaches integrate the planning into deep learning, using differentiable planning algorithms, extending the backpropagation fully from reward to observation in all parts of the model-based approach.

How can a neural network learn to plan? While conceptually exciting and appealing, there are challenges to overcome. Among them are finding suitable differentiable planning algorithms and the increase in computational training complexity, since now the planner must also be learned.

The idea of planning by gradient descent exists for some time, several authors explored learning approximations of state transition dynamics in neural networks (Kelley, 1960; Schmidhuber, 1990a; Ilin et al., 2007). Neural networks are typically used to transform and filter, to learn selection and classification tasks. A planner unrolls a state, computes values, using selection and value aggregation, and backtracks to try another state. Although counter-intuitive at first, these operations are not that different from what classic neural networks are performing. A progression of papers has published methods on how this can be achieved.

We will start at the beginning, with convolutional neural networks (CNN) and value iteration. We will see how the iterations of value iteration can be implemented in the layers of a convolutional neural network (CNN). Next, two variations of this method are presented, and a way to implement planning with convolutional LSTM modules. All these approaches implement differentiable, trainable, planning algorithms, that can generalize to different inputs. The later methods use elaborate schemes with latent models so that the learning can be applied to different application domains.

Let us start to see what is possible with a CNN. A CNN can be used to implement value iteration. This was first shown by Tamar et al. (2016), who introduced value iteration networks (VIN). The core idea is that value iteration (VI, see Algorithm 2) can be implemented step-by-step by a multi-layer convolutional network: each layer does a step of lookahead. In this way VI is implemented in a CNN. The VI iterations for the Q-action-value-function are rolled out in the network layers $Q$ with $A$ channels. Through backpropagation the model learns the value function. The aim is to learn a general model, that can navigate in unseen environments.

VIN can be used for discrete and continuous path planning, and has been tried in grid world problems and natural language tasks. VIN has achieved generalization of finding shortest paths in unseen mazes. However, a limitation of VIN is that the number of layers of the CNN restricts the number of planning steps, restricting VINs to small and low-dimensional domains. Follow-up studies focus on making end-to-end learning and planning more generally applicable. Schleich et al. (2019) extend VINs by adding abstraction, and Srinivas et al. (2018) introduce universal planning networks, UPN, which generalize to modified robot morphologies. Value propagation (Nardelli et al., 2018) uses a hierarchical structure to generalize end-to-end methods to large problems. TreeQN (Farquhar et al., 2018) incorporates a recursive tree structure in the network, modeling the different functions of an MDP explicitly. TreeQN is applied to Sokoban and nine Atari games.

A further step is to model more complex planning algorithms, such as Monte Carlo Tree Search (MCTS), a successful planning algorithm (Coulom, 2006; Browne et al., 2012). This has been achieved to a certain extent by Guez et al. (2018) who implement many elements of MCTS in MCTSnets and (Guez et al., 2019). In this method planning is learned with a general recurrent architecture consisting of LSTMs and a convolutional network (Schmidhuber, 1990b) in the form of a stack of ConvLSTM modules (Xingjian et al., 2015). The architecture was used on Sokoban and boxworld (Zambaldi et al., 2018), and was able to perform full planning steps. Future work should investigate how to achieve sample-efficiency with this architecture.

The question whether model-based planning can be learned by a neural network has been studied by (Pascanu et al., 2017), who showed that imagination-based planning steps can indeed be learned for a small game with an LSTM. Related to this, imagination-augmented agents (I2A) has been designed as a fully end-to-end differentiable architecture for model-based imagination and model-free reinforcement learning (Weber et al., 2017). It consists of an LSTM-based encoder (Chiappa et al., 2017; Buesing et al., 2018), a ConvLSTM rollout module, and a standard CNN-based model-free path. The policy improvement algorithm is A3C. Weber et al. (2017) report that on Sokoban and Pacman I2A performs better than model-free learning and MCTS. I2A has been specifically designed to handle model imperfections well and uses a manager or meta-controller to choose between rolling out actions in the environment or by imagination (Hamrick et al., 2017).

In VIN there is a tight connection between the network architecture and the application structure. One way to remedy this restriction is with a latent model, such as the ones that were discussed earlier. One of the first attempts is the Predictron (Silver et al., 2017b), where the familiar four elements appear: a representation model, a transition model, a reward model, and a value model. The goal of the latent model is to perform value prediction (not state prediction), including being able to encode special events such as “staying alive” or “reaching the next room.” Predictron performs limited-horizon rollouts, and has been applied to procedurally generated mazes.

One of the main success stories of model-based reinforcement learning is AlphaZero (Silver et al., 2017a, 2018). AlphaZero combines planning and learning in a highly successful way. Inspired by this approach, a fully differentiable version of this architecture has been introduced by Schrittwieser et al. (2020) which is named MuZero. This system is able to learn the rules and to learn to play games as different as Atari, chess, and Go, purely from the environment, with end-to-end learning and planning. The MuZero architecture is based on Predictron, with an abstract model consisting of a representation, transition, reward, and prediction function (policy and value). For planning, MuZero uses an explicitly coded version of MCTS that uses policy and value input form the network (Rosin, 2011), but that is executed separately from the network.

End-to-end planning and learning has shown impressive results, but there are still open questions concerning the applicability to different applications, and especially the scalability to larger problems.

Sequential decision problems are well-suited to model robotic actions, such as how to move the joints in a robotic arm to pour a cup of tea, how to move the joints of a humanoid figure to stand up when lying down, and how to develop gaits of a four-legged animal. The action space of such problems is continuous since the angles over which robotic joints move span a continuous range of values. Furthermore, the environment in which robots operate mimics the real world, and is non-deterministic. Things move, objects do not always respond in a predictable fashion, and unexpected situations arise.

In reinforcement learning, where agent algorithms are trained by the feedback on their many actions, working with real robots would get prohibitively expensive due to wear. Most reinforcement learning systems use physics simulations such as offered by MuJoCo (Todorov et al., 2012). MuJoCo allows the creation of experiments that provide environments for an agent. Tasks can range form small to large.

MuJoCo tasks differ in difficulty, depending on how many joints or degrees of freedom are modeled, and which task is being learned. Figure 9 shows some of the tasks that have been modeled in MuJoCo as part of the DeepMind Control Suite (Tassa et al., 2018). Some of the small tasks are ball-in-cup and reacher. The iterative quadratic non-linear optimization method iLQG (Tassa et al., 2012) is able to teach a humanoid to stand up, and Guided policy search (Levine and Koltun, 2013) and Stochastic value gradients (Heess et al., 2015) can learn tasks such as swimmer, reacher, half-cheetah (Figure 10) and walker. Also ensemble methods such as PETS, MVE, and meta ensembles achieve good results on these applications  (Gu et al., 2016; Chua et al., 2018; Feinberg et al., 2018; Clavera et al., 2018; Janner et al., 2019).

MuJoCo has enabled progress in model-based reinforcement learning in small continuous tasks, and also in larger tasks, such as how to develop gaits of a four-legged robotic animal, or how to scale an obstacle course. PlaNet (Hafner et al., 2018), Dreamer (Hafner et al., 2019) and MBPO (Janner et al., 2019) achieve good results on more complicated MuJoCo tasks using latent models to reduce the dimensionality.

There is a long tradition in reinforcement learning to see if we can teach a computer to play complicated games and puzzles (Plaat, 2020). Games and puzzles are often played on a board with discrete squares. Actions in such games are discrete, a move to square e3 is not a move to square e4. The environments are also deterministic, we assume that pieces do not move by itself.

Most games that are used in deep model-based reinforcement learning papers fall into this category. More complex games, such as partial information (card games such as poker (Brown and Sandholm, 2019)) or games with multiple actors (real-time strategy video games such as StarCraft (Vinyals et al., 2019; Ontanón et al., 2013; Wong et al., 2021)) are not used in the approaches that we survey here.

Among low-dimensional applications that are used in model-based reinforcement learning are simple pendulum problems, Cartpole and Mountain car, where the challenge is to reverse engineer the laws of impulse and gravity (Figure 11). The action space consists of two discrete actions, push left or push right, the environment is continuous and deterministic. PILCO (Deisenroth and Rasmussen, 2011) achieves good results with Gaussian process modeling and gradient based planning on the pendulum task.

Perhaps the most frequently used low-dimensional application area is grid-world, where various navigation tasks are tested (Figure 12). VIN (Tamar et al., 2016), VProp (Nardelli et al., 2018) and the Predictron (Silver et al., 2017b) that use maze navigation to test their approaches to integrating end-to-end learning and planning. Grid worlds and mazes can be designed and scaled in different forms and sizes, making them well suited for testing new ideas.

Other low-dimensional games are board games such as chess and Go. These games are low-dimensional (their input has few atttributes, in contrast to a mega-pixel image) but they nevertheless have a large state space. Finding good policies for chess and Go was one of the most challenging feats in reinforcement learning (Campbell et al., 2002; Silver et al., 2016; Plaat, 2020). Block puzzles such as Sokoban (Figure 13), are also often used to test reinforcement learning methods. Sokoban is a block-pushing puzzle that derives much of its complexity from the facft that the agent can push a box, but cannot pull (undo) a mistake, giving rise to many dead-ends that are hard to detect. It has been used by I2A (Weber et al., 2017) and MCTS network planning (Guez et al., 2019) approaches that implement planning by unrolling steps within a neural network.

Most recent success in model-free reinforcement learning has been achieved in high-dimensional problems, such as the Arcade Learning Environment (Bellemare et al., 2013), see for example (Mnih et al., 2015; Hessel et al., 2017). Atari games were popular video games in the 1980s in game arcades. Figure 14 shows a screenshot of Q*bert, a typical Atari arcade game.

Deep learning methods are well suited to process high-dimensional inputs. The challenge for model-based methods is to learn accurate models with a low number of observations. Latent model approaches have been tried for Atari with some success (Oh et al., 2015, 2017; Kaiser et al., 2019; Hafner et al., 2020). The MuZero approach was even able to learn the rules of very different games: Atari, Go, shogi (Japanese chess) and chess (Schrittwieser et al., 2020).