Learning Transferable Reward for Query Object Localization with Policy Adaptation

We propose a reinforcement learning based approach to query object localization , for which an agent is trained to localize objects of interest specified by a small exemplary set . We learn a transferable reward signal formulated using the exemplary set by ordinal metric learning . Our proposed method enables test-time policy adaptation to new environments where the reward signals are not readily available , and outperforms fine-tuning approaches that are limited to annotated images . In addition , the transferable reward allows repurposing the trained agent from one specific class to another class . Experiments on a corrupted MNIST , the CU-Birds , and the COCO datasets demonstrate the effectiveness of our approach . 1 INTRODUCTION . There is increasing interest in designing machine learning models that can be adapted to unseen environments or repurposed for new tasks , with only minimal human guidance ( Vinyals et al. , 2016 ; Snell et al. , 2017 ; Finn et al. , 2017 ; Sun et al. , 2020 ; Hansen et al. , 2021 ) . To this end , a small set of examples can not only serve the purpose of implicitly generalizing a trained model to a new environment during training time , but also enable the learner to update its learning objectives during test time . In this paper , we focus on a reinforcement learning ( RL ) formulation to the problem of query object localization , in which an agent is trained to localize the target object specified by a small set of exemplary images . Instead of using fixed bounding box proposals , our vision-based agent can be viewed as a proactive information gatherer ( Guo , 2003 ) , which actively interacts with an image environment ; Furthermore , it follows a class-specific localization policy , and thus is more suitable for aerial imagery ( Xia et al. , 2018 ) , robotic manipulation ( Kalashnikov et al. , 2018 ) , or embodied AI ( Savva et al. , 2019 ) tasks . During the test time of query object localization , the queried object class to localize may be novel , or the background environment may undergo substantial changes , hindering the applicability of class-agnostic agents with a fixed policy . In addition , in standard RL settings for object localization , the reward signal is often available , so fine-tuning methods ( Julian et al. , 2020 ) can effectively adapt agents to new environments and yield improved performance ; On the contrary , in query object localization , the reward signal is not available at test time , because the bounding box annotations are to be found by the localization agent on test images . To address these problems , we propose an ordinal metric learning based framework for learning an implicitly transferable reward signal defined with a small exemplary set . An ordinal embedding network is pre-trained with data augmentation under a loss function designed to be relevant to the RL task . The reward signal allows explicit updates of the controller in the policy network with continual training during test time . Compared to fine-tuning approaches , the agents can get exposed to new environments more extensively with unlimited usage of test images . Informed by the exemplary set precisely , the agent is adapted to the changes of the localization target . Our contributions in this paper are summarized as follows : ( 1 ) We propose a novel ordinal metric learning based RL framework to learn a transferable reward signal defined with a small exemplary set for query object localization ; ( 2 ) Our RL framework enables test-time policy adaptation to new environments where there is no readily available reward signal ; ( 3 ) Our learned transferable reward allows repurposing a trained agent from localizing one specific class to another ; ( 4 ) Extensive experiments on several datasets demonstrate the effectiveness of our proposed approach . 2 RELATED WORK . 2.1 OBJECT LOCALIZATION . Compared to bounding-box regression approaches ( Redmon et al. , 2016 ; Ren et al. , 2016 ) , deep RL based object localization approaches ( Caicedo & Lazebnik , 2015 ; Jie et al. , 2016 ) have the advantage of being region-proposal free , with customized search paths for each image environment . The specificity of an agent purely depends on the classes of bounding-boxes used in the reward . They can be made class-specific ( Caicedo & Lazebnik , 2015 ) , but the agent for each class would need to be trained separately . Despite the rise of crowdsourcing platforms , obtaining ample amount of bounding-box annotations remains costly and error-prone . Furthermore , the quality of annotations often varies , and precise annotations for certain object class may require special expertise from annotators . The emergence of weakly supervised object localization ( WSOL ) ( Song et al. , 2014 ; Oquab et al. , 2015 ; Zhou et al. , 2016 ) methods alleviates the situation , in which image class labels are used to derive bounding box annotations . It is known that WSOL methods have drawbacks of overly relying on inter-class discriminative features and failing to generalize to classes unseen during the training phase . We argue that intra-class similarity is a more natural objective for the problem of localizing objects belonging to a target class . `` Query object localization '' emphasizes more about the class-specific nature of a localization task , with the relevance defined by the similarity to query images . A similar problem is image co-localization ( Tang et al. , 2014 ; Wei et al. , 2017 ) , in which the task is to identify the common objects within a set of images . Co-localization approaches exploit the common characteristics across images to localize objects . Being unsupervised , co-localization approaches could suffer from ambiguity if there exist multiple common objects or parts , e.g. , bird head and body ( Jaderberg et al. , 2015 ) , which may provide unwanted common objects as output . 2.2 POLICY ADAPTATION . There seems to be a contradiction between the goals of training an agent with high task-specificity and better generalization performance to new situations at the same time . The key to reconcile these two goals lies in the usage of a small set of examples . There has been a paradigm shift from training static models defined with parameters to ones defined together with a support set ( Vinyals et al. , 2016 ; Snell et al. , 2017 ) , proved to be very effective in few-shot classification . Besides the efforts of employing meta learning to adjust models , fine-tuning on a pretrained model ( Oquab et al. , 2014 ; Finn et al. , 2017 ; Julian et al. , 2020 ) has also been widely used in transferring knowledge from data-abundant to data-scarce tasks . When reward signal is not available , Hansen et al . ( 2021 ) proposed a policy adaptation approach , in which the intermediate representation is fine-tuned via optimizing a self-supervised auxiliary loss while the controller is kept fixed . Our work shares the same motivation of test-time training ( Sun et al. , 2020 ) , but we focus instead on the settings where the controller needs to be adapted or even repurposed for new tasks . 3 LEARNING TRANSFERABLE REWARD FOR LOCALIZING QUERY OBJECTS . 3.1 RL FORMULATION FOR QUERY OBJECT LOCALIZATION . Given a set of images I with known class labels , the task of query object localization is to find the location b of the bounding box most relevant to the query object , specified in a small set of exemplary images E . Following the Markov Decision Process ( MDP ) framework for object localization ( Caicedo & Lazebnik , 2015 ; Jie et al. , 2016 ) , this task can be formulated as a RL problem . In the training phase , a localization agent learns to take actions of moving the bounding box to maximize a reward function reflecting localization accuracy . In the test phase , the agent is informed by a small test exemplary set about a new task objective . Figure 1 shows an overview of the query object localization problem and the RL formulation , with components listed as follows . • Environment : Raw pixels of a single image Ii , without candidate proposal boxes . • Action : Discrete actions facilitating a top-down search , including five scaling : scale towards top left/right , bottom left/right , center ; eight translation transformations : move left/right/up/down , enlarge width/height , reduce width/heigh ; plus one stay action . • State : Pooled features from the current box and an internal state within RNN that encodes information from history observations . • Reward : Improvements the agent made in localizing the query images ( in different tasks ) . 3.2 PRETRAINING ORDINAL REWARD SIGNAL . For training an agent to localize different objects according to queries , the reward function needs to be transferable . Existing deep RL approaches for object localization ( Caicedo & Lazebnik , 2015 ; Jie et al. , 2016 ) use its ground-truth object bounding box gi as the reward signal , R = sign ( IoU ( bt , gi ) − IoU ( bt−1 , gi ) ) , ( 1 ) where IoU ( bt , gi ) denotes the Intersection-over-Union ( IoU ) between the current window bt and the corresponding ground-truth box gi1 , and IoU ( b , g ) = area ( b ∩ g ) /area ( b ∪ g ) . Similar to the bounding box regression approaches which learns a mapping φ : I 7→ g , the image and box must be paired . However , annotated image-box pairs ( I , g ) may be scarce in both the training and testing phases . The reward signal defined in Eq . 1 is instance-wise , which is not transferable across images . To address this problem , a natural idea is to define the reward signal based on learned representations of the cropped images by current window bt and the ground-truth window g. Given their M - dimensional representations bt and g produced by an embedding function f : RD 7→ RM from D- dimensional image feature vectors , a distance function d : RM × RM 7→ [ 0 , +∞ ) returns the embedding distance d ( bt , g ) . However , an embedding distance based on off-the-shelf pretrained networks is not adequate , since it may not decrease monotonically as the agent approaches to the ground-truth box g. As a result , the embedding distance based reward signal may be less effective than Eq . 1 . Therefore , an embedding function customized for the downstream RL-based localization task is needed . Ordinal metric learning . We propose to use an ordinal embedding based reward signal . For any two perturbed boxes bj , bk from ground truth g , embeddings bj , bk , g are learned , such that the relative preferences between any pair of boxes are preserved in the Euclidean space , ρj > ρk ⇔ ||bj − g|| < ||bk − g|| , ∀j , k ∈ C , ( 2 ) where ρj and ρk denote the preference ( derived from either IoU to ground-truth box or ordinal feedbacks from user ) , and C is an ordinal constraint set constructed by sampling box pairs around g. This problem is originally posed as non-metric multidimensional scaling ( Agarwal et al. , 2007 ; Jamieson & Nowak , 2011 ) . Although we apply a very simple pairwise-based approach , there exist other extensions such as the listwise-based ( Cao et al. , 2007 ) , quadruplet-based ( Terada & Luxburg , 2014 ) and landmark-based ( Anderton & Aslam , 2019 ; Ghosh et al. , 2019 ) approaches . 1Notation : bounding box are denoted by lowercase lightface letters , e.g. , g , while embedding vectors are denoted by boldface letters , e.g. , g ∈ Rd is a d-dimensional vector . Loss function . In this paper , we define preference ρ as the IoU of box b to the ground-truth box g , i.e. , ρ = IoU ( b , g ) . We choose to optimize a triplet loss ( Hermans et al. , 2017 ) for learning the desired embeddings , Ltriplet = ∑ anchor , ordinal pairs max ( m+ d ( a , p ) − d ( a , n ) , 0 ) , ( 3 ) where a is the “ anchor '' embedding that will be discussed in detail later , p , n are the “ positive '' and “ negative '' embeddings for boxes with larger and smaller IoUs with ground truth box g , respectively , m is a margin hyperparameter . We learn an embedding space consistent with local ordinal constraints on positive and negative data obtained via data augmentation such as box perturbation . Box Perturbation Assuming the training exemplary set Etrain contains both image I and box g , we adopt a tailored data augmentation scheme - box perturbation , in which a pair of boxes are randomly sampled from I with different IoUs to g. We compare the efficiency of sampling schemes on generating augmented bounding box pairs , ( a ) Random sampling , where the pair of boxes are generated completely randomly , and ( b ) Group-based sampling , where dense boxes with variant scales are first generated , and then divided into M = 10 groups according to the IoU with groundtruth box . Each group has an IoU interval of 1/M . The sampling is first done on the group level , i.e. , 2 out of M , then sample one box from each group . Thus , the sampled boxes are likely to cover more diverse cases comparing to random sampling . We have found that using IoU-based partition scheme is more effective than random sampling . The choice of anchor . The anchor embedding a in Eq . 3 is not restricted to the cropped image from the same image . For example , it could be replaced by the prototype embedding ( Snell et al. , 2017 ) of the exemplary set E , c = 1/|E| ∑ i∈E gi , where gi is the embedding of the cropped image Ii by ground-truth box gi . If images from multiple classes are available , the prototype can be further made to be class-dependent , or clustering-based . A prototype as the anchor enables test-time policy adaptation without image-box pairs in the exemplary set Etest . In comparison , supervised methods such as faster-RCNN would require image-box pairs for fine-tuning on new tasks . We find that using prototype-based embedding as the anchor leads to better generalization performance than using g in some experiments . It is also interesting to explore other treatment on the exemplary set , such as permutation-invariant architectures ( Zaheer et al. , 2017 ; Ilse et al. , 2018 ; Lee et al. , 2019 ) .

Summary: The paper proposes a reinforcement learning approach for localization in images. The idea is to consider and environment where the states are crops of an image, the actions are changes in the current coordinates of the current crops starting at the full image. The method consists of learning both a reward function and a policy that maximizes that reward function. The policy is learnt with REINFORCE with entropy regularisation. The manner in which the reward function is formulated and learnt is quite appealing. To allow for more flexibility at test time, the reward function is made conditional on some exemplars and learnt in a few shot manner based on state features. Using a triplet loss to ensure the reward function is increasing as the overlap is increasing is quite good.  The experiments are quite extensive and I find compelling for the most part. Ablation of the elements of the method is well executed and benefits of few short learning are compelling enough in my opinion. The results on COCO are a bit limited but I would be willing to overlook that.