Learning Where to Look: A Reinforcement Learning Framework for Robust Micro-Ultrasound Prostate Cancer Detection

Abstract: Micro-ultrasound ($μ$US) is a new, emerging, and promising imaging modality for prostate cancer (PCa) detection, but accurate identification of suspicious tissue remains highly dependent on clinical experience, leading to substantial inter-observer variability. Machine-learning assistance can reduce this variability; however, training reliable deep models is challenging because supervision is sparse and noisy -- typically limited to core-level histopathology outcomes (e.g., cancer grade and its percentage in a biopsy core) without pixel-level lesion annotations and under severe class imbalance. We introduce Prost-RL, which reframes $μ$US PCa detection as a spatially aware, policy-driven inference problem by learning where to look before decoding. Prost-RL integrates a lightweight reinforcement-learning policy into a foundation-model encoder-decoder to generate interpretable spatial attention maps that act as soft prompts for both cancer-likelihood heatmap prediction and image-level classification. We further propose Adaptive Policy Optimization (APO) to stabilize hybrid supervised-RL training and a noise-robust objective combining symmetric cross-entropy with negative-entropy regularization to mitigate weak-label noise and encourage sharp localization. On a cohort of 6,607 biopsy cores from 693 patients across five clinical sites, Prost-RL achieves $79.0\pm3.5$ AUROC with $64.6\pm6.3$% sensitivity at 80% specificity for core-level detection (+2.1 AUROC and +4.5 sensitivity points over the strongest baseline), and $79.3\pm5.8$ AUROC for clinically significant cancer classification. The learned policy highlights biopsy-aligned regions, providing transparent, spatially grounded evidence alongside quantitative risk predictions. Code is available at: https://github.com/DeepRCL/Prost-RL.

• The paper introduces Prost-RL, a reinforcement learning framework that learns a spatial attention policy to improve prostate cancer detection in micro-ultrasound imaging.
• The paper demonstrates statistically significant gains over existing methods, achieving improved AUROC and sensitivity for both core-level and clinically significant cancer detection.
• The paper’s modular approach with noise-robust loss functions yields interpretable and precise localization, effectively addressing challenges from weak, noisy, and imbalanced labels.

Reinforcement Learning for Policy-Driven Micro-Ultrasound Prostate Cancer Detection: The Prost-RL Framework

Introduction and Motivation

Micro-ultrasound (μUS) imaging offers increased spatial resolution for prostate cancer (PCa) detection but remains highly dependent on operator experience due to image artifacts, variable acoustic appearances, and predominance of benign tissue. Automated assistance is critical for standardizing image interpretation; however, pixel-level supervision is rarely available. Biopsy data provides only sparse, weak, and noisy core-level labels, leading to highly imbalanced datasets where benign tissue dominates. Standard models, even those using foundation models, treat the image holistically and can overfit to noisy global labels, resulting in diffuse, inaccurate predictions that are poorly localizable and lack transparency.

The presented work introduces Prost-RL, a reinforcement learning (RL)-enhanced framework designed to overcome the limitations of weak biopsy supervision by learning an explicit policy of “where to look” in μUS prior to inference. The architecture decouples the attention policy from downstream prediction heads, supports adaptive exploration via RL, and incorporates noise-robust objectives to enforce both localization sharpness and resilience to label noise. This method is validated on a multi-center, high-quality μUS dataset, establishing statistically significant performance gains over strong SOTA baselines for both core-level and clinically significant PCa detection, alongside producing interpretable evidence maps for clinical decision support.

Methodological Innovations

Prost-RL is architected as an encoder–decoder-agnostic module, instantiated atop ProstNFound+, which itself fuses pre-trained medical foundation models (MedSAM) with domain-specific adapters and clinical prompt encoders. The distinguishing features of Prost-RL are as follows:

1. Learned Spatial Attention Policy: A dedicated policy network generates a spatial probability distribution over encoder-derived features, modulated by clinical metadata, gating, and residual connections to form shared embeddings for both heatmap and classification heads (Figure 1).

   Figure 1: Prost-RL architecture integrates an image encoder, clinical prompt encoder, spatial attention policy, shared latent embedding, and dual decoding heads for heatmap and classification tasks.

2. Noise-Robust Loss Formulation: The training objective combines Symmetric Cross-Entropy (SCE) (balancing standard and reverse cross-entropy, with bounded gradients under label disagreement) and a negative-entropy regularizer to explicitly favor confident, spatially sharp predictions. This is crucial for overcoming degeneracy inherent in weak proportional labels, where average region-wise predictions can satisfy the loss identically to precise localization, unless further regularized.
3. Adaptive Reinforcement Learning (APO + DRPO): After supervised warm-up, the spatial policy is further refined using RL. A pairwise ranking-based reward is defined, prioritizing correctly ordering malignant and benign cases and emphasizing clinically significant cancers. Group-relative policy optimization strategies (with inspiration from GRPO and domain-robust policy optimization, DRPO) are incorporated, employing hierarchical temperature scaling to avoid gradient collapse under severe class imbalance and heterogeneous instance difficulty. Stochastic exploration is enabled by injecting Gaussian noise into attention logits during rollouts, avoiding degenerate deterministic gradients.

Ablative analysis confirms these contributions are complementary: spatial attention alone confers clear benefits, and the RL refinement further increases specificity and sensitivity with particular impact in the regime of low tumor involvement and label uncertainty.

Experimental Results

Prost-RL is evaluated on a cohort of 6,607 biopsy cores from 693 patients (5 clinical sites, five-fold cross-validation, patient-level splits). Hard exclusion of ambiguous negative cores and anatomically informed prompt encoding is employed. Downstream metrics include AUROC and sensitivity at fixed specificity for core-level detection (heatmap activation in needle-trace) and for image-level clinically significant PCa.

Headline results include:

• Core-level detection: $79.0 \pm 3.5$ AUROC, $64.6 \pm 6.3\%$ sensitivity at 80% specificity, representing +2.1 AUROC and +4.5 percentage point improvements in sensitivity over the strongest baseline (ProstNFound+).
• Clinically significant PCa classification: $79.3 \pm 5.8$ AUROC; sensitivity at 80% specificity: $62.8 \pm 12.6\%$.

These improvements consistently outpace MedSAM-based and prior prompt-supervised SOTA models across all benchmarks (Figure 2).

Figure 2: Comparative core-level AUROC as a function of cancer involvement threshold and ablation study results indicate additive benefits of the attention policy, prompt encoding, and RL fine-tuning.

Qualitative analysis demonstrates Prost-RL yields crisp, well-localized heatmaps concordant with radiological and histopathological ground truth, particularly in subtle or low-involvement lesions. The attention policy confers interpretability by highlighting biopsy-centric regions with high model confidence, while the noise-robust losses avoid over-activation in benign or ambiguous regions (Figure 3).

Figure 3: Representative qualitative examples comparing Prost-RL and ProstNFound+, demonstrating greater localization precision, reduced spurious activation, and improved result transparency via attention maps.

Theoretical and Practical Implications

Prost-RL exemplifies how RL-driven spatial policy learning can address granular localization under extremely weak, noisy, and imbalanced clinical supervision. Decoupling the attention policy enables explicit optimization of diagnostic utility—not only global accuracy but ranking separation relevant to real clinical workflows. The SCE and entropy-based loss further supports robustness to supervision noise and promotes interpretability.

The modularity of the approach ensures that Prost-RL is easily extensible to other weakly labeled medical imaging modalities and can integrate with larger or multimodal foundation models as they emerge. The explicit spatial policy mechanism is likely to have impact where transparency of risk localization is demanded for human-in-the-loop settings. The success of DRPO strategies for class/difficulty imbalance and the use of ranking-based RLOs may guide post-hoc fine-tuning not only in imaging but also in other weakly supervised or noisy-label AI tasks.

Future Directions

• Extension to 3D and longitudinal μUS imaging to capture inter-slice and temporal consistency.
• Incorporation of more sophisticated clinical priors (e.g., population risk maps, radiomics).
• Cross-modal alignment and transfer learning using MRI or pathology data as auxiliary supervision where available.
• Broader application of APO/DRPO-inspired RL with decoupled, hierarchical normalization for difficult, real-world medical data.

Conclusion

Prost-RL advances prostate cancer detection in micro-ultrasound by introducing a reinforcement learning-derived spatial attention policy, supported by noise-robust supervision and adaptive group-based RL fine-tuning. The combination yields statistically robust improvements in both detection and ranking of clinically significant lesions, while providing interpretable spatial evidence. The methodological innovations demonstrated have broad relevance for weak-label, high-noise medical AI, with likely influence on other imperfect-supervision scenarios and policy-driven vision tasks.