Approximate Bayes-Optimal Reinforcement Learning and Planning

Abstract

This dissertation studies the optimal decision-making under uncertainty problem from the context of sequential experiment planning in autonomous laboratories. Our focus is on deep reinforcement learning methods, which we first motivate by building up from the challenges involved in conventional optimal experiment design methods and black-box numerical optimization. Secondly, we frame deep reinforcement learning in a unifying view of commonplace Bayesian optimization methods to conceptually streamline extensions or adoption of reinforcement learning methods in existing infrastructure. Each chapter investigates challenges and questions that emerge from our framing, which more specifically relate to uncertainty quantification and proper algorithm-hardware tailoring for improving performance scaling.

Bayesian inference of large complex models is typically computationally infeasible unless stringent, but possibly unrealistic, assumptions are placed on the model of the physical system that we’re trying to capture. Deep learning has since emerged as a powerful way to approximate this computationally infeasible routine through pattern recognition based on an abundance of data. This principle underpins most of the success stories of recent artificial intelligence breakthroughs, like agentic chatbots or game-playing agents. This dissertation investigates this principle from a more fundamental basis to improve our algorithmic tools for the overarching experiment-design use case.

The field of meta-reinforcement learning, similarly, has often used deep learning to capture various sources of uncertainty for decision-making agents, such as unknown rewards or world dynamics. This approximates what is known as Bayes-optimal control, an agent that optimally trades off between information-seeking actions and actions that maximize a (subjective) value. In Chapter 2 we show how common deep learning approaches, like recurrent neural networks, for meta-reinforcement learning do not appropriately capture true uncertainty from the Bayesian viewpoint. Rather, we found that uncertainty is obfuscated in the activations of the neural networks. We present the Laplace Variational Recurrent Neural Network to extract this implicit uncertainty as a means to enable practitioners to deal with uncertainty explicitly.

From Chapter 2 we built the insight that one of the roles that deep learning plays in deep reinforcement learning is to amortize an exponentially costly policy optimization problem. A well-known consequence of amortization is that we always have to deal with an amortization gap (approximation gap) between the true values and the predicted values. Improving the performance of an algorithm, thus, requires a choice of tools that minimize the runtime and data needed to reduce this gap. Among the most successful approaches in this space are those that use adaptive computation at inference time to reduce the amortization gap and complement the neural network. Monte-Carlo tree search, for instance, achieved strong data-efficiency and enabled breakthroughs like defeating the human world champion in Go. However, Monte-Carlo tree search-based methods do not favorably scale in runtime. Thus, in Chapter 3 we improve the state-of-the-art of sequential Monte-Carlo planners, which permit much more efficient parallelization on modern GPU hardware. Our newly proposed algorithm demonstrated improved scaling in performance given additional planning budget, not only in data-efficiency but also in runtime-efficiency compared to prior approaches.

The final Chapter 4 exploits our finding that sequential Monte-Carlo-based planning can improve runtime complexity by shifting computational budget from neural network learning to additional planning budget to get higher quality data. Since in meta-learning settings we typically deal with learning from trajectory-based data, we often find that when trajectories increase in length, any updates to a neural network architecture will become more expensive. This is because the computational cost of gradient updates grows with trajectory length, which causes bottlenecks in scaling. For instance, updates can start to exceed hardware memory limits, forcing sequential processing of smaller batches and reducing throughput. Thus, we again find that the performance scaling to additional planning compute in deep meta-reinforcement learning can be much improved when equipped with our designed sequential Monte-Carlo planner. In contrast to prior approaches, our planner explicitly accounts for how uncertainty evolves in light of new data. As a result, our method directly targets an optimal policy that trades off exploration with exploitation, without ad hoc exploration heuristics.

In summary, this dissertation advances the state of probabilistic reinforcement learning for uncertainty-based decision-making. We show that Bayesian optimization and deep reinforcement learning can be formulated very similarly from a probabilistic viewpoint, and that only (relatively) minor assumptions and choice of approximation tools shift us from one setting to the other. This is a particularly useful insight for our use case of self-driving laboratories, since at the time of writing, Bayesian optimization is a more established approach in practice. Our findings establish a stronger common ground between the seemingly separate Bayesian optimization and reinforcement learning fields, and provide a principled basis for transferring approximation tools between these communities.