Joint Perception and Control as Inference with an Object-based Implementation
Authors: Minne Li
Zheng Tian
Pranav Nashikkar
Ian Davies
Ying Wen
Jun Wang
Venue:  (2021)
Keywords: model-based reinforcement learning
perception modeling
object-based reinforcement learning

Abstract: Existing model-based reinforcement learning methods often study perception modeling and decision making separately. We introduce joint Perception and Control as Inference (PCI), a general framework to combine perception and control for partially observable environments through Bayesian inference. Based on the fact that object-level inductive biases are critical in human perceptual learning and reasoning, we propose Object-based Perception Control (OPC), an instantiation of PCI which manages to facilitate control using automatic discovered object-based representations. We develop an unsupervised end-to-end solution and analyze the convergence of the perception model update. Experiments in a high-dimensional pixel environment demonstrate the learning effectiveness of our object-based perception control approach. Specifically, we show that OPC achieves good perceptual grouping quality and outperforms several strong baselines in accumulated rewards.

Under review as a conference paper at ICLR 2021
JOINT PERCEPTION AND CONTROL AS INFERENCE WITH AN OBJECT-BASED IMPLEMENTATION
Anonymous authors Paper under double-blind review
ABSTRACT
Existing model-based reinforcement learning methods often study perception mod- eling and decision making separately. We introduce joint Perception and Control as Inference (PCI), a general framework to combine perception and control for partially observable environments through Bayesian inference. Based on the fact that object-level inductive biases are critical in human perceptual learning and reasoning, we propose Object-based Perception Control (OPC), an instantiation of PCI which manages to facilitate control using automatic discovered bbjecttbbsed representations. We develop unsupervised end-to-end solution and analyze the convergence of the perception model update. Experiments in high- -i-memsiiiinal pixel environment demonstrate the learning effectiveness of our object- based per- ception control approach. Specifically, we show that OPC achieves good perceptual grouping quality and outperforms several strong baselines in accumulated rewards.
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 at.) 2015)
Baker et al., (017). In order to operate within a dynamic environment while preserving homeosta- sis (Kauffman 1193), humans maintain internal model to learn new concepts efficiently from 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 |Hatner et al. (2020) provides a unified framework for perception and control. Built upon 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 deci- sion processes (POMDPs), which formalizes many real- world recision-making problems. In this paper, therefore, we focus on the joint perception and control as inference for POMDPs and provide specialized joint objective 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 el al. 1992). leading to a unified hierarchical and rehavioral-correlaee perception model perceive events and objects from the environment (Lee and Mumford] 2003). Before taking actions, a human agent can use this model to decompose complex visual scene into distinct parts, understand relations between them, reason about their dynamics and predict the consequences of its actions (Battaglia er 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 train 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 human using its ability to think ahead, what would happen for 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 mimic human's spontaneous acquisition of inductive biases

---

Under review as a conference paper at ICLR 2021
q(s) ~ p(s|(x)
plx, ols)
p(s)
The graphical model of joint Perception and Control as Inference (PCT), where and o represent the latent state and the binary optimality binary variable, respectively. The hierarchical perception model includes a bottom-up recognition model (()) and a top-down generative model p(æ,, o, s) (decomposed into the likelihood p(ær, o o|s) and the prior belief p(s)). Control is performed by taking an action a to affect the environment state.
Figure I:
throughout its life, we propose to build model able acquire new knowledge online, rather than one which merely generates static information from offline training (Dehaene et al. 2017].
In this paper, introduce joint Perception and Control as Inference (PCI) shown in Fig. I unified framework for decision making and perception modeling to facilitate understanding of the environment while providing joint objective for both the perception and the action choice. As 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 setting inspired by real- -world scenarios; we consider a partially observable environment which agents' observations consist of 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-differene (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 (Tooorov 2008 Kappen et al. 2009 Rawlik er al. 2010
Ortega and Braun 201T: Levine 2018: Tschiatschek et al. 2018 Lee et al. 2019b a Ortega er
al.12019 Xin et al, 2020 O' D'Doooghue et al.) 2020). The generalized free energy principle (Par
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 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-male 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, focus the joint perception and control as inference for POMDPs and provide a specialized joint-objective as well as practical implementation.
Model-based Deep Reinforcement Learning MBRL algorithms have been shown to be effective in various tasks (Gu al. 2016), including operating in environments with high-dimensional raw
pixel observations (Igl el al. 2018 Shani el al. 2003 Watter el al. 2015 Levine el al. 2016 Finn and Levine, 2017). Existing methods have considered incorporating reward structure into model-
learning (Farahmand er al, 2017: Oh et al.) 2017). while our proposed PCI takes one step forward by incorporating the perception model into the caaurol-a--infernnee derivation to yield single unified objective for multiple components in 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, 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.
2

---

Under review as a conference paper at ICLR 2021
Object-based Reinforcement Learning The object-based approach, which recognizes decom- posed objects from the environment observations, has attracted considerable attention in RL as well (Schmidhuber 1992). However, most models often use pre-trained object-basck representations rather than learning them from high-dimensional observations (Diuk et al.) 2008 Kansky et al. 2017). When objects 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 er 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.
distinguishing feature of our work in relation to previous works in MBRL. and the object-based RI. is that we provide the decision-making process with object-based abstractions of high- high-dimensional observations in an unsupervised manner, which contribute to faster learning.
Unsupervised Object Segmentation Unsupervised object segmentation and representation learn- ing have seen several recent breakthroughs, such as IODINE (Grell' et al. 2019). MONet (Burgess ef al. 2019), and GENESIS (Engelcke et al. 2020). Several recent works have investigated the unsu- pervised object extraction for reinforcement learning as well (Zhu er al. 2018 Asai and Fukunaga
2017. Kulkarni et al.) 2019 Watters el al. 2019| Veerapaneni et al. (0200. Although OPC is built
upon previous unsupervised object segmentation back-end (Greff et al. 2017 van Steenkiste el
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.1 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, practical method to optimize the joint objective the context of an object-based environment, which requires the model to exploit the compositional structure of visual scene.
3.1 ENVIRONMENT SETTING
We define the environment as a POMDP represented by the tuple = (A, R),, 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 's observation x' E = IR D as a visual image (a matrix of pixels) composited of K objects, where each pixel x; is determined by exactly one object. The agent receives æ' following the conditional observation distribution U (z=(s²²) (os)::= : S + X, where the hidden state s is defined by the tuple (z', ...... 0xx). Concretely, we denote as zt E = [0, 1 X the latent variable which encodes the unknown true pixel assignments, such that 2 = 1 iff pixel 2f was generated by component k. Each pixel at is then rendered by its corresponding object representations 0 RM through pixel-wise distribution Uyg (=:)=f,,- B where vif = S (0)); is generated by feeding 0% into a differentiable non-linear function fa. When the environment receives an action a' E A. it moves to a new state g(+1 following the transition function P(s(-1|s( a²) S x A + S. We assume the transition function could be parameterized and integrate its parameter into d. To embed the control problem into the graphical model, we also introduce an additional binary random variable of to represent the optimality time step f, i.e., o 1 denotes that time step is optimal, and o' denotes that it is not optimal. We choose the distribution over of to be p(o' 1|s a²) x exp(r*),
1 Note that the PCI framework is designed for general POMDPs. We extend Sect. 3.1 object-based
POMDPs for the purpose of introducing the environment setting for OPC.
We consider U4 as Gaussian, i.c., Uyg (x:) of 1) ~ ~N(2f 0?) for some fixed o².
3

---

Under review as a conference paper at ICLR 2021
where p E R the observed reward provided by the environment according to the reward function
R(r" s :x): R. We denote the distribution over initial state as p(s') + [0.
3.2 JOINT PERCEPTION AND CONTROL. AS INFERENCE
To formalize the belief about the unobserved hidden cause of the history observation atS the agent maintains perception model (s) approximate the distribution over latent states as illustrated Fig. ① An agent's inferred belief about the latent state could serve as a sufficient statistic of the history and be used as input to its policy т, which guides the agent to act the environment.
The goal of the agent maximize the future optimality o while learning a perception model by inferring unobserved temporal hidden states given the past observations x and actions att, which achieved by maximizing the following objective:
log p(o=', x a² log
Σ
p s21 x aa
II--1 p((")((s)) IT p(s
log
(I)
where we denote 9" = *q(55t St at and use ! represent the term related control as
Σ
Il p(o"
II p(a +1 k+1 )p(s +1 aaa)p((aaaa
humt
The full derivation is presented in Appendix 4td p(s') I[ 4(s" <8 a ()), where we slightly abuse notation for q" by ignoring the fact that we sample from the model p(s)) fort We then apply Jensen's inequality to Eq and get
We denote qtr q(s a ax and assume
log IT p(ar a² p')p(s'
p(a aam)
log p(o2 S11a <t) 2 Eastngw
Eastago [log !
p(s') II (0
log II 102)
DEL
((s 5 م lbp(s
80 80,0%)
+B +E,s w [log !
E (0" .6)
where DKL represents the Kullback-Leibler divergence (Kullbakk and Leibler 1951). As intro- duced in Sect. 3.1 we parameterize the transition distribution Ps(s*t'[s' a²) and the observation distribution Pa s ø. An instantiation to optimize the above evidence lower bound (ELBO) L" (q"""""0 Φ) an environment with explicit physical properties and high-dimensional pixel-level
observations will be discussed in Sect. 33
We now derive the control objective by extending
II II. a,") Im Po (x 18 (ss 1s "r,a*")do aat)p(a (
log
Σ
a( a2: xx21)
x>1
where assume (882 at m>>1) I Ps (x / p (s 1s al (a)) and denote
of q(s> a²! x ))). We then apply Jensen's inequality to Eas [log ( and get
Im P p(""|s", "aa"""""") Ih Po (otl)+//))))dlll (o (b
Eastings [log (1) 2 Eystingu
Σ
log
Im (z ( + 1 (shttls at at
- E u
B Batnt a wopt [R(s' e')) DRu( ( Mapld
(2)
where denote Ps (= +1 ghh (ss (s+1 expectation under Pa. reward maximization also bias the learning of our perception model. We will propose an option optimize the control objective Eq. in Sect. 3.4
|| s", s aaa). Note that as this control objective is
4

---

Under review as a conference paper at ICLR 2021
3.3 OBJECT-BASED PERCEPTION MODEL UPDATE
We now introduce OPC, an instantiation of PCI in the context of an object-based environment. Following the generalized expectation maximization (Dempster al. 1977), we optimize the ELBO L" (q" Φ) by improving the perception model qth about the true posterior (s a'+ s ",," with respect to set of object representations O -0,..... ENCR RXK
E-step to compute with respect to yd at time step I, i.e., qr = pa(s; [ -8' _s'.a")= po(2; generate soft-assignment of each pixel to of the K objects
estimate q" of the posterior. Assume the ELBO is already maximized
[z; 2+1 1//, a²). we can
Th, = pol2 1|2 ..)
(3)
M-step to update the model parameter d. We then find 0 to maximize the ELBO as (x++1, x1+1, +b'tlss a²)] = arg max Ag. (0++1).
0++1 arg max E
(4)
g++1
zl Louay
Gr
See derivation in Appendix
C.2 Note that Eq. returns set of points that maximize Ag (0*+1)
and we choose a²+ be any value within this set. To update the perception model by maximizing
the evidence lower bound with respect to g(s") and 0C+1, we compute Eq.. B3 and Eq. (4) iteratively. However, an analytical solution to Eq. not available because we use a differentiable non-linear function fū to map from object representations 0 into V f0%)). Therefore, get 0%* by
0-' = h +
0f + Σ
Th ==1
D
Vi.k
x.
(5)
a0
o²
00 k
0
where a is the learning rate (see details of derivation in Appendix
We regard the iterative process step rollout of copies of recurrent neural net- work with hidden states 0g receiving mi @ (vbf x(+1) as input (see the inner loop of Algo- rithm D. Each copy generates new wbtt) which is then used to re-estimate the soft-assignments Mk +++
We parameterize the Jacobian anpt, /00 and the differentiable non-linear funetion før using convolutional encoder-decoder architecture with a recurrent neural network bottleneck, which linearly combines the out- put of the encoder with 0% from the previous time step. To fit the statistical model fo to capture the regularities corresponding to the observation and the transition distribution for given POMDPs, we back-propagate the gradient of Er~r [A (0++)) through the BPTT Werbos (1988)) into the weights ø. We demonstrate the convergence of sequence of (Cm (e"",))) generated by the perception update in Ap- pendix A The proof presented by showing that the learning process follows the Global Con- vergence Theorem Zangwill 1969).
Perception Model
Decision Making Module
Ageus (0'42)
TD-error
++2
Env
MLP
"time" (also known
Food Agent Poison
K{
Rei
gradient flow
Figure 2: Illustration of OPC and the Pixel Waterworld envi- ronment. We regard the perception update as K copies of an RNN with hidden states 0f receiving Mi (xb x'+1) input. Each copy generates
2/4+1 which is used re-estimate the soft-assignments m ++1 for the calculation of the expected log-likelihood Age+ (0--2), The hidden states are then used as the input to decision-making (before being fed into an optional relational module) to guide the agent's action.
Our implementation of the per- ceptual model is based on the un- supervised perceptual grouping method proposed by (Greff et al.
2017 Van Steenkiste et al. 2018).
However, using other advanced unsupervised object representations learning methods as the perceptual model (such as IODINE (Greff et al,) 2019) and MONet (Burgess) al. 2019)) is straightforward extension.
5

---

Under review as a conference paper at ICLR 2021
Algorithm Learning the Object-based Perception Control
Initialize 0, 17, O, Ço, Topi, fro, K
Create N. environments that will execute in parallel
while training
finished do Initialize history D.: Dz, Dy with environment rollouts for tro time-steps under current policy m for 7' = to
tr do
dø < 0 Get "y from D,
for
fro do
Get a², ,x from Da
D2 respectively
Feed 77 ④ (
x into cach of the K RNN copy to get 0j and forward-output abf.
Compute 7h by Eq.
3
dΦ + d$ + a Ao (0)+)) by Eq.
Perform a T++, according to policy Theoo
(a e +T+)
Receive reward p and new observation xx +tm++1
Store a drtrro, x @+trott 77 in Da, Dz, Dy respectively
Feed 77 T+tre ( tro (+( into cach of the RNN copy to get 0 +rat1
y <- r" + + YV,, (
dç + V, log (a Tttro (0 (2 V.s (OT+tro)
dç. + a (y Vc. (0* 1)² /0c.
dé + dó a (y Vc. (0" (("=+)))) /06
Perform synchromous update of & using déb of using dç, and of Cu using dc.
3.4 DECISION-MAKING MODULE UPDATE
Recall that the control objective in Eq is the sum of the expected total reward along trajectory р& and the negative KL divergence between the policy and our prior about the future actions. The prior can reflect preference over policies. For example, uniform prior in SAC (Haarnoja et al. 2018) corresponds the preference for a simple policy with minimal assumptions. This could prevent early convergence sub-optimal policies. However, perception with control, we do not pre-impose a preference for policies. Therefore we set the prior equal to our training policy, resulting in a zero KL. divergence term in the objective.
focus studying the benefit of combing
To maximize the total reward along trajectory PaTT, we follow the conventional lemporal-differene (TD) learning approach (Sutton 1988) by feeding the object abstractions into small multilayer perceptron (MLP) (Rumelhart et al. 1986) to produce is split into a dimgl - (the "critic'.. The गे logits are normalized using softmax function, and used as the multinomial distribution from which action sampled. The Ve, is an estimate of the state-value function the current state, which is given by the last hidden state Decision-making module, the Vc. is used to compute the temporal-different
(dimg((A) + 1)-dimensional vector, which vector of क5's 's (the 'actor') logits, and baseline scalar VG,
the step RNN rollout. On training the
error given by
CTD (v/
Ve. (0**1))² v²- ,
~V.. (V...0 (0**))
(6)
where E [0. 1) is a discount factor. LTD used both to optimize πę to generate actions with larger total rewards than VG. predicts by updating with respect to the policy gradient
Ve log me (a (0 (0* (v
(O*+))).
and to optimize Vg. to accurately estimate state values by updating Ço- Also, differentiating LTD with respect to ø enables the gradient-based optimizers to update the perception model. We provide the pseudo-code for one-step TD-leannin of the proposed model in Algorithm By grouping objects concerning the reward, our model distinguishes objects with visual similarities but different semantics, thus helping the agent to better understand the environment.
EXPERIMENTS
4.1 PIXEL WATERWORLD
We demonstrate the advantage of unifying decision making and perception modeling by applying OPC an environment similar to the one used in COBRA (Watters al. 2019). modified

---

Under review as a conference paper at ICLR 2021
Waterworld environment (Karpathy 2013). where the observations 84 x 84 grayscale raw pixel images composited of an agent and two types of randomly moving targets: the poison and the food, as illustrated in Fig. 2 The agent can control its velocity by choosing from four available actions: to apply the thruster to the left, right, up and down. negative cost given to the agent it touches any poison target, while positive reward for making contact with any food target. The optimal strategy depends on the number, speed, and size of objects, thus requiring the agent to infer the underlying dynamics of the environment within given amount of observations.
The intuition of this environment is to test whether the agent can quickly learn dynamics of new environment online without any prior knowledge. i.e., the execution- time optimization throughout the agent's s life-cycle to mimic humans' spontaneous process of obtaining inductive biases. We choose Advantage Actor-Critic (A2C) (Mnih el al.) 2016) as the decision- maasinnn module of OPC without loss of generality, although other online learning methods are also applicable. For OPC, we use K and t,o 20 except for Sect. 4.5 where we analyze the effect of the hyper-paramete setting.
4.2 ACCUMULATED REWARD COMPARISONS
To verify object-level in- ductive biases facilitating decision making, we com- pare OPC against a set of baseline algorithms, in- cluding: 1) the standard A2C, which uses convo- lutional layers to trans- form the
0.0
Period Costs (*1e3)
Costs (*163)
10
-0.5
-1.5
A2C,
VAS-A20
3
3
# of experienced observations (*1e5)
f of experienced observations ("1e5)
pixel ob- servations to low dimen- sional vectors as input for the same MLP described
(a)
(b)
0.0
Period Costs (*1e3)
Period Costs (M))
-0.5
-0.5
in Sect. 3.4, 2) the World Model (Ha and Schmid- huber 2018) with A2C for control (WM-A2C), state-of-the-art model- based approach, which separately learns a vi- sual and a memorization model to provide the in- put for standard A2C, and 3) the random policy. For both the baseline A2C, WM-A2C and the decision-male makinn module of OPC, we follow the convention of Mnih et al. 2016) by running the algorithm in the forward view and using the same mix of n-step returns to update both πę and Vg, Details of the model and of hyperparameter setting can be found in Appendix B We build the environment with two poison objects and food object, and set the size of the agent 1.5 times smaller than the target. Results are reported with the average value among separate runs of three random seeds. Note that the training procedure of WM-A2C includes independent off-line training for each component, thus requiring many more observation samples than OPC. Following its original early- carly-stopping criteria, we report that WM-A2C requires 300 times more observation samples for separately training perception model before learning a policy to give the results presented in Fig. (
-1.0
-1.0
No 20
-15
Iep 20, (tll
10
1.5
3
3 # of experienced observations (*1e5)
2
# of experienced observations ("1e5)
(c)
(d)
Figure 3: Performance comparisons between methods.
Fig. Bab shows the result of accumulated rewards after each agent has experienced the same amount of observations. Note that OPC uses the plotted number of experienced observations for training the entire model, while WM-A2C uses the amount of observations only for training the control module. It is clear that the agent with OPC achieves the greatest performance, outperforming all other agents regardless due to the abject-level inductive biases obtained by the perception model (compared entangled representations extracted by the CNN used in standard A2C), and the unified learning of perception and control of OPC (as opposed to the separate learning of different modules in WM-A2C).
model in terms of accumulated reward. We believe this advantage is
To illustrate each agent's 's learning process through time, we also present the period reward, which is the accumulated reward during a given period of environment interactions (2e4 of experienced
7

---

Under review as a conference paper at ICLR 2021
observations). As illustrated in Fig. (b)). OPC significantly improves the sample efficiency of A2C, which enables the agent acting the environment to find an optimal strategy more quickly than agents with baseline models. We also find that the standard version of A2C with four parallel threads gives roughly the same result as the single-threaded version of A2C (the same as the decision-making module of OPC), eliminating the potential drawback of single-thread learning.
4.3 PERCEPTUAL GROUPING RESULTS
To demonstrate the performance of the perception model, we provide an example rollout at the early stage of training in Fig. 4 We assign different color to each 2/y. for better visualization and show through the soft-assigmment 1 that all objects are grouped semantically: the agent in blue, the food in green, and both poisons red. During the joint training of perception and policy, the soft-assignment gradually consti- tutes a semantic segmentation as the TD signal improves the recognition when grouping pixels into different types based the interaction with the environment. Consequently, the learned ob- ject representations becomes a semantic inter- pretation of raw pixels grouped into perceptual objects.
Figure 4: A sample rollout of perceptual grouping by OPC. The observation (top row), the next-step prediction uh of each copy of the K RNN copy (rows to 7), the Σn the (row 2). and the soft- assignment 77 of the pixels to each of the copies (row 3). We assign a different color to each 4hk for better visualization. This sample rollout shows that all objects are grouped semantically.
4.4 BENEFITS OF JOINT INFERENCE
To gain insights into the benefit of the control objective for joint inference, we further compare OPC against OP, perception model with the same architecture as OPC but no guidance of the TD signal from RL, i.e., using the same unsupervised object segmentation back-end as in Greff' el al. (2017). Fig- 5 shows the result of accumulated rewards and the soft-assignment n produced by both perception models. Results are produced by models running in environments with the same random seed.
As illustrated in Fig. (5a). OPC outperforms OP in terms of the accumulated rewards through time with a given amount of observations. This performance difference owes the fact shown in Fig. (b)). where the of OP are overlaid, and the shapes are not delineated in bold colors but are mixture, showing that OP has not learned segment the objects clearly. The pixel assignment (coloring) does not appear to converge even after extremely large number of iterations, suggesting that OP alone cannot properly show signs of any semantic adaptation to the task at hand.
On the other hand, the quality and consistency of the soft-assignment gen- erated by OPC are im- proved with the TD sig- nal, suggesting that RL facilitates the learning of the perception model. Furthermore. the inter- action between objects shown in Fig. (5b) demon- strates the agent is mov- ing away from the poison (even when the food is nearby). Because of the poison objects, believe
14000
14500
Costs (M1e3)
2
X
-6
DPC
m
1.5
experienced observations (*1e5)
OPC
OP
OPC
OP
(a)
(b)
Figure 5: (a) Performance comparison between OPC and OP. (b) The soft- assignment after experiencing 14000 and 14500 observations. Each th assigned with a different color as in Fig. original observation.
the binary-processed
environment setting of the high moving speed, larger target sizes and more that agents are given strong indications to stay away from all stimuli. Thus, the result shows an understanding of visual reasoning that the agent can separate itself from the rest of the objects.
8

---

Under review as a conference paper at ICLR 2021
4.5 ABLATION STUDY ON HYPER-PARAMETE PARAMETERS
We also investigate the effect of hyper-parameters to the learning of OPC, by changing: 1) the number of recurrent copies K, the rollout steps for recurrent iteration 1ro. and the use of relational mechanism across object representations 0 (see Sect. 2.2 of van Steenkiste el al. (2018) for more 3db show period reward results across different hyper-parameler settings. details). Fig. (3c and Fig.
As illustrated in Fig. (3). the number of recurrent copics affects the stability of OPC learning, as OPC with K 3 has experienced larger variance during the agent's s life-cycle. We believe the difference comes from the environment dynamics as we have visually four objects in the environment. During the earlier stage of interacting with the environment, OPC tries to group each object into distinct class; thus, a different number of K against the number of objects in the environment confuse the perception model and lead to unstable learning. Although different K settings might affect the learning stability and slow down the convergence. OPC can still find an optimal strategy within given amount of observations.
In Fig. Edl. compare OPC with different steps of recurrent rollout & A smaller tro means fewer rounds of perception updates and therefore slower convergence in terms of the number of experienced observations. We believe that the choice of tro depends the difficulty of the environment, e.g., smaller t.o can help to find the optimal strategy more quickly for simpler environments in terms of wall training time.
Meanwhile, results in Fig. (c0 and Fig. 3d show that the use of a relational mechanism has limited impact on OPC, possibly because the objects can be well distinguished and perceived by their object-level inductive biases, i.c. shapes in our experiment. We believe that investigating whether the relational mechanism will have impact on environments where entities with similar object-level inductive bias have other different internal properties is an interesting direction for future work.
5 CONCLUSIONS
In this paper, we propose joint Perception and Control as Inference (PCT), a general framework to combine perception and control for POMDPs through Bayesian inference. We then extend PCI the context of a typical pixel-level environment with compositional structure and propose Object-based Perception Control (OPC), an instantiation of PCI which manages to facilitate control with the help of automatically discovered object-based representations. We provide the convergence proof of OPC perception model update and demonstrate the execution-time optimization ability of OPC in high-dimensional pixel environment. Notably, our experiments show that OPC achieves high quality and consistent perceptual grouping and outperforms several strong baselines in terms of accumulated rewards within the agent's life-cycle. OPC agent can quickly learn the dynamics of new environment without any prior knowledge, imitating the inductive bias acquisition process of humans. For future work, we would like to investigate OPC with more types of inductive biases and test the model performance in a wider variety of environments.
REFERENCES
Masataro Asai and Alex Fukunaga. Classical planning in deep latent space: Bridging the subsymbolic-
symbolic boundary, 2017.
Chris L. Baker, Julian Jara-Ettinger, Rebecca Saxe, and Joshua B. Tenenbaum. Rational quantitative attribution of beliefs, desires and percepts in human mentalizing. Nature Human Behaviour, 1:0064 EP 03 2017.
Peter W Battaglia, Jessica B Hamrick, and Joshua B Tenenbaum. Simulation an engine of physical scene understanding. Proceedings of the National Academy of Sciences, page 201 30637277, 2013.
Christopher Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matthew M Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. ArXiv, abs/1901.1 1390, 2019.
Xi Chen, Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Picter Abbeel. InfoGAN: Interpretable representation learning by information maximizing generative adversarial

---

Under review as a conference paper at ICLR 2021
nets. In D. D. Lee, U. V. Luxburg, I. Guyon, and R. Garnett, editors, Advances In Neural Information Processing Systems 29, pages 2172-2180. Curran Associates, Inc., 2016.
Stanislas Dehaene, Hakwan Lau, and Sid Kouider. What is consciousness, and could machines have it? Science, 35863622)4866499 2017.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society. Series B (methodological), ( pages 1-38, 1977.
Carlos Diuk, Andre Cohen, and Michael L. Littman. An object-oriented representation for efficient reinforcement learning. In ICML, pages 1400247.-4 2008.
Martin Engelcke, Adam R. Kosiorek, Oiwi Parker Jones, and Ingmar Posner. GENESIS: Generative Scene Inference and Sampling of Object-Centric Latent Representations. International Conference on Learning Representations (ICIRR) 2020.
umii-MMasoou Farahmand, Andre Barreto, and Daniel Nikovski. Value-Aware Loss Function for Model-based Reinforcement Learning. volume 54 of Proceedings of Machine Learning Research, pages 1486-1494, Fort Lauderdale, FL, USA, 20-22 Apr 2017. PMLR.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In Robotics and Automation (ICRA), 2017 IEEE International Conference on, pages 2786-2793. IEEE, 2017.
Karl Friston. A theory of cortical responses. Philosophical Transactions of the Royal Society of London B: Biological Sciences, 360(1 (50(1446):815-8666 2005.
Karl Friston. The free-energy principle: a unified brain theory? Nature Reviews Neuroscience, 11:127 EP-, 01 2010.
Marta Garnelo, Kai Arulkumaran, and Murray Shanahan. Towards deep symbolic reinforcement learning. arXiv preprint 160.05560. 2016.
Vikash Goel, Jameson Weng, and Pascal Poupart. Unsupervised video object segmentation for deep reinforcement learning. In NIPS, pages 5688-5699. 2018.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. In NIPS, pages 6691 2017.
Klaus Greff, Raphaël Lopez Kaufmann, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. ICML, pages 2424-2444 243324333222 2019.
Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with model-based acceleration. In ICML, pages 2829-2838, 2016.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NIPS, pages 2451 2451-2463. 2018
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft critic:: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. volume 80 of Proceedings of Muchine Learning Rexearch, pages 1861-1870, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMIRR
Danijar Hafner, Pedro A Ortega, Jimmy Ba, Thomas Parr, Karl Friston, and Nicolas Heess. Action and perception divergence minimization. arXiv preprint arXiv:2009.01991, 2020.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for POMDPs. In ICML, pages 21 17-12226. 10-15 Jul 2018.
Leslie Pack Kaelbling, Michael Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134 1998.
10

---

Under review as a conference paper at ICLR 2021
Ken Kansky, Tom Silver, David A. Mély, Mohamed Eldawy, Miguel Lázaro-Gredille Gredilla, Xinghua Lou, Nimrod Dorfiman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema networks: Zero-shot transfer with generative causal model of intuitive physics. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1809-1818, International Convention Centre, Sydney, Australia, 06-11 1 Aug 2017. PMLR.
Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control graphical model inference problem. Machine learning, 1592)159-1182 2009.
Andrej Karpathy. Reinforcejs: Waterworld
https://cst staanford... wor..eoodddd/ooooooo wwwp//wpppppppppo/o le/ e/
2015. Accessed: 2018-11-20.
Stuart A Kauffman. The origins of order: Self-organization and selection in evolution. OUP USA, 1993.
Ramtin Keramati, Jay Whang, Patrick Cho, and Emma Brunskill. Strategic object oriented reinforce- ment learning. arXiv preprint 1rXiv:1806.00// 00175, 2018.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1412666442.6990, 2014.
Tejas Kulkarni, Ankush Gupta, Catalin lonescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. arXiv preprint 1rXiv:906.11889) 2019.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathemati- cal statistics, 22(1):79-86, 1951.
Brenden M. Lake, Ruslan Salakhutdinov, and Joshua B. Tenenbaum. Human-level concept learning through probabilistic program induction. Science, (5006266):1332-1115 2015.
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40, 2017.
Brenden M Lake, Towards more human-like concept learning in machines: Compositionality, causality, and learning-to- -llaam.. PhD thesis, Massachusetts Institute of Technology, 2014.
Tai Sing Lee and David Mumford. Hierarchical bayesian inference in the visual cortex. JOSA A,
(0(7):1434-1440 2003.
Alex X Lee, Anusha Nagabandi, Picter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. arXiv preprint arXiv:1907.00953 2019.
Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, and Ruslan Salakhutdinov. Efficient exploration via state marginal matching. arXiv preprint arxiv: 2019.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep
visuomotor policies. The Journal of Machine Learning Research, 17(1)113341133,, 2016.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprini 1 0001100050000..0 2018
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint /32.5223 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. ICML, pages 1928-1937 1937, 20-22 Jun 2016.
Brendan O' DDnoghuee Ian Osband, and Catalin Ionescu. Making sense of reinforcement learning and probabilistic inference. In International Conference Learning Representations, 2020.
11

---

Under review as a conference paper at ICLR 2021
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30, pages 6118-6128. Curran Associates, Inc., 2017.
Daniel Alexander Ortega and Pedro Alejandro Braun. Information, utility and bounded rationality. In International Conference Artificial General Intelligence, pages 269-274. Springer, 2011.
Pedro A. Ortega, J. X. Wang, M. Rowland, Tim Genewein, Zeb Kurth-Nelson, Razvan Pascanu, N. Heess, J. Veness, A. Pritzel, P. Sprechmann, Siddhant M. Jayakumar, T. McGrath, K. Miller, Mohammad Gheshlaghi Azar, Ian Osband, Neil C. Rabinowitz, A. György, S. Chiappa, Simon Osindero, Y. Teh, H. V. Hasselt, N. D. Freitas, M. Botvinick, and S. Legg, Meta-learning of sequential strategies. ArXiv, abs/1905.03030, 2019.
Thomas Parr and Karl Friston. Generalised free energy and active inference. Biological cybernetics, 113(5-6)495--1333 2019.
Konrad Rawlik, Mare Toussaint, and Sethu Vijayakumar. Approximate inference and stochastic 2010.
optimal control. arXiv preprint
D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Parallel distributed processing: Explorations in the microstructure of cognition, vol. 1. chapter Learning Internal Representations by Error Propagation, pages 318-362. MIT Press, Cambridge, MA, USA, 1986.
Jürgen Schmidhuber. Learning factorial codes by predictability minimization. Neural Computation, 4(6):863-779, 1992.
Guy Shani, Ronen I. Brafman, and Solomon E. Shimony. Model-based online learning of pomdps. In Machine Learning: ECML 2005, pages 353-364, Berlin, Heidelberg, 2005. Springer Berlin Heidelberg
Elizabeth S Spelke, Karen Breinlinger, Janet Macomber, and Kristen Jacobson. Origins of knowledge. Psychological review, 99(4):605, 1992.
Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning. 3(1):9-44, Aug 1988.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26-31, 2012.
Emanuel Todorov. General duality between optimal control and estimation. 2008 47th IEEE Conference Decision and Control, pages 0286-4292. IEEE, 2008.
Sebastian Tschiatschek, Kai Arllkumaran, Jan Stühmer, and Katja Hofmann. Variational inference for data-efficient model learning in pomdps. arXiv preprint arXiv:1805.0922 2018.
Sjoerd van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expectation maximization: Unsupervised discovery of objects and their interactions. In JCLR, 2018.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang. Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model- based reinforcement learning. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura, editors, Proceedings of the Conference Robot Learning, volume 100 of Proceedings of Machine Learning Research, pages 1439-1456. PMLR, 30 Oct-Ol Nov 2020.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, pages 2746-2754, 2015.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration, arXiv preprint arXiv: 1905.0925, 2019.
12

---

Under review as a conference paper at ICLR 2021
Paul J. Werbos. Generalization of backpropagation with application to recurrent gas market model. Neural Networks, 1(4):339 356, 1988.
CF Jeff Wu, On the convergence properties of the em algorithm. The Annals of statistics, pages 95-103, 1983.
Bo Xin, Haixu Yu, You Qin, Qing Tang, and Zhangqing Zhu. Exploration entropy for reinforcement learning. Mathematical Problems in Engineering, 2020, 2020.
Vinícius Flores Zambaldi, David Raposo, Adam Santoro, Victor Bapst. Yujia Li, Igor Babuschkin, Karl Tuyls, David P. Reichert, Timothy Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals, and Peter W. Battaglia. Relational deep reinforcement learning. CoRR, abs/ foss/880600130,, 2018.
Willard I Zangwill. Nonlinear programming: a unified approach, volume 196. Prentice-Hal Englewood Cliffs, NJ, 1969.
Guangxiang Zhu, Zhiao Huang, and Chongjie Zhang. Object-oriented dynamics predictor. In NIPS, pages 9826-9837. 2018
13

---

Under review as a conference paper at ICLR 2021
A CONVERGENCE OF THE OBJECT-BASED PERCEPTION MODEL UPDATE
Under main assumptions and lemmas as introduced below, we demonstrate the convergence of sequence of L (q"), } generated by the perception update. The proof is presented by showing that the learning process follows the Global Convergence Theorem (Zanwwill 1969..
Assumption 1. Qgo = {0 E Q Co (qt Q) I Cgo (g" .)) is compact for any Co (q" +)) < cc.
Assumption 2. C is continuous in S2 and differentiable in the interior of
The above assumptions lead to the fact that {Cg. ,")) is bounded for any 0" E Z.
Lemma 1. Let Us be the set of stationary points in the interior of Q, then the mapping
arg maxgr 1 Ag (0++1) Eq. closed over Q2 (the complement of Ptg).
Proof. See Wu 1983 A sufficient condition is that Ag (0*+1) is continuous in both 0++ and 0'.
Proposition 1. Let Sg be the set of stationary points the interior of Q then (i.) Ao Us Ce (q" 6) S Cg. (1t 6) and (ii.) v0 E SSASs ( .6) < Lge (4" 6).
Proof. Note that (i.) holds true given the condition. To prove (ii), consider any 0e E Sgss, we have
i aCg (1") x)
Bgg (0(-1)
# 0.
00
00 +1
-
-++1=0
Hence Ag:(( 10(+1) is not maximized at ac+1 A', Given the perception update described by Eq.
we therefore have Agr (0(+1)> Ao (0'), which implies cy (4th 6) Cgg (4" 6).
Theorem 1. Let (0') be sequence generated by the mapping from Eq. 4 Ug be the set of stationary points in the interior of Q. If Afssmmpttosss 7 & 2 Lemma [ and Proposition 7 are met, then all the limit points of {0'} are stationary points (local minima) and L (q .6) converges monotonically to J = Lg. (q 6) for some stationary point 0* E Ugg.
Proof. Suppose that 0* is a limit point of the sequence {0'). Given Assumptions I & and Proposition i)), have that the sequence {Ø'} are contained in compact set Ωx C. Ω. Thus, there a subsequence (0') (e)tez of { such that 0' + 0 as l + co and E L.
We first show that Cyj (q" ②) Lw. (q" Φ) as t + cc. Given L is continuous in Ω (Assumption
we have C (q" ②) + Lg (q") Φ) as Do and I E L, which means
Ye 0, ((E) E I. s.t. 2 I(e). E I,, fg ( .)) L. (q" () <.
Given Proposition and Eq. L is therefore monotonically decreasing the sequence (0*) )- which gives
VI, Co (g 100,, Φ) Cg-(1" .d) 0.
Given Eq. 7.. for any 12 I(e), we have
E (4" .)) LT (4") +) CT. Φ) Car (4") + (4 .6) L (a" < <e.
(9)
5o
Given Eq. 8 and Eq. 9 therefore have Cg. (g "0 Φ) (g" Φ) as - We then prove
that the limit point 0r is stationary point. Suppose 0* is not stationary point, i... 0" E 2,sss we consider the sub-sequence {0++1) HEL, which are also contained in the compact set 2g. Thus, there is subsequence (0 JUEL of (0'+1) such that 0¹ + 0² as l'
(0'+))
and 1' E L', yielding
Cy 12 (q" ,6) + Lg (4th 7", Φ) as l' + co and l' L', which gives
Cy (g" x) lim S 1 (9 .))
lim tEN
2o (q" .)) L (qth ,)).
(10)
ret
On the other hand, since the mapping from Eq. we therefore have 8" E arg maxg++1 Ag- (e yielding C (qt Φ) C (q" Φ) (Proposition
4 is closed over Ω Uggg (Lemmal and 0* E (2ss
Ii), which contradicts Eq. 10
14

---

Under review as a conference paper at ICLR 2021
B EXPERIMENT DETAILS
OPC In all experiments we trained the perception model using ADAM Kingma and Ba (2014) with default parameters and batch size of 32. Each input consists of sequence of binary 84 84 images containing two poison objects (two circles) and one food object (a rectangle) that start in random positions and move within the image for tro steps. These frames were thresholded 0.0001 obtain binary images and added with bit-flip noise (p = 0.2). We used convolutional encoder- architecture inspired by recent GANs Chen al. (2016) with recurrent neural network as bottleneck, where the encoder used the same network architecture from Mnih el al. (2013) as
1. 8 x conv. 16 ELU. stride 4. layer norm 2. 4 x conv. 32 ELU. stride 2. layer norm 3. fully connected. 256 ELU, layer norm
4. recurrent. 250 Sigmoid. layer norm on the output
5. fully connected. 256 RELU. layer norm
6. fully connected. 10 x 10 x 32 RELU. layer norm
7. 4 x reshape neireot-n conv. 16 RELU. layer norm
8. 8 reshape nearest-neighborr, conv. Sigmoid
We used the Advantage Actor-Critic (A2C) Mnih al (2016) with an MLP policy the decision making module of OPC. The MLP policy added a 512-unit fully connected layer with rectifier nonlinearity after layer 4 of the perception model. The decision making module had two set of outputs: 1) softmax output with one entry per action representing the probability of selecting the action, and 2) a single linear output representing the value function. The decision making module was trained using RMSProp Tieleman and Hinton (2012) with a learning rate of 7e 4, reward discount factor 7 = 0.99, an RMSProp decay factor of 0.99, and performed updates after every actions.
A2C We used the same convolutional architecture as the encoder of the perception model of OPC (layer 1 to 3), followed by fully connected layer with 512 hidden units followed by rectifier nonlinearity. The A2C was trained using the same setting as the decision making module of OPC.
WM-A2C We used the same setting as Ha and Schmidhuber 2018) to separately train the model and the M model. The experience was generated off-line by random policy operating in the Pixel Waterworld environment.
concatenated the output of the model and the M model as the A2C input, and trained A2C using the same setting introduced above.
15

---

Under review as a conference paper at ICLR 2021
C DETAILS OF DERIVATION
C.I.
DERIVATION OF EQ.
logp(o²*
x
log
Σ
P(o 18 xx (a21(a5))
=log Σ
p(o 21 x2 )p(s") x² ,021)
=log
E
II
I . "a")aa""'>sell))
_ p(s x .a (a (2
22 2t ga not
= llog
Σ
I a") II
pl (e))ppp>>>>" s² )²)))((²1 a22
sbl 2 1 nmd
5=1
= =log
Σ
Troo"|' a²))
p(a
II p(a="1*" )="))(>>"11>>> )(as 22 (0 kmt+1
x21 22,a?? mx21 nz!
1=1
= llg
II p(o"|s" a")
p( (s')
II p(= 1s p(s )p(s') II
p(s 1 (m)) II (t)
a21 aa 2a
==1
k=t+1
m=1
log
EII p(ar (s pss") II p(s amm
Σ a 2 m=21
Ilo'"' a=) II p(l |
++ )(( l aaal dl(aaaaa
aSt jm1
Σ 9(8 1 , IT, p( (s))p(")) II m
15 a
log
(185) (n axt)
C.2
DERIVATION OF EQ.
4
II; p(x0 | 1p(s'
1I 1-1 p( am a)
g++1 arg max q" log 0'+1
p(s!)
1==
1-1 ((5 12009 aS9)
( )p(x' ss)
|1 p(x sit 1 as)
j=1
arg max q" log g++1
p(s*)
9= a(50+1 100 aso)
p( 18
p(xt+1,
gi+l 8) a²)
= -arg max of log 0:11
drop terms which are constant with respect to 0 g
1==1 1o+1
[log (= 21+1 26 (2tt1
= arg max
E 1oy
0'+1
arg max Agr (0*+1)
16

---

Under review as a conference paper at ICLR 2021
C.3
DERIVATION OF EQ.
(5)
0A 0Ag-11(0022) of v +
00%11
016% 10t 1
80,
(0
0A (0 D D.*
0v1,1
D 1 30
ggg (0c42)
Ousti i,
00
Ouiti
i
Σ Po(2 | 2/ a²) Pog ( 1 1s" .0)
00%+1
i
D
Ouit!
Σ po(2)
+ 1o 1/,, a') log P i (1/ )p(2 ++ 18 1s (s)
i.k
++I
00
0xb
Σ Lk
r+1 1|x wi a² log Por (a) 12 "=(()p(((]))11 |p(2|*1|s",00 "s",
i.
Po
Jult1 i,k
ae
Po ( t-1
1
1|x, ++1 wb
logPacu (a + |2 )(( 11" |
iik
a00
D
a
Σ Po
++1 = 1
1| ++1 w/.
logp se
(af ++ | + log p(2 x 11 aaa)
00
Σ
A+1
1|
w/, a²)
log PogPaeti (z ++ | 1
++
00tt1
1,k
Σ P
2/. 1|1
Dult+1
log
V2TT2
exp
2
1
(x
(+1
Vikk
00
0
(x
w/
11+112 )
du 1+
Σ Po
= 1 1| 1+1 3/.
log o
log 2n
0rie41
2
0²
00
1
2xx
1 ) (-1)
Ou
P 1-1 = 1
1|
4bf a²)
Fi,k
2
σ²
00.+1
Eg.
Σ
1
i,k
x
[+1
066-1
i,
1
a²
00 +1
17


Reviews:
[{"review_id": "w9Y6yh2M45W", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI’m quite conflicted about this paper. \nOn the one hand, I find the problem that they’re trying to tackle (i.e. object-centric RL) to be interesting, important and rather current. Their approach is also interesting, theoretically guided and sound. \nHowever on the other hand I found the choice of introducing yet another framework (which isn’t terribly useful or novel either, see below) to be counter-productive for the paper; the derivations to be slightly strenuous (and one can argue already quite known in the field and expressed in simpler fashion elsewhere, see [2]); and perhaps of unclear novelty (i.e. what is novel in the perception aspect compared to NEM?). Finally, the results aren’t extremely convincing and they do not compare to a model which is extremely related to their proposal (OP3 [3]).\n\nComments/questions:\n      \n1. I do not feel like the way the abstract and introduction argue for a general framework “Joint Perception and Control as Inference” is productive. As mentioned during the introduction, this isn’t a particularly novel framework (considering RL as inference, all the work done in temporal generative model for RL, or the rather new formulation by Hafner 2020 [4]) and it doesn’t appear like it provides a very interesting fit to the object-centric RL model proposed. \n   1. The derivations of section 3.2 are heavy, introduce a lot of notation, but ultimately aren’t that complicated. \nHow do they differ from the similar derivations performed by Hafner et al, 2019’s Appendix F (see [2])?\n   2. I understand that performing them for POMDP would be valuable, but then it would be more useful to address what is missing from existing methods and just show the required changes.\n   3. I would argue that the paper would be stronger if you only expressed the object-centric RL derivation, and put it into context with NEM and OP3, to express what is different about it.\n   4. I found the E and M steps derivations interesting, but they are effectively the same as shown in NEM [1]? Did I miss something? It would be more helpful to contrast them, instead of rederiving them.\n   5. Finally, the novel aspect of the control objective isn’t actually used in practice, as a simple TD error method is used to train the policy, which others have done without taking a “RL as inference” stance.  \n2. The overall idea of the model is extremely close to what has been done in OP3 [3], where they combine sequential IODINE with policy heads to solve tasks while simultaneously learning an iterative inference generative model.\n   1. This work should be cited and thoroughly discussed. It also should be a baseline in the best of situations.       \n3. Does the proposed model also inherit the limitations of NEM? \n   1. For example, are you limited to black and white observations? \n   2. How computationally expensive is it to represent $\\frac{\\partial\\psi_k}{\\partial \\theta_k}$ with a neural network?      \n4. Figure 3 is relatively hard to interpret and quite noisy, which doesn’t make the results extremely promising?\n   1. Starting with accumulated rewards makes it confusing, as one could be made to believe that the agents never successfully learn to perform the task and avoid poisons.\n   2. “Period costs” should really be called “Episodic returns”, as this is really more standard.\n   3. The variance is very high and would benefit from extra seeds.      \n5. The performance of the learnt generative model in Figure 4 does not seem particularly great, especially given the simplicity of the domain considered? The segmentations are fine, but the reconstructions seem quite too blurry compared to current SOTA methods.      \n6. Given the similarity of the environments, COBRA [5] might have been a better baseline than the relatively unrelated VAE / World Model baseline selected.\n   1. I would also argue that COBRA should be cited in the “Model-based Deep Reinforcement Learning” as a closely related object-centric model-based RL method, it is more relevant than the World Model paper.\n\nSo overall, I’d recommend rewriting the paper to be more explicit about its contributions and differences to existing works, as this would make it much easier to follow.\n\n* [1] Greff et al, 2017, https://arxiv.org/abs/1708.03498 \n* [2] Hafner et al, 2019, https://arxiv.org/abs/1811.04551 \n* [3] Veerapaneni et al, 2019, https://arxiv.org/abs/1910.12827 \n* [4] Hafner et al, 2020, https://arxiv.org/abs/2009.01791 \n* [5] Watters et al, 2019, https://arxiv.org/abs/1905.09275 ", "extra_scores": {}, "extra_text": {}}, {"review_id": "Slp-a1OvBfx", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n-------------\nThe main contribution of the paper is the incorporation of the object-based inference module (Greff et al., 2017; van Steenkiste et al., 2018) into the control-as-inference framework under partial observability. The method is evaluated on a version of Waterworld environment.\n\nDecision\n-----------\nI think the paper is borderline, and I tend towards rejecting it in its current form. But I am open to increase my rating if the author address my concerns. Below are the main reasons for my decision.\n\n1. The separation between what is novel and what is already known is not clear from the paper. For example, it is claimed that the paper introduces a novel framework, however there is a number of papers working with the same framework but using different approximations, e.g., [1,2,3]. There is also a close relation to active inference papers, e.g., [4,5]. In my opinion, it should be stated more clearly in the paper that the main contribution is the application of the object-based perception model within the control-as-inference framework, and not an introduction of a novel framework.\n2. Following from the previous point, since the perception module is the main innovation, its evaluation seems insufficient. It is tested only on one environment. Would it generalize to other environments?\n\nReferences\n---------------\n[1] Ortega, P. A., Wang, J. X., Rowland, M., Genewein, T., Kurth-Nelson, Z., Pascanu, R., ... & Jayakumar, S. M. (2019). Meta-learning of sequential strategies. arXiv preprint arXiv:1905.03030.\n[2] Tschiatschek, S., Arulkumaran, K., Stühmer, J., & Hofmann, K. (2018). Variational inference for data-efficient model learning in pomdps. arXiv preprint arXiv:1805.09281.\n[3] O'Donoghue, B., Osband, I., & Ionescu, C. (2020). Making sense of reinforcement learning and probabilistic inference. arXiv preprint arXiv:2001.00805.\n[4] Tschantz, A., Millidge, B., Seth, A. K., & Buckley, C. L. (2020). Reinforcement Learning through Active Inference. arXiv preprint arXiv:2002.12636.\n[5] Friston, K., Da Costa, L., Hafner, D., Hesp, C., & Parr, T. (2020). Sophisticated Inference. arXiv preprint arXiv:2006.04120.", "extra_scores": {}, "extra_text": {}}, {"review_id": "CtBDtNP0bA", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper features two complementary contributions:\n1.  The perception and control as inference (PCI) framework, which describes the graphical model of a POMDP with auxiliary optimality variables, allowing for the derivation of an objective for optimizing both a perception model and policy jointly to maximize $\\log p(\\mathcal{O}, \\mathbf{x} \\mid \\mathbf{a})$. As the name suggests, this is similar to the control as inference formulation with a focus on partial observability and the addition of a perception model $q^w$.\n2. Object-based perception control (OPC), which is an instance of PCI using an object-factorized model.\n\nIncorporating a perception model into the the control-as-inference derivation to yield a single unified objective for multiple components in a pipeline is an interesting idea and complements other work on incorporating reward structure into model-learning nicely. (A discussion of some of these, such as [[Farahmand et al, 2017]](http://proceedings.mlr.press/v54/farahmand17a.html) or [[Oh et al, 2017]](https://arxiv.org/abs/1707.03497) could serve to better highlight the difference between PCI and existing hybrid approaches, but these references are not \"missing\" in the sense that their absence is an issue.) \n\nHowever, the experimental evaluation does not really do the generality of this idea justice. It seems that this sort of approach would be favored only when the perception problem is too difficult for a standard maximum likelihood objective, as the joint objective would (presumably) bias the perception model to be useful for control. The main evaluation of the perception model, though, only shows four primitive white shapes on a black background being assigned labels. Unless there is underlying complexity not apparent in the image, this sort of input appears to be in line with what binary segmentation algorithms could work with, and likely would require no learning. The implicit argument made is that optimizing the joint objective causes these labels to be semantically-informed, but this is not supported by an experiment; the segmentation results in figure 4 are also what you would get by segmenting based on boundaries. \n\nSomewhat better support could come from showing what labels are assigned when two different object types are contiguous with each other (would food and poison be assigned the same label, or does the model separate them?), but in general it is difficult to support the claimed benefits with this particular environment. This is likely one of the reasons for the result described in the figure 3(c-d) ablations: when the environment is too simple (this time in terms of dynamics, likely, and not perceptual complexity), many of the design choices can become less consequential.\n\nAside from experiments, the presentation of the core ideas could also be improved somewhat. Section 3.2 in particular is difficult to parse, and presented in such a way that many readers will likely just skip over most of it. It is certainly a good idea to work out an approach in its full generality, but some of the components taking up space in the presented derivation are immediately cut in the experiments section (like the $D_\\text{KL}(\\pi~ || ~p(a^t \\mid a^{<t})$ term), so you might as well present a simplified version of the result that focuses on the method actually used. \n\n\n**Miscellaneous questions**\n1. What is the role of $p(a \\mid a^{<t})$? Is this to account for a recursive policy?\n2. Why do you think the world models baseline never appreciably improves upon its random initialization in figure 2b? I could imagine that the joint objective would cause OPC to perform a bit better than world models at convergence, but I am surprised that the perception is apparently so difficult that a separate modeling objective never gets off the ground at all.\n\n\n**Minor**\n1. Section 3.4: \"resulting a\" $\\rightarrow$ \"resulting in a\"\n2. There are a number of claims about human cognition that are not supported. For example:\n\n> OPC agent can quickly learn the dynamics of a new environment without any prior knowledge, imitating the inductive bias acquisition process of humans.\n\n> Moreover, in order to mimic a human’s spontaneous acquisition of inductive biases 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).\n\nIt is hard to say whether OPC mimics humans' knowledge acquisition without any sort of evidence.\n\n\n**Summary**\n\nThis generalization of the control as inference formulation could be influential given a more didactic presentation of the PCI derivation and an evaluation better suited to its claimed strengths, but it does not seem quite ready for publication yet. I would encourage the authors to try to test their approach on an environment that is perceptually complicated enough to motivate a non-standard perceptual modeling objective.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "2BXK4NYh3zy", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a framework for joint perception and control as inference (PCI) to combine perception and control for the case of POMDPs. The authors particularly focus on the case of hidden perceptual states linked to small image observations, which are composed of pixels belonging to up to exactly one object each. Their main proposal is denoted as OPC, which stands for object based perception and control, which serves the purpose of automatically discovering objects from pixels while controlling the system. \n\nIn order to posit the OPC model, the authors adopt one of the various recent unsupervised grouping models which assign one of K latent object identities to each pixel and hence group/segment images into K objects.  The rest of the model is a common POMDP setting set up to perform control as inference via usage of latent optimality variables per timepoint.\n\nOne of the major contributions in the paper is a well worked-out joint inference derivation, which performs amortized inference in this shared model utilizing an RNN.\n\nExperimentally, the authors verify their joint model on a waterworld environment, which consists of simple object shapes with semantics. They compare their model against baselines such as A2C with convolutional networks as feature extractors , A2C with the World Model (Ha et al) as a feature extractor (also based on convolutional networks) and a random policy.\nThey unsurprisingly beat these baselines.\n\nCriticisms:\nMy main problem with the paper is that the authors do not show that the joint model does anything. If they use one of these perceptual grouping models as feature extractors without joint inference, would A2C work well? Is it that feature layer that buys performance or a deeper interactions through the joint model they propose such that the actions resolve uncertainty about objectness in a nontrivial way they would not if used in a pipeline that doe snot perform joint inference. \nThis is not demonstrated convincingly and is a huge missed opportunity.\nI am left without a firm understanding if the highlighted contribution of joint inference benefits this model in any way.\nIn addition, all the baselines have 'monolithic' feature extractors , while pretty clearly object based feature extractors appear useful.\nI would also want to see a baseline where objects are known and we understand what the rest of their framework would do. \nIn general I find that even though the experimental framework is extremely toy, the authors have left a lot on the table empirically to demonstrate the utility of their method.\nIs the lengthy derivation of updates for the model worth doing if we don't collect this evidence?\nWhat is the time complexity in terms of EM steps that are needed?\nI am less concerned about the toy nature of the chosen experiment, as it would suffice to prove the concept the authors are after, if used sufficiently for empirical rigor.\n\nA secondary criticism is the lack of discussion and comparison to systems such as the Schema networks (Schema Networks: Zero-shot Transfer with a Generative Causal Model of Intuitive Physics by Kansky et al) which has long ago proposed object based planning modules and has not used object based models in the loop (yet).\nThere may be more such papers in the related literature which are on that thread of combining RL with object based representations that have not found their way in the discussions here and clearly should.\n\nOverall:\nI consider this paper interesting and agree with the authors that the proposal of joint models is appealing and worthwhile investigating, but I would require much stronger evaluation to show that the joint aspect of OPC is driving the performance here, instead of just the object-based featurization. In its current form the paper feels interesting, but unfinished.\nThe lapses in scholarship are also important to fix in order to fairly slot the paper correctly into the wider field.\nHowever, this paper is promising and I look forward to seeing it mature.\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces an object perception and control method for RL, derived from a control-as-inference formulation within a POMDP.  The paper provides a theoretical derivation and experiments where the proposed joint-inference approach outperforms baselines.

The discussion focussed on understanding the paper's contribution relative to prior work. The reviewers highlighted the similarities with earlier systems (R1, R2, R4), the unclear benefits of joint inference over independently trained modules in the experiments (R3), and the lack of clarity of the presentation (R1, R2, R3).  The authors responded to some of these criticisms, bolstering the paper with additional experiments to show the benefits of joint inference and increasing the discussion of related work.  The reviewers examined the revisions and rebuttal and found the paper still did not resolve all their original concerns.  Two limitations mentioned in the final phase of the discussion were the use of a single environment to evaluate the general framework, and continuing doubts on the contribution of joint inference mechanism to the measured performance.

Four knowledgeable reviewers indicate reject as their concerns were not adequately resolved.  The paper is therefore rejected.
Decision: Reject