nec_agent.py 文件源码

def _search_ann(self, search_keys, dnd_keys, update_LRU_order):
        batch_indices = []
        for act, ann in self.anns.items():
            # These are the indices we get back from ANN search
            indices = ann.query(search_keys)
            log.debug("ANN indices for action {}: {}".format(act, indices))
            # Create numpy array with full of corresponding action vector index
            action_indices = np.full(indices.shape, self.action_vector.index(act))
            log.debug("Action indices for action {}: {}".format(act, action_indices))
            # Riffle two arrays
            tf_indices = self._riffle_arrays(action_indices, indices)
            batch_indices.append(tf_indices)
            # Very important part: Modify LRU Order here
            # Doesn't work without tabular update of course!
            if update_LRU_order == 1:
                _ = [self.tf_index__state_hash[act][i] for i in indices.ravel()]
        np_batch = np.asarray(batch_indices)
        log.debug("Batch update indices: {}".format(np_batch))

        # Reshaping to gather_nd compatible format
        final_indices = np.asarray([np_batch[:, j, :, :] for j in range(np_batch.shape[1])], dtype=np.int32)

        return final_indices