Actor-Critic Instance Segmentation

TL;DR

This paper introduces an actor-critic reinforcement learning approach for instance segmentation, where the actor predicts masks sequentially and the critic models long-term effects, improving accuracy over baseline methods.

Contribution

The paper proposes a novel actor-critic framework with a learned mask representation for sequential instance segmentation, addressing limitations of previous global assignment methods.

Findings

Consistent accuracy improvements over recurrent baseline

Effective exploration via learned mask representation

Enhanced segmentation performance on standard benchmarks

Abstract

Most approaches to visual scene analysis have emphasised parallel processing of the image elements. However, one area in which the sequential nature of vision is apparent, is that of segmenting multiple, potentially similar and partially occluded objects in a scene. In this work, we revisit the recurrent formulation of this challenging problem in the context of reinforcement learning. Motivated by the limitations of the global max-matching assignment of the ground-truth segments to the recurrent states, we develop an actor-critic approach in which the actor recurrently predicts one instance mask at a time and utilises the gradient from a concurrently trained critic network. We formulate the state, action, and the reward such as to let the critic model long-term effects of the current prediction and incorporate this information into the gradient signal. Furthermore, to enable effective exploration in the inherently high-dimensional action space of instance masks, we learn a compact representation using a conditional variational auto-encoder. We show that our actor-critic model consistently provides accuracy benefits over the recurrent baseline on standard instance segmentation benchmarks.

Figures

Tables

Table 1. Table 1 : Evaluation on CVPPP val. We compare our baseline with fully-unrolled ( BL ) and truncated BPTT ( BL-Trunc ) to the actor-critic with Dice-based reward, with ( AC-Dice ) and without ( AC-Dice-NoKL ) exploration, and without the State Pyramid ( AC-Dice-NoSP ).

Model	SBD $\uparrow$	$\mid$DiC$\mid$ $\downarrow$
BL	80.0	1.08
BL-Trunc	79.4	1.32
AC-Dice	80.5	0.88
AC-Dice-NoKL	75.4	1.36
AC-Dice-NoSP	61.3	1.52

Table 2. Table 2 : Contribution of recurrent states to mask quality measured by Dice and absolute Difference in Counting ∣ ∣ \mid DiC ∣ ∣ \mid on CVPPP val.

Model	LSTM + Mask		Mask only		LSTM only
	Dice∗ $\uparrow$	$\mid$DiC$\mid$ $\downarrow$	Dice∗ $\uparrow$	$\mid$DiC$\mid$ $\downarrow$	Dice∗ $\uparrow$	$\mid$DiC$\mid$ $\downarrow$
BL	78.6	1.04	76.6	4.36	6.5	3.96
BL-Trunc	77.9	1.72	77.5	6.24	6.0	4.8
AC-Dice	78.4	0.88	78.5	1.92	5.8	4.36
∗ computed by max-matching and ground-truth stopping

Table 3. Table 3 : Segmentation quality of our actor-critic model on CVPPP test with Dice-based reward ( AC-Dice ) in terms of Symmetric Best Dice (SBD) and absolute Difference in Counting ( ∣ ∣ \mid DiC ∣ ∣ \mid ).

Model	SBD $\uparrow$	$\mid$DiC$\mid$ $\downarrow$
RIS [32]	66.6	1.1
MSU [33]	66.7	2.3
Nottingham [33]	68.3	3.8
IPK [29]	74.4	2.6
DLoss [9]	84.2	1.0
E2E [31]	84.9	0.8
Ours (AC-Dice)	79.1	1.12

Table 4. (a) KITTI test set

Model	MWCov $\uparrow$	MUCov $\uparrow$	AvgFP $\downarrow$	AvgFN $\downarrow$
DepthOrder [42]	70.9	52.2	0.597	0.736
DenseCRF [41]	74.1	55.2	0.417	0.833
AngleFCN+D [35]	79.7	75.8	0.201	0.159
E2E [31]	80.0	66.9	0.764	0.201
Ours (BL-Trunc)	72.2	50.7	0.393	0.432
Ours (AC-IoU)	75.6	57.3	0.338	0.309

Table 5. (a) KITTI test set

Model	MWCov $\uparrow$	MUCov $\uparrow$	AvgFP $\downarrow$	AvgFN $\downarrow$
DepthOrder [42]	70.9	52.2	0.597	0.736
DenseCRF [41]	74.1	55.2	0.417	0.833
AngleFCN+D [35]	79.7	75.8	0.201	0.159
E2E [31]	80.0	66.9	0.764	0.201
Ours (BL-Trunc)	72.2	50.7	0.393	0.432
Ours (AC-IoU)	75.6	57.3	0.338	0.309

Table 6. (b) KITTI validation set

Model	MWCov $\uparrow$	MUCov $\uparrow$	AvgFP $\downarrow$	AvgFN $\downarrow$
E2E (Iter-1)	64.1	54.8	0.200	0.375
E2E (Iter-3)	71.3	63.4	0.417	0.308
E2E (Iter-5)	75.1	64.6	0.375	0.283
Ours (BL-Trunc)	70.4	55.8	0.313	0.339
Ours (AC-IoU)	71.9	59.5	0.262	0.253

Table 7. Table 5 : Architecture of the actor network. SP m subscript SP 𝑚 \text{SP}_{m} denotes the additional channels in the State Pyramid provided at resolution scaled by 2 m superscript 2 𝑚 2^{m} .

Section	Type	(Kernel) size	Stride	# of channels
Encoder	Conv $\to$ ReLU	$3\times 3$	1	32
	MaxPooling	$2\times 2$	2	–
	Conv $\to$ ReLU	$3\times 3$	1	48
	MaxPooling	$2\times 2$	2	–
	Conv $\to$ ReLU	$3\times 3$	1	64
	MaxPooling	$2\times 2$	2	–
	Conv $\to$ ReLU	$3\times 3$	1	96
	MaxPooling	$2\times 2$	2	–
	Conv $\to$ ReLU	$3\times 3$	1	128
	MaxPooling	$2\times 2$	2	–
Bottleneck	FC $\to$ LeakyReLU	$128\times h\times w$	–	$\text{size}\left(h_t\right)$
	LSTM	$\text{size}\left(h_t\right)$	–	$\text{size}\left(h_t\right)$
	FC $\to$ LeakyReLU	$\text{size}\left(h_t\right)$	–	$z$
	FC ($\mu$, $\sigma$)	$z$	–	$2\times l$
	FC $\to$ LeakyReLU	$l$	–	$z$
	FC $\to$ LeakyReLU	$z$	–	$128\times h\times w$
Decoder	Deconv $\to$ ReLU	$3\times 3$	2	128 + ${\text{SP}}_5$
	Deconv $\to$ ReLU	$3\times 3$	2	96 + ${\text{SP}}_4$
	Deconv $\to$ ReLU	$3\times 3$	2	64 + ${\text{SP}}_3$
	Deconv $\to$ ReLU	$3\times 3$	2	48 + ${\text{SP}}_2$
	Deconv $\to$ ReLU	$3\times 3$	2	32 + ${\text{SP}}_1$
	Deconv $\to$ ReLU	$3\times 3$	1	1 + ${\text{SP}}_0$

Table 8. Table 6 : Parameter values for the actor-critic.

Dataset	$h\times w$	$size\left(h_t\right)$	$l$	$z$
CVPPP	$7\times 7$	$512$	$16$	$256$
KITTI	$8\times 24$	$512$	$64$	$512$

Table 9. Table 7 : Architecture of the critic network.

Type	(Kernel) size	Stride	# of output channels
Conv $\to$ BN $\to$ ReLU	$3\times 3$	1	32
MaxPooling	$2\times 2$	2	–
Conv $\to$ BN $\to$ ReLU	$3\times 3$	1	64
MaxPooling	$2\times 2$	2	–
Conv $\to$ BN $\to$ ReLU	$3\times 3$	1	128
MaxPooling	$2\times 2$	2	–
Conv $\to$ BN $\to$ ReLU	$3\times 3$	1	256
MaxPooling	$2\times 2$	2	–
Conv $\to$ BN $\to$ ReLU	$3\times 3$	1	512
Global MaxPooling	$h\times w$	1	512
FC $\to$ LeakyReLU	512	–	1024
FC $\to$ LeakyReLU	1024	–	1024
FC $\to$ LeakyReLU	1024	–	512
FC	512	–	1

Equations

$$\underset{P\in \mathcal{P}}{\mathrm{argmax}}\mathrm{tr}(U_{\theta }P),$$

$$\operatorname*{arg\,max}_{P\in\mathcal{P}}\>\operatorname{tr}\big{(}(U_{\theta}+\epsilon)P\big{)}.$$

$$T\big{(}(I,M_{t}),a_{t}\big{)}=\big{(}I,\max(M_{t},\mathcal{D}(a_{t}))\big{)}$$

$${\varphi }_t:=\underset{k\in \mathcal{P}(N)}{\max }\sum _{i=1}^t\mathcal{F}({\mathcal{S}}_i,{\mathcal{T}}_{k_i}),$$

$$r_t:=\varphi (s_{t+1})-\varphi (s_t).$$

$$Q_{\phi}(s_{t},a_{t})=\mathbb{E}_{a_{j}\sim\mu(s_{j}),j>t}\Bigg{[}\sum_{i=t}^{N}\gamma^{i-t}r_{i}(s_{i},a_{i})\Bigg{]}.$$

$$\displaystyle\mathcal{L}_{\text{Critic},t}=\Bigg{\|}Q_{\phi}(s_{t},a_{t})-\sum_{i=t}^{N}\gamma^{i-t}r_{i}\Bigg{\|}_{2}^{2},$$

Full text

1 Introduction

Methods for instance segmentation have for the most part relied on the idea of parallel processing of the image elements and features within images [13]. However, previous work [31, 32] suggests that instance segmentation can be formulated as a sequential visual task, akin to human vision, for which substantial evidence has revealed that many vision tasks beyond eye movements are solved sequentially [36]. While the segmentation accuracy of feed-forward pipelines hinges on a large number of object proposals, proposal-free recurrent models have a particular appeal for instance segmentation where the number of instances is unknown. Also, the temporal context can facilitate a certain order of prediction: segmenting “hard” instances can be improved by conditioning on the masks of “easy” instances segmented first (e.g., due to occlusions, ambiguities in spatial context etc.; [20]).

A pivotal question of a recurrent formulation for instance segmentation is the assignment of the ground-truth segments to timesteps, since the order in which they have to be predicted is unknown. Previously this was addressed using the Kuhn-Munkres algorithm [18], computing the max-matching assignment. We provide some insight, however, that the final prediction ordering depends on the initial assignment. Furthermore, the loss for every timestep is not informative in terms of its effect on future predictions. Intuitively, considering the future loss for the predictions early on should improve the segmentation accuracy at the later timesteps. Although this can be achieved by unrolling the recurrent states for gradient backpropagation, such an approach quickly becomes infeasible for segmentation networks due to high memory demands.

In the past years, reinforcement learning (RL) has been showing promise in solving increasingly complex tasks [23, 26, 27]. However, relatively little work has explored applications of RL outside its conventional domain, which we attribute to two main factors: (1) computer vision problems often lack the notion of the environment, which provides the interacting agent with the reward feedback; (2) actions in the space of images are often prohibitively high-dimensional, leading to tough computational challenges.

Here, we use an actor-critic (AC) model [5] to make progress regarding both technical issues of a recurrent approach to instance segmentation. We use exploration noise to reduce the influence of the initial assignments on the segmentation ordering. Furthermore, we design a reward function that accounts for the future reward in the objective function for every timestep. Our model does not use bounding boxes – often criticised due to their coarse representation of the objects’ shape. Instead, we built on an encoder-decoder baseline that makes pixelwise predictions directly at the scale of the input image. To enable the use of RL for instance segmentation with its associated high-dimensional output space, we propose to learn a compact action-space representation through the latent variables of a conditional variational auto-encoder [16], which we integrate into a recurrent prediction pipeline.

Our experiments demonstrate that our actor-critic model improves the prediction quality over its baseline trained with the max-matching assignment loss, especially at the later timesteps, and performs well on standard instance segmentation benchmarks.

2 Related Work

Instance segmentation has received growing attention in the recent literature. One family of approaches focuses on learning explicit instance encodings [4, 9, 17, 24, 35], which are then clustered into individual instance masks using post-processing. Another common end-to-end approach is to first predict a bounding box for each instance using dynamic pooling and then to produce a mask of the dominant object within the box using a separate segmentation network [13, 22]. These methods are currently best-practice, which can be attributed to the maturity of deep network-based object detection pipelines. However, this strategy is ultimately limited by the detection performance, proposal set, and the need of additional processing to account for pixel-level context [2, 8].

Making the predictions sequentially points at an alternative line of work. Romera-Paredes & Torr [32] used a convolutional LSTM [39] with a spatial softmax, which works well for isotropic object shapes and moderate scale variation. At each timestep, the recurrent model of Ren & Zemel [31] predicts a box location and scale for one instance. However, the extent of the available context for subsequent segmentations is limited by the box. Some benefits of the temporal and spatial context have been also re-asserted on the task of object detection [7, 21] and, much earlier, on image generation [12] and recognition [19]. In contrast to these works, our method obviates the need for the intermediate bounding box representation and predicts masks directly at the image resolution.

We cast the problem as a sequential decision process, as is studied by reinforcement learning (RL; [34]). Using the actor-critic framework [5], we define the actor as the model that sequentially produces instance masks, whereas the critic learns to provide a score characterising the actor’s performance. Leveraging this score, the actor can be trained to improve the quality of its predictions. This is reminiscent of the more recent Generative Adversarial Networks (GANs; [11]), in which a generator relies on a discriminator to improve. In particular, our model is similar to Wasserstein GANs [1] in that the discriminator is trained on a regression-like loss, and to SeqGAN [40] in that the generator’s predictions are sequential.

One obstacle is the action dimensionality, since the sampling complexity required for exploration grows exponentially with the size of the actions. A naive action representation for dense pixelwise predictions would lead to an action space of dimension in the order of $O(2^{H\times W})$ for images with resolution $H\times W$. This is significantly higher than the action spaces of standard problems studied by reinforcement learning (usually, between 1 and 20), or even its applications to natural language processing [3, 30]. To address this, we suggest learning a compact representation using variational auto-encoders [16] to enable the crucial reduction of the problem from a high-dimensional discrete to a lower-dimensional continuous action space.

3 Motivation

As discussed above, we follow previous work in modelling instance segmentation as a sequential decision problem [31, 32], yielding one instance per timestep $t$.

To motivate our work, we revisit the standard practice of using the Kuhn-Munkres algorithm [18] to assign the ground-truth instances to the predictions of a recurrent model. Let $\theta$ parametrise the model and $U_{\theta}\in\mathbb{R}^{n\times n}$ denote a matrix of elements $u_{ij}$ measuring the score of the $i^{\text{th}}$ prediction w.r.t. the $j^{\text{th}}$ ground truth (e.g., the IoU). The Kuhn-Munkres algorithm finds a permutation matrix as a solution to the max-matching problem

[TABLE]

where $\mathcal{P}$ is the set of $n$-dimensional permutation matrices, i.e. such that for all $P\in\mathcal{P}$, we have $\sum_{j}P_{ij}=1,\sum_{i}P_{ij}=1$, $P_{ij}\in\{0,1\}$. Given a differentiable loss function $l_{\theta}(i,j)$ (e.g., the binary cross-entropy), the model parameters $\theta$ are then updated to minimise $\sum_{ij,P_{ij}=1}l_{\theta}(i,j)$.

Consider a simple case of two ground-truth segments, $a$ and $b$, illustrated in Fig. 1. Without loss of generality, assume that the initial (random) model parameters yield $u_{1a}+u_{2b}<u_{1b}+u_{2a}$, i.e. the sum of scores for segmenting instance $a$ first and $b$ second is lower than in the opposite order. This implies that max-matching will perform a gradient update step maximising the second sum, i.e. $u_{1b}+u_{2a}$, but not the first. As a consequence, for the updated parameters, the score for the ordering $b\rightarrow a$ is likely to dominate also at the later iterations of training.111Note that a formal proof is likely non-trivial due to the stochastic nature of training.

Previous work [20, 37] suggests that sequential models are not invariant to the order of predictions, including object segments (c.f. supplemental material). The implication from the example above is that $\sup\;u_{1a}+u_{2b}\neq\sup\;u_{1b}+u_{2a}$ (the $\sup$ is w.r.t. $\theta$). One conceivable remedy to alleviate the effect of the initial assignment is to introduce noise $\epsilon$ to the score matrix $U$ (e.g., i. i. d. Gaussian), such that Eq. 1 becomes

[TABLE]

However, the noise in the loss function will not account for the inherent causality of the temporal context in recurrent models: perturbation of one prediction affects the consecutive ones.

In this work, we consider a more principled approach to encourage exploration of different orderings. We inject exploration noise at the level of individual predictions made by an actor network, while a jointly trained critic network keeps track of the long-term outcome of the early predictions. This allows to include in the gradient to the actor not only the immediate loss, but also the contribution of the current prediction to the future loss. We achieve this by re-formulating the instance segmentation problem in the RL framework, which we briefly introduce next.

4 Notation and Definitions

In the following, we define the key concepts of Markov decision processes (MDPs) in the context of instance segmentation. For a more general introduction, we refer to [34].

We consider finite-horizon MDPs defined by the tuple $(S,A,T,r)$, where the state space $S$, the action space $A$, the state transition $T:S\times A\rightarrow S$, and the reward $r:S\times A\rightarrow\mathbb{R}$ are defined as follows.

State.  The state $s_{t}\in S$ of the recurrent system is a tuple of the input image (and its task-specific representations) and an aggregated mask, i.e. $s_{t}=(I,M_{t})$. The mask $M_{t}$ simply accumulates previous instance predictions, which encourages the model to focus on yet unassigned pixels. Including $I$ enables access to the original input at every timestep.

Action.  To limit the dimensionality of the action space, we define the action $a_{t}\in A$ in terms of a compact mask representation. To achieve this, we pre-train a conditional variational auto-encoder (cVAE; [16]) to reproduce segmentation masks. As a result, the action $a_{t}\in A=\mathbb{R}^{l}$ is a continuous latent representation of a binary mask and has dimensionality $l\ll H\cdot W$, while the decoder $\mathcal{D}:\mathbb{R}^{l}\rightarrow\mathbb{R}^{H\times W}$ “expands” the latent code to a full-resolution mask.

State transition.  As implied by the state and action definitions above, the state transition

[TABLE]

uses a pixelwise max of the previous mask and the decoded action, i.e. integrating the currently predicted instance mask into the previously accumulated predictions.

Reward.  We design the reward function to measure the progress of the state transition towards optimising a certain segmentation criterion. The building block of the reward is the state potential [28], which we base on the max-matching assignment of the current predictions to the ground-truth segments, i.e.

[TABLE]

where $\mathcal{T}_{i}$ and $\mathcal{S}_{1\leq i\leq t}$ are the $N$ ground-truth masks and $t$ predicted masks; $\mathcal{P}(N)$ is a collection of all permutations of the set $\{1,2,...,N\}$. $\mathcal{F}(\cdot,\cdot)$ denotes a distance between the prediction and a ground-truth mask and can be chosen with regard to the performance metric used by the specific benchmark (e.g., IoU, Dice, etc.). We elaborate on these choices in the experimental section.

The state potential in Eq. 4 allows us to define the reward as the difference between the potentials of subsequent states

[TABLE]

Note that since the $(t+1)^{\text{st}}$ prediction might re-order the optimal assignment (computed with the Kuhn-Munkres algorithm), our definition of the reward is less restrictive w.r.t. the prediction order compared to previous work [31, 32], which enforces a certain assignment to compute the gradient. Instead, our immediate reward allows to reason about the relative improvement of one set of predictions over another.

5 Actor-Critic Approach

5.1 Overview

The core block of the actor model, shown in Fig. 2, is a conditional variational auto-encoder (cVAE; [16]). The encoder computes a compact vector of latent variables, encoding a full-resolution instance mask. The decoder recovers such a mask from the latent code. Using the transition function defined by Eq. 3, the latest prediction updates the state, and the procedure repeats until termination.

The actor relies on two types of context with complementary properties. As discussed above, the mask $M_{t}$ is a component of the state $s_{t}$, which accumulates the masks produced in the previous steps. It provides permutation-invariant temporal context of high-resolution cues, encouraging the network to focus on yet unlabelled pixels. The hidden state of actor $h_{t}$ is implemented by the LSTM [14] at the bottleneck section and is unknown to the critic. In contrast to state $s_{t}$, the representation of the hidden state is learned and can be sensitive to the prediction ordering due to the non-commutative updates of the LSTM state. The hidden state, therefore, contributes to the temporal context and is shown to be particularly helpful for counting in the ablation study.

We train our model in two stages as described next.

5.2 Pre-training

We pre-train the actor cVAE to reconstruct the mask of the target segment. The input to the network consists of the image and the binary mask of a randomly chosen ground-truth instance. To account for the loss of high-resolution information at the latent level, the decoder is conditioned on the input image and auxiliary channels of instance-relevant representations supplied at multiple scales, which we term State Pyramid. One channel contains the foreground prediction, while the other 8 channels encode the instance angle quantisation of [35], thereby binning the pixels of the object segment into quantised angles relative to the object’s centroid. These features assist in instance detection and disambiguation, since a neighbourhood of orthogonal quantisation vectors indicates occlusion boundaries and object centroids. Following Ren & Zemel [31], we predict the angles with a pre-processing network [25] trained in a standalone fashion.

The auto-encoder uses the binary cross-entropy (BCE) as the reconstruction loss. For the latent representation, we use a Gaussian prior with zero mean and unit variance. The corresponding loss function is taken as the Kullback-Leibler divergence [16].

5.3 Training

During training, we learn a new encoder to sequentially predict segmentation masks. In addition to the image, the encoder also receives the auxiliary channels used during pre-training. In contrast, however, the encoder is encouraged to learn instance-sensitive features, since the decoder expects the latent code of ground-truth masks.

Algorithm 1 provides an outline of the training procedure. The actor is trained jointly with the critic from a buffer of experience accumulated in the episode execution step. In the policy evaluation, the critic is updated to minimise the error of approximating the expected reward, while in the policy iteration the actor receives the gradient from the critic to maximise the Q-value.

Episode execution.  For an image with $N$ instances, we define the episode as the sequence of $N$ predictions. The algorithm randomly selects a mini-batch of images without replacement and provides it as inputs to the actor. Using the reparametrisation trick [16], the actor samples an action corresponding to the next prediction of the instance mask. The results of the predictions and corresponding rewards are saved in a buffer. At the end of each episode, the target Q-value can be computed for each timestep $t$ as a sum of immediate rewards.

Policy evaluation.  The critic network parametrised by $\phi$, maintains an estimate of the Q-value defined as a function of the state and action guided by the policy $\mu$:

[TABLE]

Note the finite sum in the expectation due to the finite number of instances in the image. The critic’s loss $\mathcal{L}_{\text{Critic},t}$ for timestep $t$ is defined as the squared $\mathcal{L}_{2}$-distance with a discounted sum of rewards

[TABLE]

where $\gamma\in(0,1)$ is a discount factor that controls the time horizon, i.e. the degree to which future rewards should be accounted for. The hyperparameter $\gamma$ allows to trade off the time horizon for the difficulty of the reward approximation: as $\gamma\rightarrow 1$, the time horizon extends to all states, but the critic has to approximate a distant future reward based only on the current state and action. We update the parameters of the critic to minimise Eq. 7 using the samples of the state-actions and rewards in the buffer, and set $\gamma=0.9$ throughout our experiments.

Policy iteration.  The actor samples an action $a_{t}\in A$ from a distribution provided by the current policy $\mu_{\theta}:S\rightarrow A$, parametrised by $\theta$ and observes a reward $r_{t}$ computed by Eq. 5. Given the initial state $s_{1}$, the actor’s goal is to find the policy maximising the expected total reward, \theta^{\ast}=\operatorname*{arg\,max}_{\theta}\>\mathbb{E}_{a_{j}\sim\mu_{\theta}(s_{j})}\big{[}\sum_{i=1}^{N}\gamma^{i}r_{i}(s_{i},a_{i})\big{]}, approximated by the critic. To achieve this, the state $s_{t}=(I,M_{t})$ and the actor’s mask prediction $m_{t}$ are passed to the critic, which produces the corresponding Q-value. The gradient maximising the Q-value is computed via backpropagation and returned to the actor for its parameter update.

We found that fixing the decoder during training led to faster convergence. Since the critic only approximates the true loss, its gradient is biased, which in practice can break the assumption we maintain during training – that an optimal mask can be reconstructed from a normally-distributed latent space. We fix the decoder and maintain the KL-divergence loss $\mathcal{L}_{\text{KL}}$ while sampling new actions, thus encouraging exploration of the action space. In our ablation study, we verify that such exploration improves the segmentation quality. Note that we do not pre-define the layout of the actions, but only maintain the Gaussian prior.

To further improve the stability of the joint actor-critic training, we use a warm-up phase for the critic: episode execution and update of the critic happen without updating the actor for a number of epochs. This gives the critic the opportunity to adapt to the current action and state space of the actor. We could confirm in our experiments that pre-training the decoder was crucial; omitting this step resulted in near-zero rewards from which it proved to be difficult to train the critic even with the warm-up phase.

Termination.  We connect the hidden state $h_{t}$ and the last layer preceding it (via a skip connection) to a single unit predicting “1” to continue prediction, and “0” to indicate the terminal state. Using the ground-truth number of instances, we train this unit with the BCE loss.

Inference.  We recurrently run the actor network until the termination prediction. To obtain the masks, we discard the deviation part and only take the mean component of the action predicted by the encoder and pass it through the decoder. We do not use the critic network at inference time.

Implementation details.222Code is available at https://github.com/visinf/acis/.  We use a simple architecture similar to [32] for both the critic and the actor networks trained with Adam [15] until the training loss on validation data does not improve (c.f. supplemental material).

5.4 Discussion

In the actor-critic model the critic plays the role of modelling the subsequent rewards for states $s_{i>t}$ given state $s_{t}$. Hence, if the critic’s approximation is exact, the backpropagation through time (BPTT; [38]) until the first state is not needed: to train the actor, we need to compute the gradient w.r.t. the future rewards already predicted by the critic. The implication of this property is that memory-demanding networks, such as those for dense prediction, can be effectively trained with truncated BPTT and the critic, even in case of long sequences. Moreover, using the critic’s approximation allows the reward be a non-differentiable, or even a discontinuous function tailored specifically to the task.

6 Experiments

In our experiments, we first quantitatively verify the importance of the different components in our model and investigate the sources of the accuracy benefits of the actor-critic over the baseline. Then, we use two standard datasets of natural images for the challenging task of instance segmentation, and compare to the state of the art.

6.1 Ablation study

We design a set of experiments to investigate the effect of various aspects of the model using the A1 benchmark of the Computer Vision Problems in Plant Phenotyping (CVPPP) dataset [33]. It comprises a collection of 128 images of plants taken from a top view with leaf annotation as ground-truth instance masks. We downsized the original 128 images in the training set by a factor of two and used a centre crop of size $224\times 224$ for training. For the purpose of the ablation study, we randomly select 103 images from the CVPPP A1 benchmark for training and report the results on the remaining 25 images.

To compute the reward for our actor-critic model (Eq. 4), we use the Dice score computed as $\mathcal{F}(\mathcal{S},\mathcal{T})=\frac{2\sum_{i}{\mathcal{S}_{i}\mathcal{T}_{i}}}{\sum_{i}\mathcal{S}_{i}+\sum_{i}\mathcal{T}_{i}}$. The dimensionality of the latent action space is fixed to 16.

In the first part of the experiment, we look into how different terms in the loss influence the segmentation quality, measured in Symmetric Best Dice (SBD), and the absolute Difference in Counting ($\mid$DiC$\mid$). Specifically, we train five models: BL is an actor-only recurrent model trained with BPTT through all states. We use the BCE loss and Dice-based max-matching as a heuristic for assigning the ground truth to predictions, similar to [31, 32]. BL-Trunc is similar to BL, but is trained with a truncated, one-step BPTT. We train our actor-critic model AC-Dice with the gradient from the critic approximating the Dice score. AC-Dice-NoKL is similar to the AC-Dice model, i.e. the actor is trained jointly with the critic, but we remove the KL-divergence term, which encourages exploration, from the loss of the actor. Lastly, we verify the benefit of the State Pyramid, the multi-res spatial information provided to the decoder, by comparing to a baseline without it (AC-Dice-NoSP).

The side-by-side comparison of these models summarised in Table 1 reveals that AC-Dice exhibits a superior accuracy compared to the baselines, both in terms of Dice and counting. Using the KL-divergence term in the loss improves the actor, which shows the value of action exploration in a consistent action space. We also observed that training AC-Dice-NoKL would sometimes diverge and require a restart with the critic warm-up. Furthermore, the State Pyramid aids the decoder, as removing it leads to a significant drop in mask quality. Surprisingly, BL-Trunc is only slightly worse than BL, which however has by far higher memory demands than both AC-Dice and BL-Trunc in the setting of long sequences and high resolutions.

To further investigate the accuracy gains of the actor-critic model, we report the average Dice score w.r.t. the corresponding timestep of the prediction in Fig. 3. The histogram confirms our intuition that incorporating the future reward into the loss function for every timestep, as modelled by the critic, should improve the segmentation quality at later stages of prediction: the Dice score of the actor-critic model tends to be tangibly higher especially at the later timesteps. Note that the contribution of this benefit to the average score across the dataset is moderated by not all images in the dataset having many instances.

In the next part of the experiment, we are interested in the reliance of the model on the recurrent state. Recall that our model maintains the mask accumulating the previous predictions as well as the hidden LSTM state. We alternately “block” either of the states by providing a zero tensor at every timestep. We consider only the first $n$ predictions to compute the Dice score, where $n$ is the number of ground-truth masks. We stop the iterations if no termination was predicted after 21 timesteps, since the largest number of instances in our validation set is 20. The results in Table 2 show that the LSTM plays an important role for counting (or, termination prediction), while having almost no effect on the mask quality. The networks have learned a sequential prediction strategy given only the binary mask of previously predicted pixels. Note that in contrast to the baseline models, actor-critic training reduced the dependence on the LSTM state for counting (AC-Dice), which suggests that the actor makes a better use of the state mask to make the next prediction.

6.2 Instance segmentation

We compare our method to other approaches on two standard instance segmentation benchmarks, each containing a rich variety of small segments as well as occlusions.

CVPPP dataset.  For the CVPPP dataset used in our ablation study, this time we evaluate on the official 33 test images and train only our actor-critic model (AC-Dice) on the total 128 images in the training set.

The results on the test set in Table 3 show that our method is on par with the state of the art in terms of counting while maintaining competitive segmentation accuracy. From a qualitative analysis, see examples in Fig. 4(a), we observe that the order of prediction follows a consistent, interpretable pattern: large leaves are segmented first, whereas small and occluded leaves are segmented later. This follows our intuition for an optimal processing sequence: “easy”, more salient instances should be predicted first to alleviate consecutive predictions of the “harder” ones. We also note, however, that the masks miss some fine details, such as the stalk of the leaves, which limits the benefits of the context for occluded instances. We believe this stems from the limited capacity of the critic network to approximate a rather complex reward function.

KITTI benchmark.  We use the instance-level annotation of cars in the KITTI dataset [10] to test the scalability of our method to traffic scenes. We used the same data split as in previous work [31, 35], which provides 3712 images for training, 144 images for validation, and 120 images for testing. While the validation and test sets have high-quality annotations [6, 41], the ground-truth masks in the training set are largely ($>95\%$) coarse or incomplete [6]. Hence, good generalisation from the training data would indicate that the algorithm can cope well with inaccurate ground-truth annotation.

The evaluation criteria for this dataset are: the mean weighted coverage loss (MWCov), the mean unweighted coverage loss (MUCov), the average false positive rate (AvgFP), and the average false negative rate (AvgFN). MUCov is the maximum IoU of the ground truth with a predicted mask, averaged over all ground-truth segments in the image. MWCov additionally weighs the IoUs by the area of the ground-truth mask. AvgFP is the fraction of mask predictions that do not overlap with the ground-truth segments. Conversely, AvgFN measures the fraction of the ground-truth segments that do not overlap with the predictions.

We use an IoU-based score function to compute the rewards, i.e. $\mathcal{F}(\mathcal{S},\mathcal{T})=\frac{\sum_{i}{\mathcal{S}_{i}\mathcal{T}_{i}}}{\sum_{i}\mathcal{S}_{i}+\sum_{i}\mathcal{T}_{i}-\sum_{i}{\mathcal{S}_{i}\mathcal{T}_{i}}}$. To show the benefits of our Actor-Critic model (AC-IoU) for structured prediction at higher resolutions, we also train and report results for a baseline, the actor-only model trained with one step BPTT (BL-Trunc). Considering the increased variability of the dataset compared to CVPPP, we used 64 latent dimensions for the action space.

The results on the test are shown in Table 4(a). Given the relatively small size of the test set, we also report the results on the validation set in Table 4(b), and use the available results from the equivalent evaluation of a state-of-the-art method [31] for reference.

The results indicate that our method scales well to larger resolutions and action spaces and shows competitive accuracy despite not using bounding box representations. Similar to our results on CVPPP, our model does not quite reach the accuracy of a recurrent model using bounding boxes [31] and a non-recurrent pipeline. We believe the segmentation accuracy is currently limited by the degree of the reward approximation by the critic and the representational power of the network architecture used by the actor model. As can be seen in some examples in Fig. 4(b), without post-processing the masks are not always well aligned with the object and occlusion boundaries. However, we note that the prediction order also follows a consistent, interpretable pattern: nearby instances are segmented first, while far-away instances are segmented last. Without hard-coding such constraints, the network appears to have learned a strategy that agrees with human intuition to segment larger, close-by objects first and exploits the resulting context to make predictions in the order of increasing difficulty.

7 Conclusions

In the current study, we formalised the task of instance segmentation in the framework of reinforcement learning. Our proposed actor-critic model utilises exploration noise to alleviate the initialisation bias on the prediction ordering. Considering the high dimensionality of pixel-level actions, we enabled exploration in the action space by learning a low-dimensional representation through a conditional variational auto-encoder. Furthermore, the critic approximates a reward signal that also accounts for future predictions at any given timestep. In our experiments, it attained competitive results on established instance segmentation benchmarks and showed improved segmentation performance at the later timesteps. Our model predicts instance masks directly at the full resolution of the input image and does not require intermediate bounding box predictions, which stands in contrast to proposal-based architectures [13] or models delivering only a preliminary representation for further post-processing, e.g.  [9, 35].

These encouraging results suggest that actor-critic models have potentially a wider application spectrum, since the critic network was able to learn a rather complex loss function to a fair degree of approximation. In future work, we aim to improve our baseline model of the actor network, which currently limits the attainable accuracy.

Acknowledgements.  The authors would like to thank Stephan R. Richter for helpful discussions.

Appendix A Ordering of Prediction

As we discussed in the main text, previous work suggests that the overall accuracy of a recurrent model is not invariant to the prediction order \citesuppli2017not,OrderMatters. To confirm this experimentally for instance segmentation, we decouple the localisation and the segmentation aspect of a recurrent model in an oracle experiment. Reserving the location of the ground-truth segments as the oracle knowledge (e.g., with a bounding box) allows us to control the prediction order and study its role for the quality of the pixelwise prediction.

The setup for our oracle experiment is illustrated in Fig. 5. We supply the image and context patches (4 input channels) to a segmentation network in random order, while at the same time accumulating its predictions into the (global) context mask. We repeat this procedure 20 times for every image, with a different random order each. We use the CVPPP train/val split from our ablation study, and report the results on the 25 images in the validation set. The architecture of the network is the same as the segmentation module used by Ren and Zemel \citesuppren2016end.

Figure 6(a) shows the statistic of the results (mean and standard deviation) for each of the 25 images. We observe that there is a tangible difference between the different segmentation orderings. The mean of the maximum and the minimum Dice across the images are $86.5$ and $78.1$ Dice, respectively, which corresponds to a potential gain of $8.4$ mean Dice for the optimal prediction order. Notably, the deviation between the runs varies depending on the image. To investigate this further, we show two examples from the validation set with the highest (index #1, Fig. 6(b)) and the lowest (index #19, Fig. 6(c)) deviation. Example #1 appears to contain many more occlusions and more complex segment shapes (e.g., a single stalk) compared to example #19. This aligns well with our expectation: the benefit of the context should be more pronounced in more complex scenes.

Appendix B Architecture

The architecture of the actor network and its parameters depending on the dataset used are shown in Tables 5 and 6. The critic has a similar architecture, but makes use of batch normalisation (BN) \citesuppIoffeS15. Additionally, one fully-connected layer (FC) from the last convolution is replaced by global maximum pooling. Further details are summarised in Table 7.

We use a vanilla LSTM \citesuppgraves2005framewise without peephole connections to represent the hidden state due to its established performance \citesuppgreff2017lstm and more economic use of memory compared to convolutional LSTMs \citesuppromera2016recurrent,xingjian2015convolutional.

The pre-processing network, which predicts the foreground and the angle quantisation of instances, is an FCN \citesupplong2015fully with the same architecture as used by Ren and Zemel \citesuppren2016end.

Appendix C Training

To train our actor-critic model, we empirically found that gradually increasing the maximum sequence length leads to faster convergence than training on the complete length from the start. Specifically, we train the actor to predict $n$ remaining masks, where $n$ starts with $1$ and is gradually extended until the maximum number of instances in the dataset plus the terminations state (21 for CVPPP and 16 for KITTI) is reached. For example, with $n=1$ the model learns to predict the last instance, i.e. the initial state $s_{1}$ input to the network is an accumulation of all but one ground-truth masks on a given image, which we choose randomly. We extend $n$ by 5 once the segmentation accuracy on the validation set (either Dice or IoU) stops improving for several epochs.

We decrease the learning rate by a factor of 10 if we observe either no improvement or high variance of the results on the validation set. To trade-off the critic’s gradient, we used a constant scalar $\beta_{\text{act}}=10^{-3}$ for the KL-divergence loss. The actor and critic networks use weight decay regularisation of $10^{-5}$ and $10^{-4}$, respectively. We use Adam \citesuppkingma14adam for optimisation, as our experiments with SGD lead to considerably slower convergence. To manage training time, we downscale the original images with bilinear interpolation. The ground-truth masks are scaled down using nearest-neighbour interpolation. For CVPPP, we reduce the original resolution from $530\times 500$ to $224\times 224$, and for KITTI from the original $1242\times 375$ to $768\times 256$. We implemented our method in Lua Torch \citesuppCollobert2011Torch7AM. Both pre-training and training were performed on a single Nvidia Titan X GPU. The code is available at https://github.com/visinf/acis/.

Appendix D Qualitative Examples

We provide qualitative segment-by-segment visualisations in Figs. 7 and 8 from the CVPPP and KITTI validation sets. Supporting our analysis in the main paper, the order of prediction exhibits a consistent pattern. Furthermore, we observe that our model copes well with inaccuracies of intermediate predictions – a common failure mode of recurrent networks \citesuppRanzato:2016:SLT.

\bibliographystylesupp

ieee_fullname \bibliographysuppegbib_supp