Joint Perception and Control as Inference with an Object-based Implementation

1 INTRODUCTION . Human-like computing , which aims at endowing machines with human-like perceptual , reasoning and learning abilities , has recently drawn considerable attention ( Lake , 2014 ; Lake et al. , 2015 ; Baker et al. , 2017 ) . In order to operate within a dynamic environment while preserving homeostasis ( Kauffman , 1993 ) , humans maintain an internal model to learn new concepts efficiently from a few examples ( Friston , 2005 ) . The idea has since inspired many model-based reinforcement learning ( MBRL ) approaches to learn a concise perception model of the world ( Kaelbling et al. , 1998 ) . MBRL agents then use the perceptual model to choose effective actions . However , most existing MBRL methods separate perception modeling and decision making , leaving the potential connection between the objectives of these processes unexplored . A notable work by Hafner et al . ( 2020 ) provides a unified framework for perception and control . Built upon a general principle this framework covers a wide range of objectives in the fields of representation learning and reinforcement learning . However , they omit the discussion on combining perception and control for partially observable Markov decision processes ( POMDPs ) , which formalizes many real-world decision-making problems . In this paper , therefore , we focus on the joint perception and control as inference for POMDPs and provide a specialized joint objective as well as a practical implementation . Many prior MBRL methods fail to facilitate common-sense physical reasoning ( Battaglia et al. , 2013 ) , which is typically achieved by utilizing object-level inductive biases , e.g. , the prior over observed objects ’ properties , such as the type , amount , and locations . In contrast , humans can obtain these inductive biases through interacting with the environment and receiving feedback throughout their lifetimes ( Spelke et al. , 1992 ) , leading to a unified hierarchical and behavioral-correlated perception model to perceive events and objects from the environment ( Lee and Mumford , 2003 ) . Before taking actions , a human agent can use this model to decompose a complex visual scene into distinct parts , understand relations between them , reason about their dynamics and predict the consequences of its actions ( Battaglia et al. , 2013 ) . Therefore , equipping MBRL with object-level inductive biases is essential to create agents capable of emulating human perceptual learning and reasoning and thus complex decision making ( Lake et al. , 2015 ) . We propose to train an agent in a similar way to gain inductive biases by learning the structured properties of the environment . This can enable the agent to plan like a human using its ability to think ahead , see what would happen for a range of possible choices , and make rapid decisions while learning a policy with the help of the inductive bias ( Lake et al. , 2017 ) . Moreover , in order to mimic a human ’ s spontaneous acquisition of inductive biases o1 s1 a1 o2 a2 ……θ1 s2 θ 2 x1 x2 q ( s ) p ( s |x , o ) p ( x , o |s ) p ( s ) throughout its life , we propose to build a model able to acquire new knowledge online , rather than a one which merely generates static information from offline training ( Dehaene et al. , 2017 ) . In this paper , we introduce joint Perception and Control as Inference ( PCI ) as shown in Fig . ( 1 ) , a unified framework for decision making and perception modeling to facilitate understanding of the environment while providing a joint objective for both the perception and the action choice . As we argue that inductive bias gained in object-based perception is beneficial for control tasks , we then propose Object-based Perception Control ( OPC ) , an instantiation of PCI which facilitates control with the help of automatically discovered representations of objects from raw pixels . We consider a setting inspired by real-world scenarios ; we consider a partially observable environment in which agents ’ observations consist of a visual scene with compositional structure . The perception optimization of OPC is typically achieved by inference in a spatial mixture model through generalized expectation maximization ( Dempster et al. , 1977 ) , while the policy optimization is derived from conventional temporal-difference ( TD ) learning ( Sutton , 1988 ) . Proof of convergence for the perception model update is provided in Appendix A . We test OPC on the Pixel Waterworld environment . Our results show that OPC achieves good quality and consistent perceptual grouping and outperforms several strong baselines in terms of accumulated rewards . 2 RELATED WORK . Connecting Perception and Control Formulating RL as Bayesian inference over inputs and actions has been explored by recent works ( Todorov , 2008 ; Kappen et al. , 2009 ; Rawlik et al. , 2010 ; Ortega and Braun , 2011 ; Levine , 2018 ; Tschiatschek et al. , 2018 ; Lee et al. , 2019b ; a ; Ortega et al. , 2019 ; Xin et al. , 2020 ; O ’ Donoghue et al. , 2020 ) . The generalized free energy principle ( Parr and Friston , 2019 ) studies a unified objective by heuristically defining entropy terms . A unified framework for perception and control from a general principle is proposed by Hafner et al . ( 2020 ) . Their framework provides a common foundation from which a wide range of objectives can be derived such as representation learning , information gain , empowerment , and skill discovery . However , one trade-off for the generality of their framework is the loss in precision . Environments in many real-world decision-making problems are only partially observable , which signifies the importance of MBRL methods to solving POMDPs . However , relevant and integrated discussion is omitted in Hafner et al . ( 2020 ) . In contrast , we focus on the joint perception and control as inference for POMDPs and provide a specialized joint-objective as well as a practical implementation . Model-based Deep Reinforcement Learning MBRL algorithms have been shown to be effective in various tasks ( Gu et al. , 2016 ) , including operating in environments with high-dimensional raw pixel observations ( Igl et al. , 2018 ; Shani et al. , 2005 ; Watter et al. , 2015 ; Levine et al. , 2016 ; Finn and Levine , 2017 ) . Existing methods have considered incorporating reward structure into modellearning ( Farahmand et al. , 2017 ; Oh et al. , 2017 ) , while our proposed PCI takes one step forward by incorporating the perception model into the control-as-inference derivation to yield a single unified objective for multiple components in a pipeline . One of the methods closely related to OPC is the World Model ( Ha and Schmidhuber , 2018 ) , which consists of offline and separately trained models for vision , memory , and control . These methods typically produce entangled latent representations for pixel observations whereas , for real world tasks such as reasoning and physical interaction , it is often necessary to identify and manipulate multiple entities and their relationships for optimal performance . Although Zambaldi et al . ( 2018 ) has used the relational mechanism to discover and reason about entities , their model needs additional supervision of location information . Object-based Reinforcement Learning The object-based approach , which recognizes decomposed objects from the environment observations , has attracted considerable attention in RL as well ( Schmidhuber , 1992 ) . However , most models often use pre-trained object-based representations rather than learning them from high-dimensional observations ( Diuk et al. , 2008 ; Kansky et al. , 2017 ) . When objects are extracted through learning methods , these models usually require supervised modeling of the object property , by either comparing the activation spectrum generated from neural network filters with existing types ( Garnelo et al. , 2016 ) or leveraging the bounding boxes generated by standard object detection algorithms ( Keramati et al. , 2018 ) . MOREL ( Goel et al. , 2018 ) applies optical flow in video sequences to learn the position and velocity information as input for model-free RL frameworks . A distinguishing feature of our work in relation to previous works in MBRL and the object-based RL is that we provide the decision-making process with object-based abstractions of high-dimensional observations in an unsupervised manner , which contribute to faster learning . Unsupervised Object Segmentation Unsupervised object segmentation and representation learning have seen several recent breakthroughs , such as IODINE ( Greff et al. , 2019 ) , MONet ( Burgess et al. , 2019 ) , and GENESIS ( Engelcke et al. , 2020 ) . Several recent works have investigated the unsupervised object extraction for reinforcement learning as well ( Zhu et al. , 2018 ; Asai and Fukunaga , 2017 ; Kulkarni et al. , 2019 ; Watters et al. , 2019 ; Veerapaneni et al. , 2020 ) . Although OPC is built upon a previous unsupervised object segmentation back-end ( Greff et al. , 2017 ; van Steenkiste et al. , 2018 ) , we explore one step forward by proposing a joint framework for perceptual grouping and decision-making . This could help an agent to discover structured objects from raw pixels so that it could better tackle its decision problems . Our framework also adheres to the Bayesian brain hypothesis by maintaining and updating a compact perception model towards the cause of particular observations ( Friston , 2010 ) . 3 METHODS . We start by introducing the environment as a partially observable Markov Decision Process ( POMDP ) with an object-based observation distribution in Sect . 3.11 . We then introduce PCI , a general framework for joint perception and control as inference in Sect . 3.2 and arrive at a joint objective for perception and control models . In the remainder of this section we propose OPC , a practical method to optimize the joint objective in the context of an object-based environment , which requires the model to exploit the compositional structure of a visual scene . 3.1 ENVIRONMENT SETTING . We define the environment as a POMDP represented by the tuple = hS , P , A , X , U , Ri , where S , A , X are the state space , the action space , and the observation space , respectively . At time step t , we consider an agent ’ s observation xt 2 X ⌘ RD as a visual image ( a matrix of pixels ) composited of K objects , where each pixel xi is determined by exactly one object . The agent receives xt following the conditional observation distribution U ( xt|st ) : S ! X , where the hidden state st is defined by the tuple ( zt , ✓t1 , . . . , ✓tK ) . Concretely , we denote as z t 2 Z ⌘ [ 0 , 1 ] D⇥K the latent variable which encodes the unknown true pixel assignments , such that zti , k = 1 iff pixel z t i was generated by component k. Each pixel xti is then rendered by its corresponding object representations ✓tk 2 RM through a pixel-wise distribution U ti , k ( x t i|zti , k = 1 ) 2 , where ti , k = f ( ✓tk ) i is generated by feeding ✓tk into a differentiable non-linear function f . When the environment receives an action a t 2 A , it moves to a new state st+1 following the transition function P ( st+1|st , at ) : S⇥A ! S . We assume the transition function could be parameterized and we integrate its parameter into . To embed the control problem into the graphical model , we also introduce an additional binary random variable ot to represent the optimality at time step t , i.e. , ot = 1 denotes that time step t is optimal , and ot = 0 denotes that it is not optimal . We choose the distribution over ot to be p ( ot = 1|st , at ) / exp ( rt ) , 1Note that the PCI framework is designed for general POMDPs . We extend Sect . 3.1 to object-based POMDPs for the purpose of introducing the environment setting for OPC . 2We consider U as Gaussian , i.e. , U ti , k ( x t i|zti , k = 1 ) ⇠ N ( xti ; µ = ti , k , 2 ) for some fixed 2. where rt 2 R is the observed reward provided by the environment according to the reward function R ( rt|st , at ) : S ⇥A ! R. We denote the distribution over initial state as p ( s1 ) : S ! [ 0 , 1 ] .

Summary: This paper proposes an extension of the RL as Inference framework, and demonstrates how to use it to express an object-centric RL model and train it on simple environments. It appears to be a combination of NEM [1] with a simple TD-learning objective on top. Results are a bit hard to interpret but seem promising.