Reinforcement learning for instance segmentation with high-level priors

Instance segmentation is an important computer vision problem which remains1 challenging despite impressive recent advances due to deep learning-based meth-2 ods . Given sufficient training data , fully supervised methods can yield excellent3 performance , but annotation of ground-truth data remains a major bottleneck , espe-4 cially for biomedical applications where it has to be performed by domain experts.5 The amount of labels required can be drastically reduced by using rules derived6 from prior knowledge to guide the segmentation . However , these rules are in7 general not differentiable and thus can not be used with existing methods . Here,8 we relax this requirement by using stateless actor critic reinforcement learning,9 which enables non-differentiable rewards . We formulate the instance segmentation10 problem as graph partitioning and the actor critic predicts the edge weights driven11 by the rewards , which are based on the conformity of segmented instances to12 high-level priors on object shape , position or size . The experiments on toy and real13 datasets demonstrate that we can achieve excellent performance without any direct14 supervision based only on a rich set of priors.15 1 Introduction16 Instance segmentation is the task of segmenting all objects in an image and assigning each of them17 a different label . It forms the necessary first step to the analysis of individual objects in a scene18 and is thus of paramount importance in many practical applications of computer vision . Over the19 recent years , fully supervised instance segmentation methods have made tremendous progress both20 in natural image applications and in scientific imaging , achieving excellent segmentations for very21 difficult tasks [ 1 , 2 ] .22 A large corpus of training images is hard to avoid when the segmentation method needs to take23 into account the full variability of the natural world . However , in many practical segmentation24 tasks the appearance of the objects can be expected to conform to certain rules which are known a25 priori . Examples include surveillance , industrial quality control and especially medical and biological26 imaging applications where full exploitation of such prior knowledge is particularly important as the27 training data is sparse and difficult to acquire : pixelwise annotation of the necessary instance-level28 groundtruth for a microscopy experiment can take weeks or even months of expert time . The use of29 shape priors has a strong history in this domain [ 3 , 4 ] , but the most powerful learned shape models30 still require groundtruth [ 5 ] and generic shapes are hard to combine with the CNN losses and other,31 non-shape , priors . For many high-level priors it has already been demonstrated that integration of32 the prior directly into the CNN loss can lead to superior segmentations while significantly reducing33 the necessary amounts of training data [ 6 ] . However , the requirement of formulating the prior as34 a differentiable function poses a severe limitation on the kinds of high-level knowledge that can35 be exploited with such an approach . The aim of our contribution is to address this limitation and36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . establish a framework in which a rich set of non-differentiable rules and expectations can be used to37 steer the network training.38 To circumvent the requirement of a differentiable loss function , we turn to the reinforcement learning39 paradigm , where the rewards can be computed from a non-differentiable cost function . We base40 our framework on a stateless actor-critic setup [ 7 ] , providing one of the first practical applications41 of this important theoretical construct . In more detail , we solve the instance segmentation problem42 as agglomeration of image superpixels , with the agent predicting the weights of the edges in the43 superpixel region adjacency graph . Based on the predicted weights , the segmentation is obtained44 through ( non-differentiable ) graph partitioning and the segmented objects are then evaluated by the45 critic , which learns to approximate the rewards based on the object- and image-level reasoning ( see46 Fig . 1 ) .47 The main contributions of this work can be summarized as follows : ( i ) we formulate instance segmen-48 tation as a RL problem based on a stateless actor-critic setup , encapsulating the non-differentiable step49 of instance extraction into the environment and thus achieving end-to-end learning ; ( ii ) we exploit50 prior knowledge on instance appearance and morphology by tying the rewards to the conformity of51 the predicted objects to pre-defined rules and learning to approximate the ( non-differentiable ) reward52 function with the critic ; ( iii ) we introduce a strategy for spatial decomposition of rewards based on53 fixed-sized subgraphs to enable localized supervision from combinations of object- and image-level54 rules . ( iv ) we demonstrate the feasibility of our approach on synthetic and real images and show55 an application to an important segmentation task in developmental biology , where our framework56 delivers an excellent segmentation with no supervision other than high-level rules.57 2 Related work58 Reinforcement learning has so far not found significant adoption in the segmentation domain . The59 closest to our work are two methods in which RL has been introduced to learn a sequence of60 segmentation decision steps as a Markov Decision Process . In the actor critic framework of [ 8 ] , the61 actor recurrently predicts one instance mask at a time based on the gradient provided by the critic.62 The training needs fully segmented images as supervision and the overall system , including an LSTM63 sub-network between the encoder and the decoder , is fairly complex . In [ 9 ] , the individual decision64 steps correspond to merges of clusters while their sequence defines a hierarchical agglomeration65 process on a superpixel graph . The reward function is based on Rand index and thus not differentiable,66 but the overall framework requires full ( super ) pixelwise supervision for training.67 Reward decomposition was introduced for multi agent RL by [ 10 ] where a global reward is decom-68 posed into a per agent reward . [ 11 ] proves that a stateless RL setup with decomposed rewards requires69 far less training samples than a RL setup with a global reward . In [ 12 ] reward decomposition is70 applied both temporally and spatially for zero-shot inference on unseen environments by training on71 locally selected samples to learn the underlying physics of the environment.72 The restriction to differentiable losses is present in all application domains of deep learning . Common73 ways to address it are usually based on a soft relaxation of the loss that can be differentiated . The74 relaxation can be designed specifically for the loss , such as , for example , Area-under-Curve [ 13 ] for75 classification or Jaccard Index [ 14 ] for semantic segmentation . These approaches are not directly76 applicable to our use case as we aim to enable the use of a variety of object- and image-level priors77 which can easily be combined without handcrafting an approximate loss for each case . More generally,78 but still for a concrete task loss , Direct Loss Minimization has been proposed for CNN training in79 [ 15 ] . For semi-supervised learning of a classification or ranking task , Discriminative Adversarial80 Networks have been proposed as a means to learn an approximation to the loss [ 16 ] . Most generally,81 Grabocka et al . in [ 17 ] propose to train a surrogate neural network which will serve as a smooth82 approximation of the true loss . In our setup , the critic can informally be viewed as a surrogate network83 as it learns to approximate the priors through the rewards by Q-learning.84 Incorporation of rules and priors is particularly important in biomedical imaging applications , where85 such knowledge can be exploited to augment or even substitute scarce groundtruth annotations.86 For example , the shape prior is explicitly encoded in the popular nuclear [ 18 ] and cellular [ 19 ] 87 segmentation algorithms based on spatial embedding learning . Learned non-linear representations88 of the shape are used in [ 5 ] , while in [ 20 ] the loss for object boundary prediction is made topology-89 aware . Domain-specific priors can also be exploited in post-processing by graph partitioning [ 21 ] .90 Interestingly , the energy minimization procedure underlying the graph partitioning can also be91 incorporated into the learning step [ 22 , 23 ] .92 3 Methods93 The task of instance segmentation can be formalized as transforming an image x into a labeling y,94 where y maps each pixel to a label value . An instance corresponds to the maximal set of pixels with95 the same label value . Typically , the instance segmentation problem is solved via supervised learning,96 i.e . using a training set with ground-truth labels ŷ . Note that y is invariant under the permutation97 of label values . In general , it is difficult to formulate instance segmentation in a fully differentiable98 manner . Most approaches first predict a `` soft '' representation with a CNN , e.g . affinities [ 1 , 24 , 25 ] ,99 boundaries [ 26 , 27 ] or embeddings [ 28 , 29 ] and apply non-differentiable post-processing , such as100 agglomeration [ 27 , 30 ] , clustering [ 31 , 32 ] or partitioning [ 33 ] , to obtain the instance segmentation.101 Alternatively , proposal-based methods predict a bounding-box per instance and then predict the102 instance mask for each bounding-box [ 34 ] . Furthermore , the common evaluation metrics for instance103 segmentation [ 35 , 36 ] are also not differentiable.104 Our main motivation to explore RL for the instance segmentation task is to circumvent the restriction105 to differentiable losses and - regardless of the loss - to make the whole pipeline differentiable end-to-106 end even in presence of non-differentiable steps which transform pixelwise CNN predictions into107 individual instances.108 We formulate the instance segmentation problem using a region adjacency graph G = ( V , E ) ,109 where the nodes V correspond to superpixels ( homogeneous clusters of pixels ) and the edges E110 connect nodes which belong to spatially adjacent superpixels . Given edge weights W , an instance111 segmentation can be obtained by partitioning the graph , here using an approximate multicut solver112 [ 37 ] . Together , the image data , superpixels , graph and the graph partitioning make up the environment113 E of our RL setup . Based on the state s of E , the agent A predicts actions a , which are used to114 compute the partitioning . The reward r is then computed based on this partitioning . Our agent A is a115 stateless actor-critic [ 38 ] , represented by two graph neural networks ( GNN ) [ 39 ] . The actor predicts116 the actions a based on the graph and its node features F . The node ( superpixel ) features are computed117 by pooling together the corrresponding pixel features based on the raw image data.118 Here , we make use of two different setups : Method 1 , where the per-pixel features are computed119 based on the image data with the feature extractor being part of the agent A and Method 2 where the120 feature extractor is part of the environment E . The feature extractor is trained end-to-end in Method121 1 , whereas it is fixed and thus needs to be pre-trained in Method 2 . We use a U-Net [ 40 ] as feature122 extractor and can use hand-crafted features in addition to the learned features . More details about123 the pre- training can be found in the Appendix . The agent - environment interaction for Method 1 is124 depicted in Figure 1 . For Method 2 we refer to the Appendix.125 Importantly , this setup enables us to use both a non-differentiable instance segmentation step and126 reward function , by encapsulation of the “ pixels to instances ” step in the environment and learning a127 policy based on the rewards with a stateless actor critic.128 3.1 Stateless Reinforcement Learning Setup129 Unlike most RL settings [ 41 ] , our approach does not require an explicitly time dependent state : the130 actions returned by the agent correspond to the real-valued edge weights in [ 0 , 1 ] , which are used to131 compute the graph partitioning . Any state can be reached by a single step from the initial state and132 there exists no time dependency in the state transition . Unlike [ 9 ] , we predict all edge values at once133 which allows us to avoid the iterative strategy of [ 8 ] and deliver and evaluate a complete segmentation134 in every step . We implement a stateless actor critic formulation with episodes of length 1.135 To the best of our knowledge , stateless RL was introduced in [ 7 ] to study the connection between136 generative adversarial networks and actor critics and our method is one of the first practical applica-137 tions of this concept . Here , the agent consists of an actor , which predicts the actions a and a critic,138 which predicts the action value Q ( expected future discounted reward ) given the actions . The stateless139 approach simplifies the action value function : the action value has to estimate the reward for a single140 step instead of estimating the expected sum of discounted future rewards for many steps . We have141 explored a multi-step setup as well , but found that it yields inferior results for our application ; details142 can be found in the Appendix . As described in detail in 3.2 , we compute localized sub-graph rewards143 instead of relying on a single global reward.144 The actor corresponds to a single GNN , which predicts the mean and variance of a normal distribution145 for each edge . The actions a are determined by sampling from this distribution and applying a146 sigmoid to the result to obtain continuous edge weights in the value range [ 0 , 1 ] . The GNN takes the147 state s = ( G , F ) as input arguments and its graph convolution for the ith node is defined as in [ 39 ] :148 fi = γπ fi , 1|N ( i ) | ∑ j∈N ( i ) φπ ( fi , fj ) ( 1 ) where γπ as well as φπ are MLPs , ( · , · ) is the concatenation of vectors andN ( i ) is the set of neighbors149 of node i . The gradient of the loss for the actor is given by:150 ∇θLactor = ∇θ 1 |SG| ∑ sg∈G α∑ â∈sg log ( πθ ( â|s ) ) −Qsg ( s , a ) ( 2 ) This loss gradient is derived following [ 38 ] . We adapt it to the sub-graph reward structure by151 calculating the joint action probability of the policy πθ over each sub-graph sg in the set of all152 sub-graphs SG . Using this loss to optimize the policy parameters θ minimizes the Kullback-Leibler153 divergence between the Gibbs distribution of action values for each sub-graph Qsg ( s , a ) and the154 policy with respect to the parameters θ of the policy . α is a trainable temperature parameter which is155 optimized following the method introduced by [ 38 ] .156 157 The critic predicts the action value Qsg for each sub-graph sg ∈ SG . It consists of a GNN Qsg ( s , a ) 158 that takes the state s = ( G , F ) as well as the actions a predicted by the actor as input and predicts a159 feature vector for each edge . The graph convolution from Equation 2 is slightly modified:160 fi = γQ fi , 1|N ( i ) | ∑ j∈N ( i ) φQ ( fi , fj , a ( i , j ) ) ( 3 ) again γQ and φQ are MLPs . Based on these edge features Qsg is predicted for each sub-graph via an161 MLP . Here , we use a set of subgraph sizes ( typically , 6 , 12 , 32 , 128 ) to generate a supervison signal162 for different neighborhood scales . A given MLP is only valid for a fixed graph size , so we employ a163 different MLP for each size . The loss for the critic is given by:164 Lcritic = 1 |SG| ∑ sg∈G 1 2 ( Qδsg ( s , a ) − r ) 2 ( 4 ) Minimizing this loss with respect to the action value function ’ s parameters δ minimizes the difference165 between the expected reward and action values Qδsg ( s , a ) .166 3.2 Localized Supervision Signals167 The RL paradigm is to provide a global reward for a given state transition [ 41 ] . However , we find168 that for our application it is possible and desirable to instead provide several more localized rewards169 per state transition : Given a large action space with a policy represented by a complex multivariate170 probability distribution , it is beneficial to learn from rewards for the specific actions rather than from171 a scalar global reward for the union of all actions . Of course then requirement arises that the union of172 local rewards must resemble to the global reward . E.g . the optimal policy is the same for local as for173 the global reward.174 Our actor critic setup ( Section 3.1 ) expects rewards per sub-graph . A good set of sub-graphs should175 fulfill the following requirements : each sub-graph should be connected so that the information176 presented to the MLP computing the action value for this sub-graph is correlated . The size of177 the sub-graphs , given by the number of edges , should be a parameter and all sub-graphs should178 be extracted with exactly that size to serve as valid input for one of the MLPs . The union of all179 sub-graphs should cover the complete graph so that each edge contributes to at least one action180 value Qsg . The sub-graphs should overlap to provide a smooth sum of action values . We introduce181 Algorithm 1 to extract such a set of sub-graphs ( see Appendix ) . Figure 2 shows the sub-graphs for a182 small example graph.183 While some of the rewards used in our experiments can be directly defined for the sub-graphs , most184 are instead defined per object ( see Appendix for details on reward design ) . We use the following185 general procedure to map object-level rewards to sub-graphs : first assign to each superpixel the186 reward of its corresponding object , then determine the reward per edge as the maximum value of its187 two incident superpixels ’ rewards and average the edge rewards to obtain the reward per sub-graph.188 Here , we use the maximum because high object scores indicate that all actions contributing to the189 respective object should get a high reward . However , for low object scores it is not possible to localize190 the specific action responsible for the low score . Hence , by taking the maximum we assign the191 higher score to edges whose incident superpixels belong to different objects , because they probably192 correspond to a correct split . Note that the uncertainty in the assignment of low rewards can lead to193 a noisy reward signal , but the averaging of the edge rewards over the sub-graphs and the overlaps194 between the sub-graphs smooth and partially denoise the rewards . We have also explored a different195 actor critic setup that can use object level rewards directly , eliminating the need for the sub-graph196 extraction and mapping . However , this approach yields inferior results , see the Appendix for details.197 4 Experiments198 The agent of our setup acts on the superpixel graph and thus depends on the features assigned to the199 nodes of the graph . We introduced two variants of our algorithm : in the base variant ( Method 1 ) 200 we start from random features and make them part of the agent , allowing them to change through201 back-propagation ( Fig . 1 ) . In contrast , Method 2 acts on predefined features which are provided202 as part of the environment and are computed before training , e.g . through unsupervised clustering.203 A very accurate clustering in the features produces an easy problem for the agent to solve where204 even a global reward for all actions might be sufficient . However , in a real-world setting with no205 supervision , the noisier the features become the more local the reward has to be . We evaluate Method206 2 on synthetic data where self-supervised pretraining can deliver noisy , but meaningful node features.207 Our full setup with Method 1 is evaluated on a dataset from a light microscopy experiment , where208 highly regular object shapes are to be expected , but no good feature pre-training is possible.209 To transform the edge weight predictions of the agent into an instance segmentation we use the210 Multicut [ 42 ] algorithm . Here , other options are also possible such as hierarchical clustering used in211 [ 9 ] , but we choose the Multicut for its global optimality property . Hyperparameters of the pipeline212 were found by cross-validation ( see Appendix ) .213 4.1 Synthetic dataset : circles on structured ground214 To evaluate the feasibility of our approach , we create a synthetic dataset with prominent structured215 background . Our aim is to segment irregular disks on such background using only rule-based216 supervision . We generate the superpixels by the mutex watershed algorithm [ 25 ] which we run on217 the Gaussian gradient image . The node features of the superpixel graph were computed through218 self-supervised pretraining with contrastive loss as described in Appendix and fixed as part of the219 environment.220 As we aim to segment disks , we compute the circularity of the segmented objects for the rewards221 using the Circle Hough Transform [ 43 ] . This object-level reward is combined with the global rough222 estimate of the number of objects in the image to create the reward surface depicted in Fig . 3 . The223 reward for the number of objects provides useful gradient during early training stages : for example,224 when too few potential objects are found in the prediction , a low reward can be given to what is225 thought to be the background object . On the other hand , if too many potential objects are found , a226 low reward can be given to all the foreground objects with a low CHT value.227 In more detail , the object rewards rfg are composed as follows . We define a threshold γ on the CHT228 value ( γ = 0.8 in the reward surface shown in Fig . 3 ) . Let c ∈ [ 0 , 1 ] be the CHT value corresponding229 to the object and let k be the total number of objects that we expect and n be the number of predicted230 objects . Then231 rlocal = { σ ( ( c−γ1−γ − 0.5 ) 6 ) 0.4 , if c ≥ γ 0 , otw ( 5 ) rglobal = { rexp ( k n ) , if n ≥ k 0.6 , otw ( 6 ) rfg = rlocal + rglobal ( 7 ) Here σ ( · ) is the sigmoid function . The input to the sigmoid function is normalized to the interval232 [ −3 , 3 ] which was empirically found to be a good range . The rewards are always in [ 0 , 1 ] here this is233 split up into [ 0 , 0.5 ] for the local reward as well as for the global reward.234 For the largest predicted object we strongly suspect the background object . For this object background235 rewards rbg are calculated by236 rbg = { rexp ( n k ) , if n ≤ k 1 , otw ( 8 ) Note that this rewards have a large globally calculated part which makes this setup not fit for Method237 1 . It needs some feature representation that already gives a good idea for the clustering . The only238 useful local information in the reward is the CHT value . Therefore , if the features have a fairly239 distinct structure for circles , the agent should be able to find and to correctly cluster them.240 Fig . 4 shows the output of all algorithm components on a sample image . For comparison , we also241 computed mutex watershed [ 25 ] predictions . Texture within objects and structured background are242 inherently difficult for region-growing algorithms , but our approach can exploit higher-level reasoning243 along with low-level information and achieve a good segmentation.244 4.2 Real dataset : light microscopy imaging245 Biomedical applications often require segmentation of objects of known morphology which are246 positioned in regular patterns , while extensive prior knowledge is available on variability of both247 under normal experimental conditions [ 44 ] . Such data presents the best use case for our algorithm as248 the reward function can leverage the known characteristics of individual object shape and texture and249 the overall similarity of the objects.250 The dataset used for this experiment contains 317 2D images extracted from a video of a developing251 fruitfly embryo acquired with a light-sheet microscope [ 45 ] ( Fig . 5 ) . The image shows boundaries252 ( plasma membranes ) of the embryo cells . Across the dataset , 10 images were fully segmented by an253 expert , we use those for validation.254 Fruitfly embryo is a well-studied system for which we can exploit the prior knowledge on the expected255 cell shape and the radial pattern of cells . Furthermore , as the analysis of cell shape dynamics is256 a paramount part of many biological experiments , multiple pre-trained networks are available for257 the cell segmentation task [ 18 , 19 , 46 , 47 ] . Due to the differences in sample preparation and image258 acquisition settings , none of these would work out-of-the-box for our data . However , the CNNs in259 [ 47 ] which are trained to predict boundaries in confocal microscope images of plant tissue , can serve260 as a strong edge detector to create superpixels in our images . The superpixels are obtained using the261 seeded watershed algorithm on seeds at the local minima of the predicted edge map.262 The rewards for this experiment are designed as follows : we set a high reward for merging the263 superpixels which are certain to lie in the background ( close to the image boundary or the image264 center ) . For the background edges near the foreground area we modulate the reward by the circularity265 of the overall foreground contour . Finally , for the edges which are likely to be in the foreground266 we compute object-level rewards by fitting a rotated bounding box to each object and comparing its267 side lengths as well as its orientation to predefined template values . We do not perform semantic268 segmentation to define precise foreground/background boundaries , but instead use a soft weighting269 scheme with Gaussian weights to combine object and background rewards based on on the prior270 knowledge of the embryo width . An image of the weights for different locations in the image can be271 found in the appendix.272 More formally the edge rewards redge are calculated as follows . For each edge , we define the distance273 h between the edge and the center of the image as the average distance of the incident objects ’ center274 of mass and the center c of the image . j is the approximate radius of the circle that lies within the275 foreground and m is the maximal distance between c and the image boarder . Let further K ( · ) be the276 Gaussian kernel function . Then redge yields277 rbg = K ( ||h−c|| γ ) ( 1− a ) , if h ≤ j K ( ||m−h|| η ) ( 1− a ) , otw ( 9 ) rfg = K ( ||h− j|| δ ) max ( ro1 , ro2 ) ( 10 ) redge = rfg + rbg ( 11 ) Here γ , η , δ are normalization constants . Equation 9 first determines the background probability for278 an edge by the kernel values . 1− a constitutes a reward that directly favors merges which is scaled279 by the background probability . For each edge , ro1 and ro2 are the rewards corresponding to the two280 objects connected to that edge . The object rewards are given by fitting a rotated bounding box to the281 object and then compare rotation and dimensions to template values.282 Note that in this experiment no self-supervised pretraining is used for the node features in the283 agent ’ s GNNs . Unlike the “ Circles ” dataset , all objects in these images have very similar intensity284 distributions and can only be separated through the detection of boundaries between them . Instead285 of the pretraining , we experiment with using a few hand-crafted features like the polar coordinate286 of the node ’ s respective superpixel ’ s center of mass with respect to the coordinate system sitting287 at the center of the image as well as the superpixel ’ s mass , and with learning other features by288 back-propagation from the agent . The handcrafted features are normalized , concatenated to the289 learned features and used as input to the GNN . The projection of the first 3 PCA components of these290 features into RGB space is shown in Fig . 5 respectively for learned feature maps , their projection291 to node features through the accumulation procedure and finally the concatenation of those and the292 handcrafted features . Note that the learned features converge to a representation which resembles a293 semantic segmentation of boundaries in the image.294 We train the complete setup for Method 1 end-to-end on a Nvidia GeForce RTX 3090 GPU for 4295 days . For comparison we keep the model which achieved the highest reward on the test set . This296 makes training as well as the validation independent from ground truth annotations . The evolution297 of the rewards on the validation set for different random seeds is shown in the Appendix . All of the298 conducted trainings show a stride for high rewards regardless of different random seeding.299 For the validation scores we use the variation of information ( VI ) for both input combinations ( merge300 and split ) and the Symmetric Best Dice score . To show the influence of the imperfect superpixels on301 the final clustering , we project the superpixels to their respective ground truth clustering ( `` sp gt '' in302 Table 1 ) which sets an upper ( lower in case of VI ) bound for our method . In this study we use several303 versions of our approach . In Table 1 ( ours ) refers to method 1 as described in section 4.2 , ( ours +304 augmentation noise ) is the same method but add some noise to the input data during training , ( ours305 without edges ) is our method but without the additional edge prediction as an input and ( ours only306 handcrafted ) is our method where we only use the handcrafted features as described in section 4.2.307 We find that learned features significantly contribute to the performance of our method.308 We compare to the following baseline approaches : edge + mc , which solves the Multicut graph309 partitioning based on edge weights derived from boundary predictions used for superpixel creation,310 contrastive , which predicts a pixel-wise embedding space that is clustered into instances using311 k-means and for which the embeddings are trained using the discriminative loss function of [ 28 ] on312 the ovules dataset from [ 47 ] and contrastive + edge , which is similar to contrastive , but receives the313 [ 47 ] boundary predictions as additional input channel.314 5 Discussion315 We introduced an end-to-end instance segmentation algorithm which can exploit non-differentiable316 loss functions and high-level prior information . Our RL approach is based on stateless actor-critic317 and predicts the full segmentation at every step , allowing us to assign rewards to all objects and318 reach stable convergence . The segmentation problem is formulated as graph partitioning ; we design319 a reward decomposition algorithm which maps object- and image-level rewards to sub-graphs for320 localized supervision.321 We performed proof-of-concept experiments to demonstrate the feasibility of our approach on322 synthetic and real data and showed in particular that our setup can segment microscopy images323 with no direct supervision other than high-level reasoning . In the future , we plan to explore other324 problems and reward functions as well as a semi-supervised setup ( briefly introduced in Appendix ) 325 where we think our approach can be very beneficial . Furthermore , even in case of full supervision326 with ample groundtruth , our RL-based formulation enables end-to-end instance segmentation with327 direct object-level reasoning , which will allow for post-processing-aware training of the CNN which328 predicts object boundaries or embeddings.329 References330 [ 1 ] Kisuk Lee , Jonathan Zung , Peter Li , Viren Jain , and H Sebastian Seung . Superhuman accuracy331 on the snemi3d connectomics challenge . arXiv preprint arXiv:1706.00120 , 2017.332 [ 2 ] Liang-Chieh Chen , Huiyu Wang , and Siyuan Qiao . Scaling wide residual networks for panoptic333 segmentation , 2021.334 [ 3 ] S. Osher and N. Paragios . Geometric Level Set Methods in Imaging , Vision , and Graphics.335 Springer New York , 2007 . ISBN 9780387218106 . URL https : //books.google.de/books ? 336 id=ZWzrBwAAQBAJ.337 [ 4 ] Ricard Delgado-Gonzalo , Virginie Uhlmann , Daniel Schmitter , and Michael Unser . Snakes on338 a plane : A perfect snap for bioimage analysis . IEEE Signal Processing Magazine , 32 ( 1 ) :41–48,339 2014.340 [ 5 ] Ozan Oktay , Enzo Ferrante , Konstantinos Kamnitsas , Mattias Heinrich , Wenjia Bai , Jose341 Caballero , Stuart A. Cook , Antonio de Marvao , Timothy Dawes , Declan P. O ‘ Regan , Bernhard342 Kainz , Ben Glocker , and Daniel Rueckert . Anatomically constrained neural networks ( acnns ) :343 Application to cardiac image enhancement and segmentation . IEEE Transactions on Medical344 Imaging , 37 ( 2 ) :384–395 , 2018. doi : 10.1109/TMI.2017.2743464.345 [ 6 ] Hoel Kervadec , Jose Dolz , Meng Tang , Eric Granger , Yuri Boykov , and Ismail Ben Ayed.346 Constrained-cnn losses for weakly supervised segmentation . Medical Image Analysis , 54:88–99,347 2019 . ISSN 1361-8415. doi : https : //doi.org/10.1016/j.media.2019.02.009.348 [ 7 ] David Pfau and Oriol Vinyals . Connecting generative adversarial networks and actor-critic349 methods . CoRR , abs/1610.01945 , 2016 . URL http : //arxiv.org/abs/1610.01945.350 [ 8 ] Nikita Araslanov , Constantin Rothkopf , and Stefan Roth . Actor-critic instance segmentation.351 In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition ( CVPR ) ,352 2019.353 [ 9 ] Viren Jain , Srinivas Turaga , Kevin Briggman , Moritz Helmstaedter , Winfried Denk , and Hyun-354 june Seung . Learning to agglomerate superpixel hierarchies . Advances in Neural Information355 Processing Systems , 24 , 01 2011.356 [ 10 ] Peter Sunehag , Guy Lever , Audrunas Gruslys , Wojciech Marian Czarnecki , Vinicius Zambaldi,357 Max Jaderberg , Marc Lanctot , Nicolas Sonnerat , Joel Z. Leibo , Karl Tuyls , and Thore Graepel.358 Value-decomposition networks for cooperative multi-agent learning , 2017.359 [ 11 ] Drew Bagnell and Andrew Ng . On local rewards and scaling distributed reinforcement learning.360 In Y. Weiss , B. Schölkopf , and J. Platt , editors , Advances in Neural Information Processing361 Systems , volume 18 . MIT Press , 2006 . URL https : //proceedings.neurips.cc/paper/362 2005/file/02180771a9b609a26dcea07f272e141f-Paper.pdf.363 [ 12 ] Huazhe Xu , Boyuan Chen , Yang Gao , and Trevor Darrell . Scoring-aggregating-planning:364 Learning task-agnostic priors from interactions and sparse rewards for zero-shot generalization.365 CoRR , abs/1910.08143 , 2019 . URL http : //arxiv.org/abs/1910.08143.366 [ 13 ] Elad Eban , Mariano Schain , Alan Mackey , Ariel Gordon , Ryan Rifkin , and Gal Elidan . Scalable367 learning of non-decomposable objectives . In Artificial intelligence and statistics , pages 832–840.368 PMLR , 2017.369 [ 14 ] Maxim Berman , Amal Rannen Triki , and Matthew B Blaschko . The lovász-softmax loss:370 A tractable surrogate for the optimization of the intersection-over-union measure in neural371 networks . In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,372 pages 4413–4421 , 2018.373 [ 15 ] Yang Song , Alexander G. Schwing , Richard S. Zemel , and Raquel Urtasun . Training deep374 neural networks via direct loss minimization . International Conference on Machine Learning,375 2016.376 [ 16 ] Cicero Nogueira dos Santos , Kahini Wadhawan , and Bowen Zhou . Learning loss functions for377 semi-supervised learning via discriminative adversarial networks , 2017.378 [ 17 ] Josif Grabocka , Randolf Scholz , and Lars Schmidt-Thieme . Learning surrogate losses . arXiv379 preprint arXiv:1905.10108 , 2019.380 [ 18 ] Uwe Schmidt , Martin Weigert , Coleman Broaddus , and Gene Myers . Cell detection with381 star-convex polygons . In International Conference on Medical Image Computing and Computer-382 Assisted Intervention , pages 265–273 . Springer , 2018.383 [ 19 ] Carsen Stringer , Tim Wang , Michalis Michaelos , and Marius Pachitariu . Cellpose : a generalist384 algorithm for cellular segmentation . Nature Methods , 18 ( 1 ) :100–106 , 2021.385 [ 20 ] Xiaoling Hu , Fuxin Li , Dimitris Samaras , and Chao Chen . Topology-preserving386 deep image segmentation . In Advances in Neural Information Processing Systems,387 volume 32 , 2019 . URL https : //proceedings.neurips.cc/paper/2019/file/388 2d95666e2649fcfc6e3af75e09f5adb9-Paper.pdf.389 [ 21 ] Constantin Pape , Alex Matskevych , Adrian Wolny , Julian Hennies , Giulia Mizzon , Marion390 Louveaux , Jacob Musser , Alexis Maizel , Detlev Arendt , and Anna Kreshuk . Leveraging domain391 knowledge to improve microscopy image segmentation with lifted multicuts . Frontiers in392 Computer Science , 1:6 , 2019.393 [ 22 ] Jeremy B Maitin-Shepard , Viren Jain , Michal Januszewski , Peter Li , and Pieter Abbeel . Combi-394 natorial energy learning for image segmentation . In D. Lee , M. Sugiyama , U. Luxburg , I. Guyon,395 and R. Garnett , editors , Advances in Neural Information Processing Systems , volume 29 . Cur-396 ran Associates , Inc. , 2016 . URL https : //proceedings.neurips.cc/paper/2016/file/397 31857b449c407203749ae32dd0e7d64a-Paper.pdf.398 [ 23 ] Jie Song , Bjoern Andres , Michael J Black , Otmar Hilliges , and Siyu Tang . End-to-end learning399 for graph decomposition . In Proceedings of the IEEE/CVF International Conference on400 Computer Vision , pages 10093–10102 , 2019.401 [ 24 ] Naiyu Gao , Yanhu Shan , Yupei Wang , Xin Zhao , Yinan Yu , Ming Yang , and Kaiqi Huang.402 Ssap : Single-shot instance segmentation with affinity pyramid . In Proceedings of the IEEE/CVF403 International Conference on Computer Vision , pages 642–651 , 2019.404 [ 25 ] Steffen Wolf , Alberto Bailoni , Constantin Pape , Nasim Rahaman , Anna Kreshuk , Ullrich Köthe,405 and Fred A Hamprecht . The mutex watershed and its objective : Efficient , parameter-free graph406 partitioning . IEEE transactions on pattern analysis and machine intelligence , 2020.407 [ 26 ] Thorsten Beier , Constantin Pape , Nasim Rahaman , Timo Prange , Stuart Berg , Davi D Bock,408 Albert Cardona , Graham W Knott , Stephen M Plaza , Louis K Scheffer , et al . Multicut brings409 automated neurite segmentation closer to human performance . Nature methods , 14 ( 2 ) :101–102,410 2017.411 [ 27 ] Jan Funke , Fabian Tschopp , William Grisaitis , Arlo Sheridan , Chandan Singh , Stephan Saalfeld,412 and Srinivas C Turaga . Large scale image segmentation with structured loss based deep learning413 for connectome reconstruction . IEEE transactions on pattern analysis and machine intelligence,414 41 ( 7 ) :1669–1680 , 2018.415 [ 28 ] Bert De Brabandere , Davy Neven , and Luc Van Gool . Semantic instance segmentation with a416 discriminative loss function . arXiv preprint arXiv:1708.02551 , 2017.417 [ 29 ] Davy Neven , Bert De Brabandere , Marc Proesmans , and Luc Van Gool . Instance segmentation418 by jointly optimizing spatial embeddings and clustering bandwidth . In Proceedings of the419 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 8837–8845 , 2019.420 [ 30 ] Alberto Bailoni , Constantin Pape , Steffen Wolf , Thorsten Beier , Anna Kreshuk , and Fred A421 Hamprecht . A generalized framework for agglomerative clustering of signed graphs applied to422 instance segmentation . arXiv preprint arXiv:1906.11713 , 2019.423 [ 31 ] Leland McInnes and John Healy . Accelerated hierarchical density based clustering . In 2017424 IEEE International Conference on Data Mining Workshops ( ICDMW ) , pages 33–42 . IEEE,425 2017.426 [ 32 ] Dorin Comaniciu and Peter Meer . Mean shift : A robust approach toward feature space analysis.427 IEEE Transactions on pattern analysis and machine intelligence , 24 ( 5 ) :603–619 , 2002.428 [ 33 ] Bjoern Andres , Thorben Kroeger , Kevin L Briggman , Winfried Denk , Natalya Korogod , Graham429 Knott , Ullrich Koethe , and Fred A Hamprecht . Globally optimal closed-surface segmentation430 for connectomics . In European Conference on Computer Vision , pages 778–791 . Springer,431 2012.432 [ 34 ] Kaiming He , Georgia Gkioxari , Piotr Dollár , and Ross Girshick . Mask r-cnn . In Proceedings of433 the IEEE international conference on computer vision , pages 2961–2969 , 2017.434 [ 35 ] Marina Meilă . Comparing clusterings by the variation of information . In Learning theory and435 kernel machines , pages 173–187 . Springer , 2003.436 [ 36 ] William M Rand . Objective criteria for the evaluation of clustering methods . Journal of the437 American Statistical association , 66 ( 336 ) :846–850 , 1971.438 [ 37 ] Brian W Kernighan and Shen Lin . An efficient heuristic procedure for partitioning graphs . The439 Bell system technical journal , 49 ( 2 ) :291–307 , 1970.440 [ 38 ] Tuomas Haarnoja , Aurick Zhou , Kristian Hartikainen , George Tucker , Sehoon Ha , Jie Tan,441 Vikash Kumar , Henry Zhu , Abhishek Gupta , Pieter Abbeel , and Sergey Levine . Soft actor-critic442 algorithms and applications . CoRR , abs/1812.05905 , 2018 . URL http : //arxiv.org/abs/443 1812.05905.444 [ 39 ] Justin Gilmer , Samuel S. Schoenholz , Patrick F. Riley , Oriol Vinyals , and George E. Dahl.445 Neural message passing for quantum chemistry . CoRR , abs/1704.01212 , 2017 . URL http:446 //arxiv.org/abs/1704.01212.447 [ 40 ] Olaf Ronneberger , Philipp Fischer , and Thomas Brox . U-net : Convolutional networks for448 biomedical image segmentation . CoRR , abs/1505.04597 , 2015 . URL http : //arxiv.org/449 abs/1505.04597.450 [ 41 ] Richard S. Sutton and Andrew G. Barto . Reinforcement Learning : An Introduction . The MIT451 Press , second edition , 2018 . URL http : //incompleteideas.net/book/the-book-2nd.452 html.453 [ 42 ] Jörg Hendrik Kappes , Markus Speth , Björn Andres , Gerhard Reinelt , and Christoph Schn.454 Globally optimal image partitioning by multicuts . In Yuri Boykov , Fredrik Kahl , Victor455 Lempitsky , and Frank R. Schmidt , editors , Energy Minimization Methods in Computer Vision456 and Pattern Recognition , pages 31–44 , Berlin , Heidelberg , 2011 . Springer Berlin Heidelberg.457 ISBN 978-3-642-23094-3.458 [ 43 ] Allam Shehata Hassanein , Sherien Mohammad , Mohamed Sameer , and Mohammad Ehab Ragab.459 A survey on hough transform , theory , techniques and applications . CoRR , abs/1502.02160,460 2015 . URL http : //arxiv.org/abs/1502.02160.461 [ 44 ] D ’ Arcy Wentworth Thompson . On Growth and Form . Canto . Cambridge University Press,462 1992. doi : 10.1017/CBO9781107325852.463 [ 45 ] Sourabh Bhide , Ralf Mikut , Maria Leptin , and Johannes Stegmaier . Semi-automatic generation464 of tight binary masks and non-convex isosurfaces for quantitative analysis of 3d biological465 samples , 2020.466 [ 46 ] Lucas von Chamier , Romain F Laine , Johanna Jukkala , Christoph Spahn , Daniel Krentzel , Elias467 Nehme , Martina Lerche , Sara Hernández-Pérez , Pieta K Mattila , Eleni Karinou , Séamus Holden,468 Ahmet Can Solak , Alexander Krull , Tim-Oliver Buchholz , Martin L Jones , Loïc A Royer,469 Christophe Leterrier , Yoav Shechtman , Florian Jug , Mike Heilemann , Guillaume Jacquemet,470 and Ricardo Henriques . Democratising deep learning for microscopy with ZeroCostDL4Mic.471 Nature Communications , 4 2021.472 [ 47 ] Adrian Wolny , Lorenzo Cerrone , Athul Vijayan , Rachele Tofanelli , Amaya Vilches Barro,473 Marion Louveaux , Christian Wenzl , Sören Strauss , David Wilson-Sánchez , Rena Lymbouridou,474 et al . Accurate and versatile 3d segmentation of plant tissues at cellular resolution . Elife , 9:475 e57613 , 2020.476 [ 48 ] Bert De Brabandere , Davy Neven , and Luc Van Gool . Semantic instance segmentation with a477 discriminative loss function , 2017.478 [ 49 ] Reuben R. Shamir , Yuval Duchin , Jinyoung Kim , Guillermo Sapiro , and Noam Harel.479 Continuous dice coefficient : a method for evaluating probabilistic segmentations . CoRR,480 abs/1906.11031 , 2019 . URL http : //arxiv.org/abs/1906.11031.481 Checklist482 1 . For all authors ... 483 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s484 contributions and scope ? [ Yes ] 485 ( b ) Did you describe the limitations of your work ? [ Yes ] 486 ( c ) Did you discuss any potential negative societal impacts of your work ? [ No ] It does not487 have any negative societal impacts.488 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to489 them ? [ Yes ] 490 2 . If you are including theoretical results ... 491 ( a ) Did you state the full set of assumptions of all theoretical results ? [ N/A ] We do not492 claim any theoretical assumptions.493 ( b ) Did you include complete proofs of all theoretical results ? [ N/A ] See above.494 3 . If you ran experiments ... 495 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-496 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] Yes , see the497 Supplementary.498 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they499 were chosen ) ? [ Yes ] See the Supplementary500 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-501 ments multiple times ) ? [ Yes ] The bars are reported w.r.t . to the samples within the best502 seed , see Supplementary section for further details503 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type504 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] 505 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 506 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 507 ( b ) Did you mention the license of the assets ? [ Yes ] 508 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 509 The link to the biological dataset is given in the Supplementary510 ( d ) Did you discuss whether and how consent was obtained from people whose data511 you ’ re using/curating ? The dataset was provided from the other laboratory as a part of512 collaboration.513 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable514 information or offensive content ? [ No ] No human data is used515 5 . If you used crowdsourcing or conducted research with human subjects ... 516 ( a ) Did you include the full text of instructions given to participants and screenshots , if517 applicable ? [ N/A ] 518 ( b ) Did you describe any potential participant risks , with links to Institutional Review519 Board ( IRB ) approvals , if applicable ? [ N/A ] 520 ( c ) Did you include the estimated hourly wage paid to participants and the total amount521 spent on participant compensation ? [ N/A ] 522

Summary: The paper proposes the use of RL for instance segmentation. In particular, the paper uses super pixels together with graph partitioning and RL to learn the labels of instances. The model is validated on two datasets, one synthetic and one biomedical imaging dataset and is numerically compared to two approaches ([28] and [47]).