SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anindya Sarkar Nathan Jacobs Yevgeniy Vorobeychik  
{anindya, jacobsn, yvorobeychik}@wustl.edu, Department of Computer Science and Engineering Washington University in St. Louis
Visual active search (VAS) has been proposed as a modeling framework in which visual cues are used to guide exploration, with the goal of identifying regions of interest in a large geospatial area. Its potential applications include identifying hot spots of rare wildlife poaching activity, search-and-rescue scenarios, identifying illegal trafficking of weapons, drugs, or people, and many others. State of the art approaches to VAS include applications of deep reinforcement learning (DRL), which yield end-to-end search policies, and traditional active search, which combines predictions with custom algorithmic approaches. While the DRL framework has been shown to greatly outperform traditional active search in such domains, its end-to-end nature does not make full use of supervised information attained either during training, or during actual search, a significant limitation if search tasks differ significantly from those in the training distribution. We propose an approach that combines the strength of both DRL and conventional active search by decomposing the search policy into a prediction module, which produces a geospatial distribution of regions of interest based on task embedding and search history, and a search module, which takes the predictions and search history as input and outputs the search distribution. We develop a novel meta-learning approach for jointly learning the resulting combined policy that can make effective use of supervised information obtained both at training and decision time. Our extensive experiments demonstrate that the proposed representation and meta-learning frameworks significantly outperform state of the art in visual active search on several problem domains.
Consider a scenario where a child is abducted and law enforcement needs to scan across hundreds of potential regions from a helicopter for a particular vehicle. An important strategy in such a search and rescue portfolio is to obtain aerial imagery using drones that helps detect a target object of interest (e.g., the abductor's car) [1, 2, 3, 4, 5]. The quality of the resulting photographs, however, is generally somewhat poor, making the detection problem extremely difficult. Moreover, security officers can only inspect relatively few small regions to confirm search and rescue activity, doing so sequentially.
We can distill some key generalizable structure from this scenario: given a broad area image (often with a relatively low resolution), sequentially query small areas within it (e.g., by sending security officers to the associated regions, on the ground) to identify as many target objects as possible. The number of queries we can make is typically limited, for example, by budget or resource constraints. Moreover, query results (e.g., detected search and rescue activity in a particular region) are highly informative about the locations of target objects in other regions, for example, due to spatial correlation. We refer to this general modeling framework as visual active search (VAS). Numerous other scenarios share this broad structure, such as identification of drug or human trafficking sites, anti-poaching enforcement activities, identifying landmarks, and many others. Sarkar et al. [6]
recently proposed a visual active search (VAS) framework for geospatial exploration, as well as a deep reinforcement learning (DRL) approach for learning a search policy. However, the efficacy of this DRL approach remains limited by its inability to adapt to search tasks that differ significantly from those in the training data, with the DRL approach struggling in such settings even when combined with test-time adaptation methods.
To gain intuition about the importance of domain adaptivity in VAS, consider Figure 1. Suppose, we pre-train a policy by leveraging a fully annotated search tasks with large vehicle as a target class to learn a search policy. Now the challenge in visual active search problem is how to utilize such a policy for our current task, that is to search for small car given a broad aerial image. As depicted in Figure 1, an adaptive policy (right) initially makes a mistake but then quickly adapts to the current task by efficiently leveraging information obtained in response to queries to learning from its mistakes. In contrast, a non-adaptive policy (left) keeps repeating its
Figure 1: Comparative search strategy of non-adaptive and adaptive policy with small car as a target.
mistakes and ultimately fails to find the region containing the target object.
Indeed, traditional active search approaches have been designed precisely with such adaptivity in mind [7, 8, 9, 10] by combining an explicit machine learning model that predicts labels over inputs with custom algorithmic approaches aiming to balance exploration (which improves prediction efficacy) and exploitation (to identify as many actual objects of interest as possible within a limited budget). However, as Sarkar et al. [6] have shown, such approaches perform poorly in VAS settings compared to DRL, despite the lack of effective domain adaptivity of the latter.
We combine the best of traditional active search and the state-of-the-art DRL-based VAS approach by developing a partially-supervised reinforcement learning framework for VAS (PSVAS). A PSVAS policy architecture is comprised of two modules: (i) A task-specific prediction module that learns to predict the locations of the target object based on the task aerial image and labels resulting from queries during the search process, and (ii) a task-agnostic search module that learns a search policy given the predictions provided by the prediction module and prior search results. The key advantage of this decomposition is that it enables us to use supervised information observed at decision time to update the parameters of the prediction module, without changing the search module. Furthermore, to learn search policies that are effective in the context of evolving predictions during the search, we propose a novel meta-learning approach to jointly learning the search module with the initialization parameters of the prediction module (used at the beginning of each task). Finally, we generalize the original VAS framework to allow for multiple simultaneous queries (a common setting in practice), and develop the PSVAS framework and a meta-learning approach for such settings.
In summary, we make the following contributions:
- A novel partially supervised reinforcement learning (PSVAS) framework that enables effective adaptation of VAS search policies to out-of-distribution search tasks.
- A novel meta-learning approach (MPS-VAS) to learn initialization parameters of the prediction module jointly with a search policy that is robust to the evolving predictions during a search task.
- A generalization of the VAS problem to domains in which we can make multiple simultaneous queries, and a variant of MPS-VAS that learns how to choose a subset of queries to make in each search iteration.
- An extensive experimental evaluation on two publicly available satellite imagery datasets, xView and DOTA, in a variety of unknown target settings, demonstrating that the proposed approaches significantly outperform all baselines. Our code is publicly available at this link.
We consider a generalization of the visual active search (VAS) problem proposed by Sarkar et al. [6]. The basic building block of VAS is a task, which centers around an aerial image  $x$  divided into
$N$  grid cells, so that  $x = (x^{(1)}, x^{(2)}, \ldots, x^{(N)})$ , with each grid cell a subimage. Broadly, the goal is to identify as many target objects of interest through iterative exploration of these grid cells as possible, subject to a budget constraint  $\mathcal{C}$ . To this end, we represent the subset of grids containing the target object by associating each grid cell  $j$  with a binary label  $y^{(j)} \in \{0, 1\}$ , where  $y^{(j)} = 1$  if the grid cell  $j$  contains the target object, and 0 otherwise. The complete label vector associated with the task is  $y = (y^{(1)}, y^{(2)}, \ldots, y^{(N)})$ . When we are faced with the task at decision time, we have no direct information about  $y$ , but when we query a grid cell  $j$ , we obtain the ground truth label  $y^{(j)}$  for this cell. Moreover, if  $y^{(j)} = 1$ , we also accrue utility from exploring  $j$ . In the original variant of VAS, we can make a single query  $j$  in each time step. Here, we consider a natural generalization where we have  $R$  query resources (for example,  $R < N$  patrol units identifying traps in a wildlife conservation setting), so that we can make  $R$  queries in each time step. We assume that  $R$  is constant for convenience; it is straightforward to generalize our approach below if  $R$  is time-varying.
Let  $c(j, k)$  be the cost of querying grid cell  $k$  if we start in grid cell  $j$ . For the very first query, we can define a dummy initial grid cell  $d$ , so that cost function  $c(d, k)$  captures the initial query cost. Let  $q_{t}^{r}$  denote the set of queries performed in step  $t$  by a query resource  $r$ . Our ultimate goal is to solve the following optimization problem:
$$
\max  _ {\{q _ {t} ^ {r} \}} U (x; \{q _ {t} ^ {r} \}) \equiv \sum_ {t} y ^ {(q _ {t} ^ {r})}
$$
$$
\text {s . t .}: \sum_ {t \geq 0} \sum_ {r = 1} ^ {R} c \left(q _ {t - 1} ^ {r}, q _ {t} ^ {r}\right) \leq \mathcal {C}. \tag {1}
$$
Finally, we assume that we possess a collection of tasks (in this case, aerial images) for which we have annotated whether each grid cell contains the target object or not. This collection, which we will refer to as  $\mathcal{D} = \{(x_i,y_i)\}$ , is comprised of images  $x_{i}$  with corresponding grid cell labels  $y_{i}$ , where each  $x_{i}$  is composed of  $N$  elements  $(x_{i}^{(1)},x_{i}^{(2)},\ldots ,x_{i}^{(N)})$  representing the cells in the image, and each  $y_{i}$  contains  $N$  corresponding labels  $(y_{i}^{(1)},y_{i}^{(2)},\ldots ,y_{i}^{(N)})$ .
The central technical goal in VAS is to learn an effective search policy that maximizes the total number of targets discovered, on average, for a sequence of tasks  $x$  on which we have no prior direct experience, given a set of resources  $R$ , exploration cost function  $c(j,k)$ , and total exploration budget  $C$ . Sarkar et al. [6] proposed a deep reinforcement learning (DRL) approach in which they learned a search policy  $\pi(x,o,B)$  that outputs at each time step  $t$  the grid we should explore at the next step  $t + 1$ , given the task  $x$ , remaining budget  $B$ , and information produced from the sequence of previous queries encoded into an observation vector  $o$  with  $o^{(j)} = 2y^{(j)} - 1$  if  $j$  has been explored, and  $o^{(j)} = 0$  otherwise. The reward function for this DRL approach is naturally captured by  $R(x,o,j) = y^{(j)}$ .
Note that active search (including VAS) is qualitatively distinct from active learning [11, 12, 13]: in the latter, the goal is solely to learn to predict well, so that the entire query process serves the goal of improving predictions. In active search, in contrast, we aim to learn a search policy that balances exploration (improving our ability to predict where target objects are) and exploitation (actually finding such objects) within a limited budget. Indeed, it has been shown that active learning approaches are not competitive in the VAS context [6].
3 Partially-Supervised Reinforcement Learning Approach for VAS
The DRL approach for VAS proposed by Sarkar et al. [6] is end-to-end, producing a policy that is only partly adaptive to observations  $o$  made during exploration for each task. In particular, this end-to-end policy aims to capture the tension between exploration and exploitation fundamental in active search [7, 9], without explicitly representing the central aim of exploration, which is to improve our ability to predict which grid cells in fact contain the target object. In conventional active search, in contrast, this is directly represented by learning a prediction function  $f(x^{(j)})$  that is updated each time a grid  $j$  is explored during the search process, so that exploration directly impacts our ability to predict target locations, and thereby make better query choices in the future. While the end-to-end approach implicitly learns this, it would only do so effectively so long as tasks  $x$  we face at prediction time are closely related to those used in training. However, it does not take advantage of the supervised information we obtain about which grid cells actually contain the target object, either at training or
Figure 2: The PSVAS policy network architecture.
test time, potentially reducing the efficacy of learning as well as ability to adapt when task distribution changes.
To address this issue, we propose a novel partially-supervised reinforcement learning approach for visual active search, which we refer to as PSVAS. In PSVAS, the search policy we obtain is a composition of two modules: 1) the task-specific prediction module  $f_{\theta}(x,o)$  and 2) the task-agnostic search module  $g_{\phi}(p,o,B)$ , where  $\theta$  and  $\phi$  are trainable parameters, where  $p = f_{\theta}(x,o)$  is the vector of predicted probabilities with  $p^{(j)}$  the predicted probability of a target in grid cell  $j$  (see Figure 2, and Supplement Section A for further details about the policy network architecture). Conceptually,  $f_{\theta}$  makes predictions based solely on the task  $x$  and the prediction-relevant information gathered during the search  $o$ , while  $g$  relies solely on the information relevant to the search itself: predicted locations of objects  $p$ , observations acquired during the search  $o$ , and remaining budget  $B$ . The search policy is then the composition of these modules,  $\pi(x,o,B) = g_{\phi}(f_{\theta}(x,o),o,B)$ .
Although in principle we can still train the policy  $\pi$  above end-to-end (effectively collapsing the composition), a key advantage of PSVAS is that it enables us to directly make use of supervised information about true labels  $y$  as they are observed either at training or decision time, using these to update  $f_{\theta}$  in both cases. In prior work that relies solely on end-to-end policy representation, there was no straightforward way to take advantage of this information. Specifically, we train the policy using the following objective function:
$$
\mathcal {L} _ {\text {P A R V S}} = \mathcal {L} _ {R L} + \lambda \mathcal {L} _ {B C E},
$$
where  $\lambda$  is a hyperparameter. We represent  $\mathcal{L}_{RL}$  loss as follows:
$$
\nabla \mathcal {L} _ {R L} = - \sum_ {i = 1} ^ {M} \sum_ {t = 1} ^ {T _ {i}} \mathbb {1} _ {\sum_ {t \geq 0} c (q _ {t - 1}, q _ {t}) \leq \mathcal {C}} \nabla \log \pi_ {(\theta , \phi)} (a _ {i} ^ {t} | x _ {i}, o _ {i} ^ {t}, B _ {i} ^ {t}) \mathbb {R} _ {i} ^ {t}
$$
Where  $M$  is the number of example search task seen during training and  $\mathbb{R}_i^t$  is the discounted cumulative reward defined as  $\mathbb{R}_i^t = \sum_{k=t}^{T} \gamma^{k-t} R_i^k$  with a discount factor  $\gamma \in [0,1]$ . Note that the RL loss in this approach also updates the parameters of the prediction module,  $\theta$ , in an end-to-end fashion. We also represent  $\mathcal{L}_{BCE}$  as follows:  $\mathcal{L}_{BCE} = \sum_{i=1}^{M} - (y_i \log(p_i) + (1 - y_i) \log(1 - p_i))$ .
The PSVAS algorithm is more formally presented in Algorithm 1. The combined RL and supervised loss yields a balance between using supervised information to improve the quality of initial predictions at the beginning of the episode and ensuring that these serve the goal of producing the best episode-level policies. However, it is still crucial to note that  $f_{\theta}$  serves solely an instrumental role in the process, with learning a search policy that effectively adapts to each task the primary goal. Consequently, we ultimately wish to jointly learn  $g_{\phi}$  and  $f_{\theta}$  in such a way that  $f_{\theta}$  facilitates adaptive search as it is updated during a search task at decision time. To address this, we propose a meta-learning approach, MPS-VAS, that learns  $g_{\phi}$  along with an initial parametrization  $\theta_0$  of  $f_{\theta}$  for each task. At the beginning of each task, then,  $\theta$  is initialized to  $\theta_0$ , and then updated as labels are observed after each query.
At the high level, MPS-VAS trains over a series of episodes, where each episode corresponds to a fully labeled training task  $(x_{i},y_{i})$  and budget constraint  $\mathcal{C}$  (we vary the budget constraints during training). We begin an episode  $i$  with the current prediction function parameters  $\theta_{i} = \theta$  and search policy parameters  $\phi$ . We fix  $g_{\phi}$  over the length of the episode, and use it to generate a sequence of queries (since the policy is stochastic, it naturally induces exploration). After observing the label  $y^{(j)}$
Algorithm 1 The PSVAS algorithm.
Require: A search task instance  $(x_{i},y_{i})$ ; budget constraint  $\mathcal{C}$ ; Prediction function  $(f)$  with current parameters  $\theta_{i}$ , i.e.,  $f_{\theta_i}$ ; Search policy  $(g)$  with current parameters  $\phi_{i}$ , i.e.,  $g_{\phi_i}$ ;
1: Initialize  $o^0 = [0\dots 0];B^0 = \mathcal{C}$  ; step  $t = 0$
2: while  $B^t > 0$  do
3:  $\tilde{y} = f_{\theta_i}(x_i,o^t)$
4:  $j\gets Sample_{j\in \{UnexploredGrids\}}[g_{\phi_i}(\tilde{y},o^t,B^t)]$
5: Query grid cell with index  $j$  and observe true label  $y^{(j)}$
6: Obtain reward  $R^{t} = y^{(j)}$ , Update  $o^{t}$  to  $o^{t+1}$  with  $o^{(j)} = 2y^{(j)} - 1$ , and update  $B^{t}$  to  $B^{t+1}$  with  $B^{t+1} = B^{t} - c(k,j)$  (assuming we query  $k$ 'th grid at  $(t-1)$ ).
7: Collect transition tuple  $(\tau)$  at step t, i.e.,  $\tau^t = (\text{state} = (x_i, o^t, B^t), \text{action} = j, \text{reward} = R^t, \text{next state} = (x_i, o^{t+1}, B^{t+1}))$ .
8:  $t\gets t + 1$
9: end while
10: Update the prediction and search policy parameters, i.e.,  $\theta_{i}$  and  $\phi_{i}$  using  $(\mathcal{L}_{PSVAS})$  based on the collected transition tuples  $(\tau^{t})$  and the collected labels  $(y^{(j)})$  throughout the episode.
11: Return updated prediction and search policy parameters, i.e.,  $\theta_{i + 1}$  and  $\phi_{i + 1}$  respectively.
for each queried grid cell  $j$  during the episode, we update  $f_{\theta_i}$  using a standard supervised (binary cross-entropy) loss. At the completion of the episode (once we have exhausted the search budget  $\mathcal{C}$ ), we update the policy parameters, as well as the initialization prediction function parameters  $\theta$  by combining RL and supervised loss. For the search module, we use the accumulated sum of rewards  $R_i = \sum_j y^{(j)}$  over grids  $j$  explored during the episode, with the RL loss  $\mathcal{L}_{RL}$  based on the REINFORCE algorithm [14]. For the prediction module, we use the collected labels  $y^{(j)}$  during the episode and the standard binary cross-entropy loss. Finally, the MPS-VAS loss also explicitly trades off the RL and supervised loss:  $\mathcal{L}_{MLPARVS} = (\mathcal{L}_{RL} + \lambda \mathcal{L}_{BCE})$ . The proposed meta-learning approach thus explicitly trains the policy to account for the evolution of the prediction during the episode. The full MPS-VAS is presented more formally in Algorithm 2.
Algorithm 2 The MPS-VAS meta-learning algorithm.
Require: A search task instance  $(x_{i},y_{i})$ ; budget constraint  $\mathcal{C}$ ; Prediction function  $(f)$  with current parameters  $\theta_i^0 = \theta_i$ , i.e.,  $f_{\theta^0}$ ; Search policy  $(g)$  with current parameters  $\phi_i$ , i.e.,  $g_{\phi_i}$ ;
1: Initialize  $o^0 = [0\dots 0];B^0 = \mathcal{C}$  ; step  $t = 0$
2: while  $B^t > 0$  do
3:  $\tilde{y} = f_{\theta_1^t}(x_i,o^t)$
4:  $j\gets Sample_{j\in \{\text{Unexplored Grids}\}}[g_{\phi_i}(\tilde{y},o^t,B^t)]$
5: Query grid cell with index  $j$  and observe true label  $y^{(j)}$
6: Update  $\theta_i^t$  to  $\theta_i^{t + 1}$  using  $\mathcal{L}_{BCE}$  loss between  $\tilde{y}$  and pseudo label  $\hat{y}$ , defined as  $\hat{y} \gets \left\{ \begin{array}{ll} y^{(j)} & \text{if } y^{(j)} \text{ is Observed} \\ \tilde{y}^{(j)} & \text{if } y^{(j)} \text{ is Unobserved.} \end{array} \right.$
7: Obtain reward  $R^{t} = y^{(j)}$ , Update  $o^{t}$  to  $o^{t + 1}$  with  $o^{(j)} = 2y^{(j)} - 1$ , and update  $B^{t}$  to  $B^{t + 1}$  with  $B^{t + 1} = B^{t} - c(k,j)$  (assuming we query  $k$ 'th grid at  $(t - 1)$ ).
8: Collect transition tuple  $\tau$  at step t, i.e.,  $\tau^t = (\text{state} = (x_i, o^t, B^t)$ , action  $= j$ , reward  $= R^t$ , next state  $= (x_i, o^{t+1}, B^{t+1})$ ).
9:  $t\gets t + 1$
0: end while
11: Update search policy parameters  $\phi_{i}$  using  $(\mathcal{L}_{RL})$  based on the collected transition tuples  $(\tau^t)$  throughout the episode and update initial prediction function parameters  $\theta_{i}$  using  $(\mathcal{L}_{BCE})$  based on the collected labels  $(y^{(j)})$  throughout the episode.
12: Return updated prediction and search policy parameters, i.e.,  $\theta_{i + 1}$  and  $\phi_{i + 1}$  respectively.
The search policy  $\pi$  produces a probability distribution over grid cells. However, since in our setting no advantage can be gained by querying previously queried grid cells, we simply renormalize the probability distribution induced by  $\pi$  over the remaining grid cells, both during training and at decision time. Note that during inference, in both PSVAS and MPS-VAS framework, we freeze the parameters of the search module  $(\phi)$  and update the parameters of the prediction module  $(\theta)$  after
observing query outcomes at each step using  $\mathcal{L}_{BCE}$  loss between predicted label and pseudo label as shown in step (6) of Algorithm 2.
The discussion thus far effectively assumed that  $R = 1$ , that is, we make only a single query in each search time step. Next, we describe the generalization of our approach when we have multiple query resources  $R > 1$ . First, note that the prediction module  $f_{\theta}$  is unaffected, since the number of query resources only pertain to the actual search strategy  $g_{\phi}$ . One way to handle  $R$  queries is to simply sample the search module (which is stochastic) iteratively  $R$  times without replacement during training, and to greedily choose the most probable  $R$  grid cells to query at decision time. We refer to this as MPS-VAS-TOPK. However, since the underlying problem is now combinatorial (the choice of  $R$  queries out of  $N$  grid cells), such a greedy policy may fail to capture important interdependencies among such search decisions.
To address this, we propose a novel policy architecture which is designed to learn how to optimize such a heuristic greedy approach for combinatorial grid cell selection in a way that is non-myopic and accounts for the interdependent effects of sequential greedy choices. Specifically, let  $\psi$  be a vector corresponding to the grid cells, with  $\psi_{j} = 1$  if grid cell  $j$  has either been queried in the past (during previous search steps), or has been already chosen to be queried in the current search step, and  $\psi_{j} = 0$  otherwise. Thus,  $\psi$  encodes the choices that have already been made, and enables the policy to learn the best next sequential choice to make using a greedy procedure through the same RL framework that we described above. We refer to this approach as MPS-VAS-MQ (in reference to multiple queries).
Since the goal of active search is to maximize the number of target objects identified, we use average number of targets identified through exploration (ANT) as our evaluation metric.