summaries: In this paper, the authors propose to replace commonly-used shooting-based methods for action sequence planning in learned latent-space dynamics models by a collocation-based method. They argue that shooting-based methods exhibit problematic behavior especially for sparse-reward and long-horizon tasks, as shooting methods do not allow for planning trajectories which (slightly) violate the learned dynamics. The authors propose a collocation method based on Levenberg-Marquard optimization with a scheduled Lagrange multiplier which outperforms two shooting methods (CEM and gradient-based) on a set of robotic tasks.

summaries: This paper introduces a vision-based motion planning approach using collocation. Many existing approaches to vision-based control rely on computationally expensive planning approaches using shooting to perform model-based control, which is often only useful in simple control tasks. Collocation approaches are effective in settings with difficult path constraints, and thus exploited by this work to dramatically improve model-based reinforcement learning.

summaries: The paper studies the problem of planning in domains with sparse rewards where observations are in the form of images. It focuses on solving this problem using model-based RL with emphasis on better trajectory optimization. The proposed solution uses latent models to extract latent representations of the planning problem that is optimized using the Levenberg-Marquardt algorithm (over a horizon). The experimental results show improvements over a) zeroth-order CEM optimization, b)  PlaNet (Hafner et al., 2019) and c)  gradient-based method that optimizes the objective in Eq. 1.

abstractText: The ability to plan into the future while utilizing only raw high-dimensional observations, such as images, can provide autonomous agents with broad and general capabilities. However, realistic tasks may require handling sparse rewards and performing temporally extended reasoning, and cannot be solved with only myopic, short-sighted planning. Recent work in model-based reinforcement learning (RL) has shown impressive results using heavily shaped reward functions that require only short-horizon reasoning. In this work, we study how techniques trajectory optimization can enable more effective long-horizon reasoning. We draw on the idea of collocation-based planning and adapt it to the image-based setting by leveraging probabilistic latent variable models, resulting in an algorithm that optimizes trajectories over latent variables. Our latent collocation method (LatCo) provides a general and effective approach to longer-horizon visual planning. Empirically, our approach significantly outperforms prior model-based approaches on challenging visual control tasks with sparse rewards and long-term goals.

authors:
Masataro Asai
Alex Fukunaga

title: Classical planning in deep latent space: Bridging the subsymbolicsymbolic boundary

authors:
Somil Bansal
Anayo K Akametalu
Frank J Jiang
Forrest Laine
Claire J Tomlin

title: Learning quadrotor dynamics using neural network for flight control

venue: In 2016 IEEE 55th Conference on Decision and Control (CDC),

authors:
Lars Buesing
Theophane Weber
Sébastien Racanière
S.M. Ali Eslami
Danilo Jimenez Rezende
David P. Reichert
Fabio Viola
Frederic Besse
Karol Gregor
Demis Hassabis
Daan Wierstra

title: Learning and querying fast generative models for reinforcement learning

authors:
Yize Chen
Yuanyuan Shi
Baosen Zhang

title: Optimal control via neural networks: A convex approach

authors:
Kurtland Chua
Roberto Calandra
Rowan McAllister
Sergey Levine

title: Deep reinforcement learning in a handful of trials using probabilistic dynamics models

venue: In Advances in Neural Information Processing Systems,

authors:
Junyoung Chung
Kyle Kastner
Laurent Dinh
Kratarth Goel
Aaron C Courville
Yoshua Bengio

title: A recurrent latent variable model for sequential data

authors:
E. Denton
R. Fergus

title: Stochastic video generation with a learned prior

authors:
Yilun Du
Toru Lin
Igor Mordatch

title: Model based planning with energy based models

venue: arXiv preprint arXiv:1909.06878,

authors:
Frederik Ebert
Chelsea Finn
Alex X Lee
Sergey Levine

title: Self-supervised visual planning with temporal skip connections

authors:
Frederik Ebert
Chelsea Finn
Sudeep Dasari
Annie Xie
Alex Lee
Sergey Levine

title: Visual foresight: Model-based deep reinforcement learning for vision-based robotic control

authors:
Ben Eysenbach
Russ R Salakhutdinov
Sergey Levine

title: Search on the replay buffer: Bridging planning and reinforcement learning

authors:
Marco Fraccaro
Søren Kaae Sønderby
Ulrich Paquet
Ole Winther

title: Sequential neural models with stochastic layers

venue: In Advances in neural information processing systems,

authors:
Danijar Hafner
Timothy Lillicrap
Ian Fischer
Ruben Villegas
David Ha
Honglak Lee
James Davidson

title: Learning latent dynamics for planning

authors:
Danijar Hafner
Timothy Lillicrap
Jimmy Ba
Mohammad Norouzi

title: Dream to control: Learning behaviors by latent imagination

authors:
Charles R Hargraves
Stephen W Paris

title: Direct trajectory optimization using nonlinear programming and collocation

venue: Journal of guidance, control, and dynamics,

authors:
Mrinal Kalakrishnan
Sachin Chitta
Evangelos Theodorou
Peter Pastor
Stefan Schaal

title: Stomp: Stochastic trajectory optimization for motion planning

venue: IEEE international conference on robotics and automation,

title: An introduction to trajectory optimization: How to do your own direct collocation

venue: SIAM Review,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

authors:
Thanard Kurutach
Aviv Tamar
Ge Yang
Stuart J Russell
Pieter Abbeel

title: Learning plannable representations with causal infogan

authors:
Kara Liu
Thanard Kurutach
Christine Tung
Pieter Abbeel
Aviv Tamar

title: Hallucinative topological memory for zero-shot visual planning

venue: arXiv preprint arXiv:2002.12336,

authors:
Igor Mordatch
Emanuel Todorov
Zoran Popović

title: Discovery of complex behaviors through contact-invariant optimization

venue: ACM Transactions on Graphics (TOG),

authors:
Anusha Nagabandi
Kurt Konolige
Sergey Levine
Vikash Kumar

title: Deep dynamics models for learning dexterous manipulation

venue: In Conference on Robot Learning,

authors:
Suraj Nair
Chelsea Finn

title: Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation

venue: arXiv preprint arXiv:1909.05829,

authors:
Soroush Nasiriany
Vitchyr Pong
Steven Lin
Sergey Levine

title: Planning with goal-conditioned policies

authors:
Giambattista Parascandolo
Lars Buesing
Josh Merel
Leonard Hasenclever
John Aslanides
Jessica B Hamrick
Nicolas Heess
Alexander Neitz
Theophane Weber

title: Divide-and-conquer monte carlo tree search for goal-directed planning

authors:
Karl Pertsch
Oleh Rybkin
Frederik Ebert
Chelsea Finn
Dinesh Jayaraman
Sergey Levine

title: Long-horizon visual planning with goal-conditioned hierarchical predictors, 2020

authors:
Nathan Ratliff
Matt Zucker
J Andrew Bagnell
Siddhartha Srinivasa

title: Chomp: Gradient optimization techniques for efficient motion planning

venue: IEEE International Conference on Robotics and Automation,

authors:
Nikolay Savinov
Alexey Dosovitskiy
Vladlen Koltun

title: Semi-parametric topological memory for navigation

venue: arXiv preprint arXiv:1803.00653,

authors:
Buser Say
Ga Wu
Yu Qing Zhou
Scott Sanner

title: Nonlinear hybrid planning with deep net learned transition models and mixed-integer linear programming

venue: In IJCAI,

authors:
Julian Schrittwieser
Ioannis Antonoglou
Thomas Hubert
Karen Simonyan
Laurent Sifre
Simon Schmitt
Arthur Guez
Edward Lockhart
Demis Hassabis
Thore Graepel

title: Mastering atari, go, chess and shogi by planning with a learned model

venue: arXiv preprint arXiv:1911.08265,

authors:
John Schulman
Yan Duan
Jonathan Ho
Alex Lee
Ibrahim Awwal
Henry Bradlow
Jia Pan
Sachin Patil
Ken Goldberg
Pieter Abbeel

title: Motion planning with sequential convex optimization and convex collision checking

venue: The International Journal of Robotics Research,

authors:
Manuel Watter
Jost Springenberg
Joschka Boedecker
Martin Riedmiller

title: Embed to control: A locally linear latent dynamics model for control from raw images

authors:
Tianhe Yu
Deirdre Quillen
Zhanpeng He
Ryan Julian
Karol Hausman
Chelsea Finn
Sergey Levine

title: Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning

authors:
Marvin Zhang
Sharad Vikram
Laura Smith
Pieter Abbeel
Matthew Johnson
Sergey Levine

title: Solar: Deep structured representations for model-based reinforcement learning

venue: In International Conference on Machine Learning,

heading: 1 INTRODUCTION

text:
In order for autonomous agents to perform complex tasks in open-world settings, they must be able to process high-dimensional sensory inputs, such as images, and reason over long horizons about the potential effects of their actions. Recent work in model-based reinforcement learning (RL) has shown impressive results in autonomous skill acquisition directly from image inputs, demonstrating benefits such as learning efficiency and improved generalization. While these advancements have been largely fueled by improvements on the modeling side – from better uncertainty estimates and incorporation of temporal information (Ebert et al., 2017), to the explicit learning of latent representation spaces (Hafner et al., 2019) – they leave much room for improvement on the planning and optimization side. Most of the current best-performing deep learning approaches for vision-based planning use only gradient-free action sampling as the underlying optimizer, and are typically applied to settings where a dense and well-shaped reward signal is available. In this work, we argue that more powerful planners are necessary for longer-horizon reasoning.
With this goal of long-horizon planning, we aim to extend the myopic planning behavior of existing visual planning methods. Whether it’s to avoid local minima due to short-sightedness, or to reason further into the future in order to solve multi-step or sparse-reward tasks, this ability to perform long-horizon planning is critical. Many of the current state-of-the-art visual planning approaches use gradient-free sampling-based optimization methods such as shooting (Ebert et al., 2018; Nagabandi
et al., 2020); as shown in Figure 6, these approaches can get stuck when they must reason further into the future. Here, the curse of dimensionality in conjunction with a lack of shaped reward signal can prevent greedy planning methods from succeeding. Instead, we look to the gradient-based optimization approach of collocation, where crucially, dynamics satisfaction constraints are optimized jointly with the rewards. As shown in Figure 1, this ability to violate dynamics constraints and imagine even impossible trajectories enables collocation to explore much more effectively and take shortcuts through the optimization landscape in order to learn about high-reward regions before figuring out how to get to those regions. In contrast to shooting approaches, or even to gradient-based approaches such as backpropagating rewards through model predictions, this ability to violate dynamics greatly helps to prevent getting stuck in local minima while still avoiding the need for special dense or shaped reward signals.
Collocation, as introduced above, can provide many benefits over other optimization techniques, but it has thus far been demonstrated (Ratliff et al., 2009; Mordatch et al., 2012; Schulman et al., 2014) mostly in conjunction with known dynamics models and when performing optimization over states. In this work, we are interested in autonomous skill acquisition directly from image inputs, where both the underlying states as well as the underlying dynamics are unknown. Naı̈vely applying collocation to this visual planning setting would lead to intractable optimization, due to the high-dimensional as well as partially observed nature of images, in addition to the fact that only a thin manifold in the pixel space constitutes valid images. Instead, we draw from the representation learning literature and leverage latent dynamics models, which learn a latent representation of the observations that is not only Markovian but also compact and lends itself well to planning. In this learned latent space, we propose to perform collocation over states and actions with the joint objective of maximizing rewards as well as minimizing dynamics violations.
Bridging control theory literature with modern deep-RL techniques for learning from images, we propose in this work to perform collocation in a learned latent space in order to enable effective long-horizon planning. The main contribution of this work is an algorithm for latent-space collocation (LatCo), which is an efficient model-based RL approach for solving challenging planning tasks that works directly from image observations by leveraging collocation in a learned latent space. To the best of our knowledge, our paper is the first to scale collocation to visual observations, thus enabling longer-horizon reasoning in model-based RL. In our experiments, we analyze various aspects of our algorithm, and we demonstrate our approach significantly outperforming prior model-based approaches on visual control tasks that require longer-horizon planning with sparse rewards.

heading: 2 RELATED WORK

text:
Model-based reinforcement learning. Recent work has scaled model-based reinforcement learning to complex systems leveraging powerful neural network dynamics models (Chua et al., 2018; Nagabandi et al., 2020), while showing significant data efficiency improvements over model-free agents. Further, these neural network models can be scaled to high-dimensional image observations using convolutional architectures (Ebert et al., 2018; Hafner et al., 2019). However, despite these successes in building better predictive models, planning with these black-box neural network dynamics remains a challenge. While this prior work used simple trajectory optimization techniques like derivative-free shooting, we propose to leverage the more powerful collocation methods. Other work explored more complex approaches based on mixed-integer linear programming (Say et al., 2017) or gradient descent with input-convex neural networks (Chen et al., 2019), but it is unclear whether these approaches scale to visual observations.
Latent Planning. Other works have considered different optimization methods such as iterative Linear-Quadratic Regulator (iLQR) (Watter et al., 2015; Zhang et al., 2019). However, these approaches require specialized locally-linear predictive models, and still rely on shooting and local search in the space of actions, which is prone to local minima. Instead, our collocation approach can be used with any latent state model, and is able to optimize in the state-space, which we show often enables us escape local minima and plan better trajectories. Another line of work relied on graph-based optimization (Kurutach et al., 2018; Savinov et al., 2018; Eysenbach et al., 2019; Liu et al., 2020) tree search (Schrittwieser et al., 2019; Parascandolo et al., 2020), or other symbolic planners (Asai & Fukunaga, 2018), while we use continuous optimization, which is more suitable for continuous control. Recent work has designed hierarchical planning methods that plan over extended
periods of time by considering intermediate subgoals (Nair & Finn, 2019; Nasiriany et al., 2019; Pertsch et al., 2020). These approaches are orthogonal to our trajectory optimization method, but they present many possible synergies as collocation can be naturally used in hierarchical approaches. Many of these methods (Buesing et al., 2018; Hafner et al., 2019; 2020) used latent state-space models for improved prediction quality and reduced computational requirements. Our proposed method leverages this latent state-space design to construct an effective trajectory optimization method with collocation, and we design our method to be model-agnostic, such that it can benefit from improved latent variable models in the future.
Collocation-based planning. Collocation is a powerful optimal control technique for trajectory optimization (Hargraves & Paris, 1987; Witkin & Kass, 1988) that optimizes a sequence of states for the sum of expected reward, while eventually enforcing the constrained that the optimized trajectory conform to a dynamics model (also see Kelly (2017) for a recent tutorial). Prior work in optimal control has explored many versions of collocation for complex motion planning tasks, including Hamiltonian optimization (Ratliff et al., 2009), contact-invariant optimization (Mordatch et al., 2012), sequential convex programming (Schulman et al., 2014), as well as stochastic trajectory optimization (Kalakrishnan et al., 2011). We note that in this paper we will use the terms ”trajectory optimization” and ”planning” synonymously referring to this type of approaches. These works have demonstrated good results in controlling complex simulated characters, such as humanoid robots, contact-heavy tasks, and tasks with complex constraints. Our work is most similar to that of Schulman et al. (2014), however, all this prior work assumed availability of a ground truth model of the environment. Some recent works have attempted using collocation with learned neural network dynamics models (Bansal et al., 2016; Du et al., 2019), but it only considered simple low-dimensional dynamics. In this work, we address how to scale up collocation methods to high-dimensional image observations, where direct optimization over images is intractable. We propose to do this by utilizing a learned latent space representation.

heading: 3 BACKGROUND

text: Our method combines collocation methods for trajectory optimization with latent variable models, in order to make it possible to optimize paths over compact and Markovian latent states. We separately review collocation methods and latent state models in this section.

heading: 3.1 TRAJECTORY OPTIMIZATION WITH COLLOCATION

text:
Given a dynamics model st+1 = f(st, at) that predicts the next state given the previous state and action, a reward function r(st), and the current state s1, the problem of trajectory optimization is to select actions that maximize the total reward:
max a1:T−1 ∑ t r(st) = [r(f(s1, a1)) + r(f(f(s1, a1), a2)) + r(f(f(f(s1, a1), a2), a3) + . . . ] . (1)
Shooting methods optimize this objective directly with respect to the actions. However, this is known to be poorly conditioned due to recursive application of the dynamics function, which results in vanishing or exploding gradients. Instead, we can leverage the structure of the problem to construct an objective that only has pairwise dependencies between temporally adjacent states, and no recursive application of the model. To this end, collocation methods formulate the trajectory optimization problem in Equation 1 as a constrained optimization problem, optimizing over sequences of actions and states, while ensuring that the constraint imposed by the dynamics model f is satisfied:
max s2:T ,a1:T−1 ∑ t r(st) s.t. st+1 = f(st, at). (2)
This constrained optimization problem, under some regularity conditions, can be solved with a primal dual approach, which can be formulated as the following saddle point problem:
min λ max s2:T ,a1:T−1 ∑ t r(st)− λ||st+1 − f(st, at)||2. (3)
In practice, we can address this problem with numerical optimization, taking alternating maximization and minimization steps. This collocation approach is better conditioned, as it is able to exploit the
Markovian property of the state-space and remove the recursive application of the model. The dynamics constraint only needs to be evaluated on each pair of temporally adjacent states (see Fig 2, right). Further, the dynamics constraint can be relaxed in the beginning of the optimization, leading the method to rapidly discover high-reward state space regions (potentially violating the dynamics), and then gradually modifying the trajectory to be more dynamically consistent, as illustrated in Figure 1. This is in contrast to shooting-based approaches, which suffer from severe local optima in long horizon tasks, since the algorithm must simultaneously drive the states toward the high-reward regions and discover the actions that will get them there. The ability to disentangle these two stages by first finding the high-reward region and only then optimizing for the actions that achieve that reward allows collocation methods to solve more complex and temporally extended tasks while suffering less from local optima (Ratliff et al., 2009; Mordatch et al., 2012; Schulman et al., 2014).

heading: 3.2 LATENT VARIABLE MODELS

text:
The design of predictive models for high-dimensional visual observations is challenging. Recent work has proposed learning latent state-space models that represent the observations in a compact latent space. Specifically, this work learns a latent space zt, a latent dynamics model pθ(zt+1|zt, at), and decoding models that can be used to decode the observation pθ(ot|zt) and the reward r(zt) from the latent variable (see Fig. 2, left). This approach is powerful due to the use of high-capacity neural network latent dynamics models, and is computationally efficient as the latent space is compact. Importantly, the Markov property of the latent variable can be enforced, allowing a convenient interpretation of the latent state-space as the belief state (Hafner et al., 2019; Buesing et al., 2018).
A probabilistic latent variable model can then be trained, by maximizing the likelihood of the observed rewards r1:T and images o1:T . While maximizing the exact likelihood is often intractable, we can construct and optimize a variational lower bound on it using a variational distribution qθ(zt+1|ot+1, at, zt) (Chung et al., 2015; Fraccaro et al., 2016):
ln pθ(o2:T , r1:T |o1, a1:T ) ≥ LELBO(o1:T , a1:T , r1:T ) = Eqθ(z1:T |o1:T ,a1:T ,z0) ∑ t [ln pθ(ot+1|zt+1)− KL (qθ(zt+1|ot+1, at, zt) || pθ(zt+1|zt, at))] . (4)

heading: 4 LATENT COLLOCATION (LATCO)

text:
Our aim is to design a collocation method that can be used to produce near-optimal trajectories from raw image observations. A naı̈ve approach would learn a dynamics model over images, and directly optimize an image sequence using Eq (2). However, such as method would be impractical for several
reasons. First, the optimization over images would make the problem more difficult due to the high dimensionality of the images and the fact that valid images lie on a thin manifold. Second, images typically do not constitute a Markovian state space, violating the assumptions of the method. We propose to instead learn a Markovian and compact state space by means of a latent variable model, and then use this learned state space for collocation.
Latent state models. To make visual collocation scalable, we leverage the latent state-space dynamics model described in Section 3.2. Following Hafner et al. (2019); Denton & Fergus (2018), we implement this model with convolutional neural networks for the encoder and decoder, and a recurrent neural network for the transition dynamics. The latent state of this model includes a stochastic component with a conditional Gaussian transition function, and the hidden state of the recurrent neural network with deterministic transitions. The model is shown in Fig. 2 (left).
Latent space collocation. We adapt the collocation procedure described in Section 3.1 into the probabilistic latent state setting. Instead of optimizing a sequence of observations, we optimize a sequence of latent variables z2:T and actions a1:T . Since the latent state space model predicts the reward directly from the latent state zt using an approximation r(zt), we never need to decode images during the optimization procedure, which makes it memory-efficient. We adapt the collocation algorithm to use probabilistic dynamics by enforcing that the next latent state is the mean of the distribution predicted by the model: zt+1 = Epθ(zt+1|zt,at) [zt+1], where the expectation is given simply as the mean of the Gaussian zt+1 = µ(zt, at). Other approaches are possible, such as maximizing the likelihood of the latent state, but these require introduction of additional balance hyperparameters, while our constrained optimization approach automatically tunes this balance. We visualize the collocation procedure in Fig 2 and provide the detailed algorithm as Algorithm 1.
Algorithm 1 Latent Collocation (LatCo) 1: Start with any available data D 2: while not converged do 3: for each environment step t = 1 . . . Ttot with step Tcache do 4: Infer latent state: zt ∼ q(zt|ot) 5: Define the Lagrangian:
L(zt+1:t+H , at:t+H , λ) = ∑ t [ r(zt)− λdynt (||zt+1 − µ(zt, at)||2 − )− λactt max(0, |at| − am)2 ] 6: for each optimization step k = 1 . . .K do 7: Update plan: zt+1:t+H , at:t+H := zt+1:t+H , at:t+H + ∇̃L . Eq (5) 8: Update dual variables: λt:t+H := UPDATE(L, λt:t+H) . Eq (6) 9: Execute at:t+Tcache in environment: ot:t+Tcache , rt:t+Tcache ∼ penv
10: Add episode to replay buffer: D := D ∪ (o1:Ttot , a1:Ttot , r1:Ttot) 11: for training iteration i = 1 . . . It do 12: Sample minibatch from replay buffer: (o1:T , a1:T , r1:T )1:b ∼ D 13: Train dynamics model: θ := θ + α∇LELBO(o1:T , a1:T , r1:T )1:b . Eq (4)

heading: 5 OPTIMIZATION FOR LATENT COLLOCATION

text:
The latent collocation framework described in the previous section can be used with any optimization procedure, such as gradient descent or Adam (Kingma & Ba, 2015). However, using appropriate optimization methods is important when time constraints are a factor. We found that the choice of the optimizer for both the latent states and the Lagrange multipliers has a large influence on runtime performance. We detail our specific implementation below.
Levenberg-Marquardt optimization. We use the Levenberg-Marquardt optimizer for the states and actions, which pre-conditions the gradient direction with the matrix (JTJ)−1, where J is the Jacobian of the objective with respect to the states and actions. This preconditioner approximates the Hessian inverse, significantly improving convergence speed:
∇̃ = (JTJ + λI)−1JTρ. (5)
The Levenberg-Marquardt optimizer has cubic complexity in the number of optimized dimensions. However, by noting that the Jacobian of the problem has a block-tridiagonal structure, we can implement a more efficient optimizer that scales linearly with the planning horizon (Mordatch et al., 2012). This efficient version of the optimizer converges 10-100 times faster than naı̈ve gradient descent on wall clock time in our experiments.
The Levenberg-Marquardt algorithm optimizes the sum of squares ∑ i ρ 2 i , defined in terms of residuals ρ. Any bounded objective can be expressed in terms of residuals by a combination shifting and square root operations. For the dynamics objective, we use the zt+1 − µ(zt, at) differences as residuals directly, with one residual per state dimension. For the action objective, we use max(0, |at|−am). For the reward objective, we form residuals by passing the negative reward through a softplus operation: ρ = ln(1 + e−r).
Constrained optimization. The Lagrange multipliers control the balance between the strength of the dynamics constraint and the reward. For fast optimization, it is crucial that the Lagrange multipliers be updated fast when the dynamics constraint is not yet optimized, but the update magnitude should decrease over time to ensure smooth convergence as the dynamic violation converges to zero. We therefore design an update rule that scales with the relative difference between the current dynamics constraint ||zt+1 − µ(zt, at)||2 and the target constraint , such that the update is large when the constraint violation is large and vice versa. We found the following rule to work well in our experiments:
λdyn += 0.1 log
( ||zt+1 − µ(zt, at)||2
+ η ) λdyn, (6)
where η = 0.01 is used to ensure numerical stability. We found that using a small non-zero is beneficial for the optimization and ensures fast convergence, as the exact constraint might be hard to reach. We additionally constrain the planned actions to be within the environment range using the same constrained optimization procedure.

heading: 6 EXPERIMENTS

text: Our evaluation aims to test the effectiveness of our method on performing long-horizon MBRL from images by answering the following questions: (1) Does LatCo produce effective and feasible plans from visual inputs? (2) Can LatCo solve complex simulated robotic tasks with sparse rewards? (3) How does LatCo compare to trajectory optimization methods used in state-of-the-art visual planning algorithms? (4) What is the effect of LatCo allowing dynamics violations during optimization?

heading: 6.1 EXPERIMENTAL SETUP

text:
We evaluate our method on three visual tasks (see Fig. 3): Obstacle, a navigation environment with an obstacle, Sawyer Reach, a robotic reaching task, Sawyer Push, a robotic pushing task. The first
of these tasks is a custom environment, while the last two are from the Meta-World benchmark (Yu et al., 2020). All tasks only provide the agent with image inputs and no access to underlying state of the simulator. The Obstacle environment consists of an image-based 2D navigation task where the agent must reach a goal on the other side of a wall. The reward defined is the negative distance to the goal. The agent plans over a planning horizon H = 30 steps, executes every action Tcache = 30, Ttot = 30, and results are averaged over 100 runs. In the Sawyer Reach task, a 7-DoF Sawyer arm needs to reach the goal. For this task, H = 30, Tcache = 30, Ttot = 150. In the Sawyer Push task, a 7-DoF Sawyer arm needs to push an object (puck) on the table to the goal. For this task, H = 40, Tcache = 20, Ttot = 120. In both Sawyer tasks, we use a sparse reward of 1 for reaching the goal and 0 otherwise. In particular, for the push task, there is no special reward shaping to encourage the arm to go toward the puck. The arm is initialized to its default position, and results are averaged over 20 runs. We train all methods online according to Algorithm 1. The hyperparameters used in the experiments are detailed in Appendix A.

heading: 6.2 DOES LATCO PRODUCE EFFECTIVE PLANS ON SPARSE-REWARD VISUAL CONTROL TASKS?

text: In our first set of experiments, we evaluate how well LatCo can solve image-based control tasks from sparse rewards, where effectively reasoning into the future is critical for success. As shown in Figure 3, LatCo is indeed able to perform collocation in a learned latent space in order to successfully produce plans that solve these tasks, with effective longerhorizon reasoning. In these results, we see LatCo (top) avoiding local minima that would arise from myopic planners, (middle) extending to higherdimensional control tasks while still operating only from image inputs, and (bottom) reasoning into the future and succeeding even when the sparse task reward provides no direct incentive to go toward the object.

heading: 6.3 HOW DOES LATCO COMPARE TO ALTERNATIVE SHOOTING-BASED APPROACHES?

text:
To specifically ascertain the benefits of collocation, we must determine whether the benefits of LatCo stem from gradient-based optimization, from optimizing over states, or both. Therefore, we include a prior method based on zeroth-order CEM optimization, PlaNet (Hafner et al.,
2019), which is representative of a broader class of shooting-based methods (Ebert et al., 2018; Nagabandi et al., 2020), as well as a gradient-based method that optimizes actions directly using the objective in Eq. 1, which we denote GD actions. To provide a fair comparison that isolates the effects of different planning methods, we use the same dynamics model architecture for all agents.
From Table 1, we observe that LatCo exhibits superior performance to shooting-based methods on all three tasks. On the Sawyer Reach task, shooting fails to find a successful trajectory most of the time due to a lack of local reward signal. LatCo, on the other hand, solves this task consistently. The Sawyer Push task is challenging for all of the methods, as the dynamics are more complex and requires accurate long-term predictions to correctly optimize the sparse signal. Specifically,
there is no partial reward for reaching down to the object, so the planner has to look ahead into the high reward state and reason backwards that it needs to approach the object. As shown in Fig. 4, shooting-based methods altogether fail to solve this task, not reaching for the object in any run. LatCo outperforms them by a considerable margin, achieving a 29% success rate. It is worth noting that even in the unsuccessful runs, LatCo plans often display goal-reaching behaviors, as can be seen in Fig. 4. However, significant room for improvement for MBRL still remains on this task. In addition, to evaluate the different methods in a more controlled setup, we test them on the same model trained with offline data in Table 2 and App. C

heading: 6.4 THE EFFECT OF ALLOWING DYNAMICS VIOLATIONS

text: In this section, we analyze the ability of our method to temporarily violate dynamics in order to effectively plan for long-term reward. We show the dynamics violation costs and the predicted rewards associated with the planned trajectories over the course of optimization quantitatively in Fig. 5 and qualitatively in Figs. 1, 6. Since the dynamics constraint is only gradually enforced with the increase of the Lagrange multipliers, the first few steps of optimization allow for dynamics violations in order to focus on discovering the high-reward region (corresponding to steps 0 to 20 in Fig. 5). In the later steps of the optimization, the reward is relaxed and the constraints are optimized until the trajectory becomes feasible.

heading: 7 CONCLUSION

text:
We presented LatCo, an algorithm for latent collocation that improves performance of visual-model based reinforcement learning agents. In contrast to the commonly used shooting-based methods, LatCo performs powerful trajectory optimization and is able to plan for tasks where prior work fails, such as long-horizon and sparse reward tasks. By improving the planning capabilities of visual model-based reinforcement learning agents and removing the need for reward shaping, LatCo enables these agents to scale to complex tasks more easily and with less manual instrumentation.
Future work will examine improvements to the collocation procedure, such as Hamiltonian (Ratliff et al., 2009) or stochastic (Kalakrishnan et al., 2011) optimization. Further, collocation can be applied with variety of latent variable models, including specialized models with more structure. Finally, we believe the ability to perform reinforcement learning in state space as opposed to action space opens up many new avenues for algorithmic development, such as in imitation learning or hierarchical reinforcement learning.

heading: A MODEL ARCHITECTURE AND TRAINING DETAILS

text: We use the latent dynamics and reward models from PlaNet (Hafner et al., 2019) with default hyperparameters. We set the image size to 64x64 and action repeat to 1 for both pointmass and Meta-World models. For every N = 1 episode collected, we train for It = 15 iterations. The pointmass models are trained with episode length 30 and βKL = 0.1, while the Meta-World models are trained with episode length 150 and βKL = 1. All models are trained on a single high-end GPU.

heading: B PLANNING DETAILS

text:
CEM. we optimize for 100 iterations. In each iteration, 10000 action sequences are sampled and the distribution is refit to the 100 best samples. For this baseline, we have manually tuned the batch size and number of iterations, and we report the best results.
GD. we optimize for 500 iterations using the Adam optimizer (Kingma & Ba, 2015) (which
is a modified version of momentum gradient descent) with learning rate 0.05. We use dual descent to penalize infeasible action predictions. The Lagrange multipliers are updated every 5 optimization steps. For this baseline, we have manually tuned the learning rate and tried several first-order optimizers, and we report the best results.
LatCo. we optimize for 200 iterations using the Levenberg-Marquardt optimizer with damping 10−3. The damping parameter controls the trust region, with smaller or zero damping speeding up convergence, but potentially leading to numerical instability or divergence. The Lagrange multipliers are updated every step using the rule from Section 5, with = 10−5 and η = 0.01. The threshold directly controls the magnitude of the final dynamics and action violations. In general, we found this parameter to be important for good performance, as a large threshold may cause infeasible plans, while low threshold would make the initial relaxation of the dynamics constraint less effective. However, we observed that a single threshold of 10−5 works for all of our environments. A larger η makes the optimization more aggressive but less stable, and a smaller η diminishes the effect of multiplier updates. We initilalize λdyn0 = 1, λ act 0 = 1.
These planning hyperparameters remain fixed across the experiments as we observe that reward optimization converges in all cases. Planning a 30-step trajectory takes 12, 14, and 14 seconds for CEM, GD, and LatCo respectively on a single RTX 2080Ti graphics card. We trained the online models for 24 hours each.

heading: C ADDITIONAL EXPERIMENTAL RESULTS

text:
Offline training. In Section 6.3, we have evaluated our algorithm in the online training regime, where each agent needs to collect the entire dataset from scratch by executing the policy in the environment. To further test whether collocation is beneficial for planning in a controlled setup, we have performed offline training experiments, where a single model trained with offline data is tested with different planning algorithms. This experiment further shows the applicability of our method to offline setting without additional data collection. We collect the offline data by training an oracle agent that observes dense reward, avoiding the exploration issue. We precollect 200K steps of this data for all tasks. For the Sawyer Push task, we collect an additional 200K environments steps of online interaction data according to Algorithm 1 as none of the methods were able to solve the task with offline data only. We use the same hyperparameters as in Section 6.3 otherwise.
Table 2 shows the results of this experiment. We see that collocation outperforms shooting methods in this controlled setup, showing that it
is better able to exploit a dynamics model trained offline.
We visualize the optimization procedure on the obstacle task in Fig 6.
We visualize the reward predictor output on the Pushing task in Fig 8.
We visualize the additional curves of the optimization in Fig 7.
We further compare the predicted and achieved reward by all methods in the offline setup in Table 3. The predicted reward is often higher than the achieved reward, indicating some degree of model exploitation. However, we see that the predicted reward reflects the general trends in the achieved reward and trajectories with high predicted reward also achieve high reward for our method.


summaries: The work propose a theory suggesting that the cold posterior phenomena arises solely due the the curated nature of image benchmarks. A generative model is proposed where multiple annotators label datapoints, and only unanimously labeled datapoints are accepted into a dataset. This theory is studied under a toy-problem using VI and a relabelled version of the CIFAR-10 test set with SGLD.

summaries: This paper addresses the perplexing issue of cold posterior having better predictive performance than the ideal Bayesian posterior in Bayesian deep learning (Wenzel et al., 2020), and offers a possible explanation in terms of a mis-specified likelihood function that deviates from the true generative process of the data. By considering the data curation process and augmenting the likelihood model accordingly, the effect of cold posterior is shown to diminish significantly, and the ideal posterior is again optimal. Empirical results on both a toy problem and image classification support the theory.

summaries: The authors propose the idea that cold posteriors in Bayesian neural networks could be caused by the likelihood instead of the prior. They argue theoretically that the curation process of popular benchmark data sets would lead to a different weighting of the likelihood in the posterior. They show in some experiments that the cold posterior effect can be reduced when accounting for this.

abstractText: To get Bayesian neural networks to perform comparably to standard neural networks it is usually necessary to artificially reduce uncertainty using a “tempered” or “cold” posterior. This is extremely concerning: if the generative model is accurate, Bayesian inference/decision theory is optimal, and any artificial changes to the posterior should harm performance. While this suggests that the prior may be at fault, here we argue that in fact, BNNs for image classification use the wrong likelihood. In particular, standard image benchmark datasets such as CIFAR-10 are carefully curated. We develop a generative model describing curation which gives a principled Bayesian account of cold posteriors, because the likelihood under this new generative model closely matches the tempered likelihoods used in past work.

authors:
Ben Adlam
Jasper Snoek
Samuel L Smith

title: Cold posteriors and aleatoric uncertainty

venue: arXiv preprint arXiv:2008.00029,

title: A statistical theory of semi-supervised learning

venue: arXiv preprint,

authors:
Arsenii Ashukha
Alexander Lyzhov
Dmitry Molchanov
Dmitry Vetrov

title: Pitfalls of in-domain uncertainty estimation and ensembling in deep learning

venue: arXiv preprint arXiv:2002.06470,

authors:
Juhan Bae
Guodong Zhang
Roger Grosse

title: Eigenvalue corrected noisy natural gradient

venue: arXiv preprint arXiv:1811.12565,

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: IEEE conference on computer vision and pattern recognition,

title: The safe bayesian

venue: In International Conference on Algorithmic Learning Theory, pp. 169–183

authors:
Peter Grünwald
Thijs Van Ommen

title: Inconsistency of bayesian inference for misspecified linear models, and a proposal for repairing it

venue: Bayesian Analysis,

authors:
James Hensman
Alexander Matthews
Zoubin Ghahramani

title: Scalable variational gaussian process classification

title: Probability theory: The logic of science

venue: Cambridge university press,

title: Learning multiple layers of features from tiny images

venue: Technical report,

authors:
Alexander G de G Matthews
James Hensman
Richard Turner
Zoubin Ghahramani

title: On sparse variational methods and the kullback-leibler divergence between stochastic processes

venue: Journal of Machine Learning Research,

authors:
Sebastian W Ober
Laurence Aitchison

title: Global inducing point variational posteriors for bayesian neural networks and deep gaussian processes

venue: arXiv preprint arXiv:2005.08140,

authors:
Kazuki Osawa
Siddharth Swaroop
Mohammad Emtiyaz E Khan
Anirudh Jain
Runa Eschenhagen
Richard E Turner
Rio Yokota

title: Practical deep learning with bayesian principles

authors:
Joshua C Peterson
Ruairidh M Battleday
Thomas L Griffiths
Olga Russakovsky

title: Human uncertainty makes classification more robust

venue: In Proceedings of the IEEE International Conference on Computer Vision,

title: L2 regularization versus batch and weight normalization

venue: arXiv preprint arXiv:1706.05350,

authors:
Max Welling
Yee W Teh

title: Bayesian learning via stochastic gradient langevin dynamics

venue: In Proceedings of the 28th international conference on machine learning

authors:
Florian Wenzel
Kevin Roth
Bastiaan S Veeling
Jakub Świątkowski
Linh Tran
Stephan Mandt
Jasper Snoek
Tim Salimans
Rodolphe Jenatton
Sebastian Nowozin

title: How good is the bayes posterior in deep neural networks really

venue: arXiv preprint arXiv:2002.02405,

authors:
Kyle W Willett
Chris J Lintott
Steven P Bamford
Karen L Masters
Brooke D Simmons
Kevin RV Casteels
Edward M Edmondson
Lucy F Fortson
Sugata Kaviraj
William C Keel

title: Galaxy zoo 2: detailed morphological classifications for 304 122 galaxies from the sloan digital sky survey

venue: Monthly Notices of the Royal Astronomical Society,

authors:
Christopher KI Williams
Carl Edward Rasmussen

title: Gaussian processes for machine learning

authors:
Guodong Zhang
Shengyang Sun
David Duvenaud
Roger Grosse

title: Noisy natural gradient as variational inference

venue: arXiv preprint arXiv:1712.02390,

authors:
Ruqi Zhang
Chunyuan Li
Jianyi Zhang
Changyou Chen
Andrew Gordon Wilson

title: Cyclical stochastic gradient mcmc for bayesian deep learning

heading: 1 INTRODUCTION

text:
Recent work has highlighted that Bayesian neural networks (BNNs) typically have better predictive performance when we “sharpen” the posterior (Wenzel et al., 2020). In stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011), this can be achieved by multiplying the log-posterior by 1/T , where the “temperature”, T is smaller than 1 (Wenzel et al., 2020). Broadly the same effect can be achieved in variational inference by “tempering”, i.e. downweighting the KL term. As noted in Wenzel et al. (2020), this approach has been used in many recent papers to obtain good performance, albeit without always emphasising the importance of this factor (Zhang et al., 2017; Bae et al., 2018; Osawa et al., 2019; Ashukha et al., 2020).
These results are puzzling if we take the usual Bayesian viewpoint, which says that the Bayesian posterior, used with the right prior, and in combination with Bayes decision theory should give optimal performance (Jaynes, 2003). Thus, these results may suggest we are using the wrong prior. While new priors have been suggested (e.g. Ober & Aitchison, 2020), they give only minor improvements in performance — certainly nothing like enough to close the gap to carefully trained non-Bayesian networks. In contrast, tempered posteriors directly give performance comparable to a carefully trained finite network.
The failure to develop an effective prior suggests that we should consider alternative explanations for the effectiveness of tempering. Here, we consider the possibility that it is predominantly (but not entirely) the likelihood, and not the prior that is at fault. In particular, we note that standard image benchmark datasets such as ImageNet and CIFAR-10 are carefully curated, and that it is important to consider this curation as part of our generative model. We develop a simplified generative model describing dataset curation which assumes that a datapoint is included in the dataset only if there is unanimous agreement on the class amongst multiple labellers. This model naturally multiplies the effect of each datapoint, and hence gives posteriors that closely match tempered or cold posteriors. We show that toy data drawn from our generative model of curation can give rise to optimal temperatures being smaller than 1. Our model predicts that cold posteriors will not be helpful when the original underlying labels from all labellers are available. While these are not available for standard datasets such as CIFAR-10, we found a good proxy: the CIFAR-10H dataset (Peterson et al., 2019), in which ∼ 50 humans annotators labelled the CIFAR-10 test-set (we use these as our training set, and use the standard CIFAR-10 training set for test-data). As expected, we find strong cold-posterior effects
when using the original single-label, which are almost entirely eliminated when using the 50 labels from CIFAR-10H. In addition, curation implies that each label is almost certain to be correct, which is one way to understand the statistical patterns exploited by cold posteriors. As such, if we destroy this pattern by adding noise to the labels, the cold posterior effect should disappear. We confirmed that with increasing label noise, the cold posterior effect disappears and eventually reverses (giving better performance at temperatures close to 1).

heading: 2 BACKGROUND: COLD AND TEMPERED POSTERIORS

text:
Tempered (e.g. Zhang et al., 2017) and cold (Wenzel et al., 2020) posteriors differ slightly in how they apply the temperature parameter. For cold posteriors, we scale the whole posterior, whereas tempering is a method typically applied in variational inference, and corresponds to scaling the likelihood but not the prior,
log Pcold (θ|X,Y ) = 1T log P (X,Y |θ) + 1 T log P (θ) + const (1) log Ptempered (θ|X,Y ) = 1λ log P (X,Y |θ) + log P (θ) + const . (2) While cold posteriors are typically used in SGLD, tempered posteriors are usually targeted by variational methods. In particular, variational methods apply temperature scaling to the KL-divergence between the approximate posterior, Q (θ) and prior,
L = EQ(θ) [log P (X,Y |θ)]− λDKL (Q (θ) ||P (θ)) . (3) Note that the only difference between cold and tempered posteriors is whether we scale the prior, and if we have Gaussian priors over the parameters (the usual case in Bayesian neural networks), this scaling can be absorbed into the prior variance,
1 T log Pcold (θ) = − 1 2Tσ2cold ∑ i θ2i + const = − 12σ2tempered ∑ i θ2i + const = log Pcold (θ) . (4)
in which case, σ2cold = σ 2 tempered/T , so the tempered posteriors we discuss are equivalent to cold posteriors with rescaled prior variances.

heading: 3 METHODS: A GENERATIVE MODEL FOR CURATED DATASETS

text:
Standard image datasets such as CIFAR-10 and ImageNet are carefully curated to include only unambiguous examples of each class. For instance, in CIFAR-10, student labellers were paid per hour (rather than per image), were instructed that “It’s worse to include one that shouldn’t be included than to exclude one”, and then Krizhevsky (2009) “personally verified every label submitted by the labellers”. For ImageNet, Deng et al. (2009) required the consensus of a number of Amazon Mechanical Turk labellers before including an image in the dataset.
To understand the statistical patterns that might emerge in these curated datasets, we consider a highly simplified generative model of consensus-formation. In particular, we draw a random image X from some underlying distribution over images, P (X), and ask S humans to assign a label, {Ys}Ss=1 (e.g. using Mechanical Turk). We force every labeller to label every image and if the image is ambiguous they are instructed to give a random label. If all the labellers agree, Y1=Y2= · · · =YS , consensus is reached and we include the datapoint in the dataset. If any of the labellers disagree consensus is not reached, and we exclude the datapoint (Fig. 1), Formally, the observed random variable, Y , is taken to be the usual label if consensus was reached and None if consensus was not reached (Fig. 2B),
Y |{Ys}Ss=1 = { Y1 if Y1=Y2= · · · =YS None otherwise
(5)
Taking the human labels, Ys, to come from the set Y , so Ys ∈ Y , the consensus label, Y , could be any of the underlying labels in Y , or None if no consensus is reached, so Y ∈ Y ∪ {None}. When consensus was reached, the likelihood is,
P (Y =y|X, θ) = P ( {Ys=y}Ss=1|X, θ ) = S∏ s=1 P (Ys=y|X, θ) = P (Ys=y|X, θ)S (6)
where we have assumed labellers are IID. This would appear to directly give an account of tempering, as we have taken the single-labeller likelihood to the power S, which is equivalent to setting λ = 1/S. However, to see how the full generative model functions we need to go on to consider the case in which consensus was not reached,
P (Y =None|X, θ) = 1− ∑ y∈Y P (Y =y|X, θ) = 1− ∑ y∈Y P (Ys=y|X, θ)S . (7)
To understand the impact of our model of consensus formation, note that the probability of a particular class-label can be separated into two terms, a probability of consensus, and the probability of a particular class given that consensus was reached,
P (Y =y|X, θ) = P (Y =y|Y 6=None, X, θ) P (Y 6=None|X, θ) . (8)
Remarkably, knowing there was consensus gives us information about the weights, even in the absence of the class label,
P (Y 6=None|X, θ) = ∑ y∈Y P (Ys=y|X, θ)S , (9)
in essence, telling us that the output probability was close to 1 for one of the class labels, without telling us which one. Schematically (Fig. 3), we see that the datapoints with a consensus class-label, “cat” or “dog”, lie far from the decision boundary where the class is unambiguous, and consensus is easily reached. In contrast, in regions close to the decision boundary the inputs are ambiguous, which tends to produce disagreement in the labellers, leading to noconsensus. Thus, the existence of one or more consensus points in a region implies that decision boundaries do not go through that region, giving us information about the decision boundary location, even if the label is not known. Concurrent work has shown that this likelihood can be used to explain classical semi-supervised likelihoods (Aitchison, 2020), so this term really does give information about the neural network parameters. Finally, the label probability, conditioned on consensus,
P (Y =y|Y 6=None, X, θ) = P (Y =y|X, θ) P (Y 6=None|X, θ) = P (Ys=y|X, θ)S∑ y∈Y P (Ys=y|X, θ) S (10)
simply represents a “reparameterised” softmax.
In datasets where the noconsensus inputs are known it is clear we should use the full likelihood, (Eq. 6 and 7). The question is: in real-world datasets, where we only know the consensus inputs and the noconsensus inputs have been thrown away (Fig. 2C), can we use Eq. (10), a reparameterisation of the softmax-categorical probabilities, for the known consensus points? The answer is no because in Bayesian inference, we do not get to just pick a sensible-looking conditional probability distribution, such as Eq. 10, to use as the likelihood. Instead, we need to write down the full generative model,
and marginalise over unknown latents. To write down the generative model in the case of unknown noconsensus images (Fig. 2C), we need to take X to be an unobserved latent variable, and take Z to be an observed random variable,
Z|X, {Ys}Ss=1 = { X if Y1=Y2= · · · =Ys None otherwise
(11)
which is the underlying image, X , if consensus is reached, and None otherwise. As Z depends on θ (Fig. 2C), we cannot take the usual shortcut of using P (Y |Z, θ); we must instead use the full likelihood, P (Y, Z|θ),
P (Y,Z|θ) = ∑ {Ys}Ss=1 ∫ dX P ( X,Y, {Ys}Ss=1, Z|θ ) =
∑ {Ys}Ss=1 ∫ dX P (X) [∏S s=1 P (Ys|X, θ) ] P ( Y |{Ys}Ss=1 ) P ( Z|X, {Ys}Ss=1 ) .
(12)
for y not None,
P ( Y = y|{Ys}Ss=1 ) = { 1 if Ys = y for all s ∈ {1, . . . , S} 0 otherwise
(13)
and,
P ( Z|X, {Ys}Ss=1 ) = { δ(Z −X) if Y1 = Y2 = · · · = YS IZ=None otherwise
(14)
where δ is the Dirac-delta function, and where the indicator function, IZ=None is 1 if Z = None and 0 otherwise. Substituting Eq. (13) and Eq. (14) into Eq. (12),
P (Y = y, Z|θ) = P (X) S∏ s=1 P (Ys = y|X, θ) ∝ P (Ys = y|X, θ)S , (15)
where the proportionality arises because P (X) does not depend on the parameters of interest, θ. Note that this is proportional to Eq. (6) above, and not Eq. (10), so this does not just represent a reparameterisation of the softmax.
Finally, this explanation for the cold posterior effect would suggest that we would see cold posteriors even in maximum a-posteriori (MAP) inference, as we confirm in Appendix A. This is expected as it is “common knowledge” (but we do not know of a good reference) that while weight-decay is closely related to MAP inference with Gaussian priors, the best performing value of the weight decay coefficient tend to be lower than those suggested by untempered MAP inference.

heading: 4 DIFFERENCES BETWEEN COLD POSTERIOR AND DATASET CURATION SETUPS

text:
Our model of data curation provides a direct explanation for the effectiveness of cold posteriors, as Eq. (6) takes the underlying likelihoods, P (Ys|X, θ) to the power S, which has exactly the same
effect as 1/λ in tempered posteriors (Eq. 2). However, it is important to note two differences between the cold posterior setup and the ideal setup for our generative model. First, our generative model assumes that the noconsensus points are known, or if they are unknown, we can compute the integral,
P (Y =None, Z=None|θ) = ∫ dX P (X) 1−∑ y∈Y P (Ys = y|X, θ)S  (16)
which is obtained by evaluating Eq. (12) in the case where Y =None, and substituting Eq. (7). Of course, in reality it is not possible to compute this integral because we do not know the underlying P (X) (and usually, we do not even have samples from this distribution). In contrast, the standard cold-posterior setup entirely ignores these terms. Second, in the usual setting, the test data is also subject to the same consensus-formation process, in which case, we should use Eq. (10) for prediction. In contrast, in the standard cold-posterior setting, we use the single-labeller distribution, P (Ys|X, θ). (Note that if the test-set was drawn “from the wild”, without dataset curation, and labelled by a single labeller, then we should use P (Ys|X, θ) for prediction). Overall then, while our model of data curation offers a potential explanation of the benefits of tempering, the differences in setup imply that we cannot expect λ∗ = 1/S, to hold exactly, where λ∗ is the optimal temperature, and this is confirmed in our results on toy data below.

heading: 5 RESULTS

heading: 5.1 DATA SAMPLED FROM A KNOWN GAUSSIAN PROCESS MODEL

text:
In the introduction, we took it as given that if the prior is correct, the optimal approach is to use the true Bayesian posterior, avoiding either tempering or cold posteriors. To check that this is indeed the case, we tested the performance, measured as test-log-likelihood for tempered posteriors using data generated from a known Gaussian process model. We uniformly generated 50 input points on the 1D interval [-10, 10], and used a Gaussian process with squared exponential kernel with a standard deviation of 4, and kernel bandwidth of 1. For inference, we used reparameterised VI, with an approximate posterior given by multiplying the prior by a single Gaussian factor for each datapoint (Ober & Aitchison, 2020).
In particular, we used a Gaussian process prior for the function values, u, at the training inputs (Williams & Rasmussen, 2006),
P (f |x) = N (f ;0,K(x)) (17) and we use an approximate posterior (Hensman et al., 2015; Matthews et al., 2016; Ober & Aitchison, 2020) defined by multiplying the prior by a Gaussian with diagonal covariance, where we treat v and
Λ−1 as variational parameters.
Q (f) ∝ P (f |x)N ( v; f ,Λ−1 ) = N (f ;ΣΛv,Σ) where Σ = ( K−1(x) + Λ )−1 (18)
Note that this captures the true posterior in the case of regression (where v is set to y and Λ−1 is the output noise covariance; Ober & Aitchison, 2020). We then optimize the tempered ELBO (e.g. Zhang et al., 2017),
L = EQ(f) [log P (y|f) + λ (log P (f |x)− logQ (f))] . (19)
And we used the standard GP approach for prediction at test points (Williams & Rasmussen, 2006).
Initially, we tried using standard regression and classification generative models, without including our model of consensus formation. Unsurprisingly, λ = 1, corresponding to the Bayesian posterior, is optimal for GP regression (Fig. 4A) and classification (Fig. 4B). For GP regression, we use a Gaussian likelihood with standard deviation 1, and for classification, we use a standard sigmoid probability with a Bernoulli likelihood.
Next, we confirmed that even under our new, more complex generative model of dataset curation λ = 1 was still optimal (Fig. 4C), if we trained and tested using the correct form for the loglikelihood. In particular, we considered no-consensus inputs as known and used Eq. 7 to incorporate their likelihood. However, in standard benchmarks, only the consensus inputs are known. As such, next we generated curated data, trained using P (Ys|X, θ) and ignoring noconsensus points, but tested on the correct likelihood, including noconsensus points (Fig. 4D); the optimal temperature remains around 1. Finally, we considered the standard cold/tempered posterior setup, where we use P (Ys|X, θ), excluding noconsensus points, for training and testing. The optimal λ indeed fell with S (Fig. 4EF), giving a potential explanation for the cold posterior effect. As expected, we have S ≈ 1/λ∗, but the relation does not appear to hold exactly due to the mismatches discussed in Sec. 4. Finally, we show in Appendix B that the same patterns emerge in a Bayesian neural network, using Langevin sampling for inference.

heading: 5.2 CURATED AND UNCURATED GALAXY ZOO DATASETS

text:
The most direct test of our theory is to evaluate the cold-posterior effect on real-world curated and uncurated datasets. To this end, we used the Galazy Zoo 2 dataset (Willett et al., 2013), as the original dataset is not “curated” in our sense (the criteria for inclusion were e.g. brightness) and the dataset gives classifications from∼ 50 labellers. We used a reduced label set for simplicity (see Appendix C). For the uncurated dataset, we selected 12500 points at random from the full dataset, and for the curated dataset, we selected the 12500 most confident points (defined in terms of labeller agreement), such that the correct class-balance was maintained. These datasets were split at random into 2500 training points and 10000 test points. Note that using our generative model directly would lead to drastically different class-balance in the curated and uncurated datasets, due to different levels of certainty for different classes. To perform probabilistic inference, we used SGLD with code adapted from (Zhang et al., 2019). In particular, we used their standard settings of a ResNet18, momentum of 0.9, and a cyclic cosine learning rate schedule. Due to the smaller size of our training set, we used longer cycles (600 epochs rather than 50), and more cycles (8 rather than 4).
As expected, we found that performance on curated data was far better than that on uncurated data, both in terms of test log-likelihood and test accuracy (Fig. 5AC). We found some cold-posterior effects in uncurated data, which is not surprising because there may be other causes of cold-posteriors such as model-mismatch or biases in SGLD. Critically though, we found much stronger cold-posterior effects for curated data than for uncurated data (Fig. 5BD). Moreover, these plots tend to understate the differences between curated and uncurated data. In particular, the proportional changes to the test-log-likelihood (which has an upper bound of zero) and the test error is for curated data is very large, with both changing by a factor of ∼ 3 as temperature falls. In contrast, proportional changes for test-log-likelihood and test-error are much smaller (only ∼ 20%).

heading: 5.3 CIFAR-10H

text:
One prediction made by our framework is that if we had access to the original underlying human labels on the full dataset, including images that were rejected because consensus was not reached, then tempering should not be necessary. Obviously, all this additional information is not available for standard datasets such as CIFAR-10. However, we are able to get close by considering the CIFAR-10H dataset (Peterson et al., 2019). The authors of this work asked around 50 human labellers to label the CIFAR-10 test-set. As we might expect given the careful curation that went in to creating the original CIFAR-10 dataset (Krizhevsky, 2009), for almost half of the datapoints, all∼ 50 labellers agreed, and more than three-quarters of images had 2 or fewer disagreements (corresponding to 4%) of labellers (Fig. 6A). While it is not possible to estimate S without having information about the unknown image distribution, the high level of agreement would indicate that the effective value of S is large — potentially even larger than 10.
Next, we trained a neural network on the ∼ 50 labels provided by the CIFAR-10H dataset. As these labels are provided for the test-set with 10,000 images, this required us to swap the identities of the test and training sets (so that our training set consisted of 10,000 points, each with around 50 labels, and the test set consists of 50,000 points, with the single label from the original CIFAR-10 dataset). We compared against training with the standard CIFAR-10 test set, containing the same images but
only with a single label. To perform probabilistic inference, we used SGLD with code adapted from (Zhang et al., 2019). In particular, we used their standard settings of a ResNet18, momentum of 0.9, and a cyclic cosine learning rate schedule. Due to the smaller size of our training set, we used longer cycles (150 epochs rather than 50), and more cycles (8 rather than 4). Importantly, we kept all the parameters of the learning algorithm the same for CIFAR-10H and our CIFAR-10 comparison. The only complication was that to keep the same effective learning rate for CIFAR-10H we need to take into account the number of labellers per datapoint (if we have 50 labels for a datapoint, the loss and the gradients of the loss are 50 times larger, resulting in step-sizes that are also 50 times larger). As such, to keep the same step-sizes, we divided the actual learning rate by 50. An alternative way to look at this is that we use the average log-likelihood per labeller (rather than per-image). Importantly, this change in learning rate, leaves the stationary distribution was unchanged, as changing the learning rate is not equivalent to tempering. All other parameters were left at the values specified by Zhang et al. (2019).
As expected, when training on the single-label from the original CIFAR-10 testset (blue lines Fig. 6BC), there are very large tempering effects. In contrast, when training on the ∼ 50 labels provided CIFAR-10H (orange), the effects of tempering are far smaller. In a sense, this is not surprising — using 50 labellers in effect makes the likelihood 50 times stronger, which is very similar to applying tempering. But this simplicity is the point: in the Bayesian setting we need to condition on all data — in our case all the labels, and once we do that, the cold posterior effect disappears.
Note that while tempering-effects are dramatically reduced, they have not been eliminated entirely. This is expected in our setting, both because of the mismatch between this setup and the exact setup (Sec. 4), but also because the inference method, SGLD, becomes more accurate as the minibatch-size increases, but we can never use full-batch in practical settings due to the large size of datasets such as CIFAR-10 (Welling & Teh, 2011). In SGLD minibatch gradients are used as a proxy for the gradient for the full dataset, but these minibatch gradients contain additional noise, and there is a potential that reducing the temperature may partially compensate for this additional noise. This is particularly evident if we consider (Wenzel et al., 2020, Fig. 6), which showed that cold-posterior effects can be amplified by using very small minibatch sizes — smaller minibatches imply larger variance in gradient estimates, and hence more potential for lower temperatures to compensate for that additional noise. That said, these effects do not appear to be significant for CIFAR-10 (Wenzel et al., 2020, Fig. 5), so we are in agreement with Wenzel et al. (2020) that minibatch noise are unlikely to be the primary source of tempering effects.

heading: 5.4 CIFAR-10 WITH NOISY LABELS

text:
Curated labels are almost certain to represent the true class, and this is one way to understand the statistical patterns induced by curation that are exploited by tempered posteriors. We therefore considered disrupting this property by adding noise to the labels. In the standard classification setting, this should have very little effect (except to shrink the logit outputs somewhat to give more
uncertain outputs). However, and as expected under our theory, we find that adding noise destroys and eventually reverses the cold-posterior effect (Fig. 7). This confirms that cold-posteriors are exploiting specific properties of the labels that are destroyed by adding noise. As such, cold posteriors are not likely to arise e.g. from failing to capture the prior over neural network weights (in which case adding noise to the outputs should have little effect). We suspect that the small improvement in performance at the first value of λ below 1 might be due to partial compensation for additional noise introduced by minibatch estimates of the gradient.

heading: 6 RELATED WORK

text:
Wenzel et al. (2020) introduced the cold-posterior effect in the context of neural networks, and proposed and dismissed multiple potential explanations (although none like the one we propose here).
Other past work, though not in the neural network context, argues that tempering may be important in the context of model misspecification (Grünwald, 2012; Grünwald et al., 2017). Critically, we believe that there may be many causes of cold-posterior like effects, including but not limited to curation, model misspecification and artifacts from SGLD. Ultimately, the contribution of each of these factors in any given setting will depend on the exact dataset and inference method in question. Importantly, this also means we do not necessarily expect there to be no tempering in uncurated data, merely that we should see less tempering in the case of uncurated data.
Finally, the closest paper is concurrent work (Adlam et al., 2020) raising the possibility that BNNs overestimate aleatoric uncertainty, in part because of high-quality labels available in benchmark datasets. However, they concluded that while cold posteriors might help us to capture our priors, they do not correspond to an exact inference procedure. In contrast, here we give a generative model of dataset curation in which tempered likelihoods emerge naturally even under exact inference methods.

heading: 7 DISCUSSION AND CONCLUSIONS

text: We showed that modelling the process of data-curation can explain the improved performance of tempering or cold posteriors in Bayesian neural networks. While this work does provide a justification for future practitioners to use tempering in their Bayesian neural networks, we would urge caution. Importantly, and as we confirmed in Fig. 5, there may be other causes of the cold-posterior effect, e.g. because inference is inaccurate (compensating for additional noise added in SGLD), or may compensate for issues arising from model mispecification (Grünwald, 2012; Grünwald et al., 2017). As such, we urge practitioners to regard tempering with caution: if a very large amount of tempering is necessary to achieve good performance, it may indicate issues with either inference or the prior, and fixing these issues is of the utmost importance to obtaining accurate uncertainty estimation. More importantly, we hope that our work will prompt more careful dataset design, and further study of how data curation might impact downstream analyses in machine learning. Indeed, there are initial suggestions that semi-supervised learning methods are also exploiting the artificial clustering (Fig. 3) induced by data curation (Aitchison, 2020).

heading: ACKNOWLEDGEMENTS

text: I would like to thank Adrià Garriga-Alonso and Sebastian Ober for insightful discussions, Stoil Ganev for the GZ2 analyses which sadly came in only during the ICLR review period at which point author changes are banned (which is odd, given that revising the paper during the review period is allowed) and to Mike Walmsley and Sotiria Fotopoulou for getting us set up with GZ2. I would also like to thank Bristol’s Advanced Computing Research Centre (ACRC) for providing invaluable compute infrastructure that was used for all the experiments in this paper.

heading: A COLD POSTERIORS IN MAXIMUM A-POSTERIORI INFERENCE

text:
Our explanation for the cold posterior effect suggests that it is not just due to the stochasticity of the posterior, but should also arise in e.g. maximum a-posteriori (MAP) inference. However, MAP inference causes serious issues in networks with batchnorm. In particular, with batchnorm, the outputs become invariant to the scale of the weights, so the weights can continually decay towards zero (Van Laarhoven, 2017). Note that these considerations are not relevant when we are doing full approximate inference, as the product of the prior and likelihood is a static quantity with a well-defined scale that cannot e.g. decay towards zero. As such, we considered tempering in our usual ResNet18, but where we have deleted the batchnorm layers. We used a standard protocol for these types of network: SGD with momentum of 0.9 and an initial learning rate of 0.1, followed by decays to 0.01 and 0.001 at epochs 150 and 200. We read off performance at the final epoch (250) and used a batch size of 128. We use a standard prior over the weights, with variance 2/fan-in, to keep the scale of the inputs and outputs similar.
Remarkably, we found extremely strong cold-posterior effects (Fig. 8). These are perhaps expected as it is common-knowledge (but difficult to find a reference for) that while weight-decay is closely related to MAP inference with a Gaussian prior, good settings for the weight-decay coefficient are much smaller than those suggested by MAP inference with a sensible prior. Interestingly, these results suggest that the robustness of standard models to very low temperatures might arise due to batchnorm, and not be a fundamental property (e.g. having very large amounts of data).

heading: B TOY BAYESIAN NEURAL NETWORK EXAMPLE

text:
To confirm the results in the Gaussian processes toy example in the main text, we did the same exercise with a toy Bayesian neural network. In particular, we considered 100 training examples, drawn IID from standard Gaussians, each with 5 input dimensions and 2 class outputs. We generated the “true” or target outputs from a one hidden layer ReLU network, with 30 hidden units, drawn from the prior. We trained the same network, using Langevin dynamics (full batch, but no rejection step; initialized randomly from the prior). As in the GP example, we applied our generative model of dataset curation with different settings for S in the generative model. For test and training, we used the standard single-labeller log-likelihood, which ignores rejected noconsensus inputs. We implemented a highly parallelised Langevin sampling algorithm, which allowed us to draw 500 different datasets, and run one chain for each dataset (and each setting of λ in parallel). The large number of chains allowed us to obtain very accurate results, but did mean that it was not possible to plot lines for individual datasets. We used a learning rate of 0.01 and burn in of 4,000 steps; to compute performance, we averaged over 50 samples of neural-network parameters, which were taken from a chain that was 20,000 steps long in total.
In Fig. 9A, we plot the test-log-likelihood for different settings of λ. There is a clear trend that as S increases, the optimal temperature falls, and at the same time, the curves become flatter. Notably, these curves become flatter at the top for larger settings of S, making it difficult to identify the optimal setting for λ at higher values of S. We then plotted the optimal setting of λ against S (Fig 9B). Initially, S = 1 is equivalent to the standard classification setup, as consensus is always reached, and as such, we find that the optimal temperature was 1. As in Fig. 4E, we found that λ∗ = 1/S held only approximately. Nonetheless, these results suggest that we might need more tempering than suggested by λ∗ = 1/S, which supports our proposed explanation of the cold posterior effect. Finally, it should be noted that these results depend sensitively on the accuracy of the inference algorithm, so it is difficult to say as yet the degree to which mismatch between λ∗ and 1/S arises due to inaccuracy in inference or would emerge in the intractable true-posterior.

heading: C GALAXY ZOO 2 DETAILS

text:
The Galaxy Zoo 2 data set (Willett et al., 2013) was constructed by presenting images of galaxies to volunteers, who were asked to answer a questionnaire about the morphological features of the presented galaxy. These questions are not independent from each other but are part of a decision tree. As such, the questions a volunteer gets asked depends on their answers to earlier questions. For each galaxy, answers from multiple volunteers were collect. The raw data is recorded as counts of how many times each answer was picked across all volunteers.
From this data we wanted to derive classes in such a way that the classification performed by each volunteer was equivalent to a vote towards one of our classes. The most straightforward approach is to consider each unique path in the decision tree to be a unique class. Given that a volunteer fills in the questionnaire only once, their classification is equivalent to a single vote towards a specific class. The problem with this approach is that there are 1265 possible paths in the original decision tree. The number of volunteers per galaxy is less than a hundred, which means that the majority of these paths would receive 0 votes.
In order to address this, we simplified the tree to one where there are only 9 possible paths. The simplifications was performed by manually pruning branches of the decision tree. This way each path in the original tree could be mapped to a path in the simplified one, while maintaining specific aspects of its meaning. The simplified tree can be seen in (modified_graph.png). The following pruning operations were performed in order to create this simplified tree. Here we refer to each question using its task number, as given in the original paper.
• Task 02, Answer ’no’: leads directly to Task 05, skipping Task 03, Task 04, Task 10 and Task 11 • Task 05, Answers ’dominant’ and ’obvious’: merged together becoming a combined category
labelled ’obvious’ • Task 05, all answers: lead directly to the end of the graph, becoming exit points • Task 07, all answers: lead directly to the end of the graph, becoming exit points • Task 09, Answers ’rounded’ and ’boxy’: merged together becoming a combined category
labelled ’with bulge’ • Task 09, all answers: lead directly to the end of the graph, becoming exit points • As a consequence of the above pathing changes Task 03, Task 04, Task 06, Task 08, Task 10
and Task 11 are no longer reachable and thus removed from the graph
If we are to take each of the 9 paths in the simplified tree to form a class, then the resulting 9 classes would have the following meanings:
• smooth galaxy, completely round shape • smooth galaxy, in-between round and cigar shaped • smooth galaxy, cigar shaped • edge-on disk galaxy, with central bulge • edge-on disk galaxy, no central bulge • face-on disk galaxy, obvious central bulge • face-on disk galaxy, just noticeable central bulge • face-on disk galaxy, no central bulge • star or artifact
Using the simplified decision tree, we can decompose the data into votes for each path. The Galaxy Zoo data was recorded as counts for each specific answer, thus losing the information about which route a specific volunteer took in the tree. However, with the help of network flow algorithms it is possible to translate the data into counts of how many volunteers took a specific path. This way, by taking each path to be a class and the number of volunteers that followed a path to be the number of votes toward a class, we have a data set where each data point has received multiple classification over a fixed range of classes.
For each galaxy in the Galaxy Zoo 2 data set we also have the image that was presented to the volunteers. As such, we can use our derived classes in an image classification task. In order to do this, we performed certain pre-processing. The original images have a resolution of 424x424 with the galaxy being at the center of the image. Thus, we performed a center crop of size 212x212 and then resized the images down to 32x32. For data augmentation, we used full 360 degree continuous rotations.

title: A STATISTICAL THEORY OF COLD POSTERIORS IN DEEP NEURAL NETWORKS