Reinforcement Learning for Unsupervised Domain Adaptation in Spatio-Temporal Echocardiography Segmentation

TL;DR

RL4Seg3D tackles unsupervised domain adaptation for 2D+t echocardiography segmentation by extending reinforcement-learning-based segmentation to full 3D spatio-temporal inputs. It fuses adaptive anatomical, landmark-based, and temporal rewards, complemented by a test-time optimization and uncertainty-guided refinement, to achieve superior anatomical validity, temporal coherence, and landmark accuracy on a large unlabeled target domain. The method demonstrates strong performance against baselines and foundation models on over 30,000 videos, with reliable uncertainty estimates and practical test-time improvements, all without target-domain annotations. This approach offers a scalable path to clinically robust, label-free segmentation in dynamic ultrasound data and sets the stage for broader applications in 3D medical image segmentation under domain shifts.

Abstract

Domain adaptation methods aim to bridge the gap between datasets by enabling knowledge transfer across domains, reducing the need for additional expert annotations. However, many approaches struggle with reliability in the target domain, an issue particularly critical in medical image segmentation, where accuracy and anatomical validity are essential. This challenge is further exacerbated in spatio-temporal data, where the lack of temporal consistency can significantly degrade segmentation quality, and particularly in echocardiography, where the presence of artifacts and noise can further hinder segmentation performance. To address these issues, we present RL4Seg3D, an unsupervised domain adaptation framework for 2D + time echocardiography segmentation. RL4Seg3D integrates novel reward functions and a fusion scheme to enhance key landmark precision in its segmentations while processing full-sized input videos. By leveraging reinforcement learning for image segmentation, our approach improves accuracy, anatomical validity, and temporal consistency while also providing, as a beneficial side effect, a robust uncertainty estimator, which can be used at test time to further enhance segmentation performance. We demonstrate the effectiveness of our framework on over 30,000 echocardiographic videos, showing that it outperforms standard domain adaptation techniques without the need for any labels on the target domain. Code is available at https://github.com/arnaudjudge/RL4Seg3D.

Figures (6)

• Figure 1: Overview of the full RL4Seg3D framework and its main steps. ${\@fontswitch\mathcal{D}}_s$: labeled source domain dataset, ${\@fontswitch\mathcal{D}}_\tau$: unlabeled target domain dataset, ${\@fontswitch\mathcal{D}}_{\mathit{ANAT}}$: anatomical reward dataset, $\pi_\theta$: policy, $s$: state, $a$: action, $R(s,a):$ reward function, $V_\phi^\pi(s):$ value function, $r_{\psi}^{ANAT}$: anatomical reward network, $r_{\omega}^{LM}$: landmark reward network, $A(s,a)$: advantage function.
• Figure 2: Example of a segmentation frame (a) and reward maps highlighting errors. The anatomical reward (b) from $r_{\psi}^{\mathit{ANAT}}$ shows general segmentation issues such as the misshapen apex, while the landmark reward (c) from $r_{\omega}^{LM}$ highlights errors in mitral valve alignment. The final reward map (d), obtained with min-based fusion, illustrates a localization error at the left mitral valve commissure: the segmentation places it inside the ventricle (red circles), whereas the correct location is indicated by the yellow arrow. (e) illustrates ground truth error maps for the validated landmark (LM) subset, and (f) shows the fused reward map with the temporal penalty applied. Video of full sequence and rewards available in supplementary material.
• Figure 3: Qualitative comparison of segmentations. Displayed frames correspond to the lowest Dice score between the ground truth and baseline 3D U-Net in their respective videos. Full sequence comparisons are provided in the supplementary material.
• Figure 4: Comparison of the temporal evolution of a slice from segmentations from RL4Seg (2D) and RL4Seg3D, compared to the ground truth (white outline). Protrusions show inconsistent temporal evolution of the segmentation between frames. The example shown is representative of the test set, with metrics within one standard deviation of the mean.
• Figure 5: Calibration plots for uncertainty methods, with expected calibration error (ECE) (the smaller the better). Grey histograms in the background represent the density of pixels in each confidence bin.