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: This paper extends works performing RL-based object localization by: 1. Conditioning the localization on a exemplary set of images, instead of more classical hardcoded finite set of classes 2. Introducing a new reward signal which doesn’t explicitly use IoU but instead enforces a rank-preserving metric space (where distances reflect how well 2 bounding boxes match up).  This is shown on multi-MNIST data, CUB and COCO. They are targeting rather challenging generalisation setups, where a policy is only trained to localize 1 digit class (in the MNIST case say) and extrapolates to all other ones.