Staged Voxel-Level DRL for Segmentation

• The paper introduces a staged voxel-level deep reinforcement learning framework that formulates 3D medical segmentation as a voxelwise policy optimization task using a vA3C module.
• It integrates a Swin-UNETR encoder with discrete action choices per voxel, leveraging local intensities and global context to iteratively refine segmentation despite annotation noise.
• Empirical evaluations with controlled ablations demonstrate state-of-the-art Dice score improvements on datasets like LA and Pancreas-CT, underscoring the method's noise-robust performance.

The voxel-level Asynchronous Advantage Actor-Critic (vA3C) module is a reinforcement learning component designed for robust 3D medical image segmentation in the presence of noisy annotations. Integrated within the Staged Voxel-Level Deep Reinforcement Learning (SVL-DRL) framework, vA3C formulates segmentation as a voxelwise policy optimization task, assigning each voxel in a 3D image as an autonomous @@@@1@@@@. Through a staged learning schedule, composite reward functions, and a voxel-adapted A3C architecture, vA3C achieves state-of-the-art robustness and accuracy on challenging medical segmentation benchmarks, especially under annotation noise (Fu et al., 7 Jan 2026).

1. Voxel Agent Formulation and State Representation

vA3C treats each voxel in a 3D input volume $I \in \mathbb{R}^{H\times W\times D}$ as a separate RL agent indexed by $i=1\ldots N$ ($N = H \cdot W \cdot D$). At training step $t$, the state $s_i^{(t)}$ is initialized as the raw intensity $I_i$, but in the full procedure, each $s_i^{(t)}$ is derived from a global feature volume $F^{(t)}$ extracted by a shared Swin-UNETR encoder, where the feature vector at the voxel’s spatial location encodes both local and global context. This agent-centric state representation enables each voxel to dynamically adjust its segmentation hypothesis based not only on its intensity but also on anatomical and contextual cues aggregated across the image.

2. Discrete Voxelwise Action Space

The vA3C action space is limited to three discrete choices per voxel: $\mathcal{A} = \{0, 1, 2\}$ where $a = 0$ corresponds to "do nothing," $a=1$ to "enhance tissue/lesion," and $a=2$ to "weaken tissue/lesion." Operationally, these actions update the voxel intensity as follows:

• $I_\text{new} = I_\text{orig}$ if $a=0$ (no change),
• $$I_\text{new} = \text{clip}\left(I_\text{orig} \cdot (1 + 0.3 \cdot \epsilon), 0, 1\right)$$ if $a=1$,
• $$I_\text{new} = \text{clip}\left(I_\text{orig} \cdot (1 - 0.3 \cdot \epsilon), 0, 1\right)$$ if $a=2$,

where $\epsilon$ is a uniform random variable in $[0, 1]$, and "clip" denotes bounding to $[0, 1]$. This quantized, stochastic "medical effect" enables nuanced voxelwise corrections, emulating annotation refinement operations.

3. Network Architecture and vA3C Adaptation

The core network comprises a shared Swin-UNETR encoder and three decoder branches:

• The segmentation head $f_{\theta_s}$ outputs a probabilistic map $P^{(t)}$.
• The value head $V_{\theta_v}$ estimates voxel- or globally-pooled returns $V(s^{(t)})$.
• The policy head $\pi_{\theta_p}$ parameterizes the per-voxel action probability distribution across $\mathcal{A}$.

All encoder weights are shared, while decoders are independently parameterized. At each iteration, actions for all voxels are sampled in parallel from the policy head (using a temperature $\tau$ annealed during training), applied to generate the next observation volume, rewards are computed, and a joint gradient update is applied. The RL update process advances through three stages:

• Warmup: Supervised Dice loss only.
• Transition: Weighted sum of Dice and value loss.
• Full RL: Composite of Dice, value (MSE), and policy gradient losses with weights $(1-\alpha-\beta), \alpha, \beta$.

This staged protocol stabilizes policy initialization and convergence under label noise.

4. Advantage Actor-Critic Loss, Rewards, and Learning Dynamics

The policy is optimized with the A3C benefit-based advantage $A(s^{(t)}, a^{(t)}) = R^{(t)} - V(s^{(t)})$, where $R^{(t)}$ is the n-step discounted return using mean-voxel rewards. The policy gradient for each voxel $i$ is: $$\nabla_{\theta_p}\mathcal{L}_\text{policy} = -\nabla_{\theta_p} \left[ \log \pi_{\theta_p}(a_i^{(t)} | s_i^{(t)}) A(s_i^{(t)}, a_i^{(t)}) \right].$$

The stepwise reward at time $t$ combines:

• $$\Delta \text{Dice} = \text{Dice}(f^{(t)}, G) - \text{Dice}(f^{(t-1)}, G)$$,
• Anatomical constraint $$\mathcal{C}(f^{(t)}) = \max(N_\text{cc}(f) - 1, 0) + \sum_{i, j} |\nabla f_{i, j}|$$,

where $N_\text{cc}(f)$ is the number of connected components and the second term encourages spatial smoothness. Thus, $r^{(t)} = \Delta \text{Dice} + \mathcal{C}$, directly incentivizing segmentation performance improvements and anatomical plausibility.

The loss formulations are:

Stage	Loss Function
Warmup	$$\mathcal{L}_{\mathrm{warm}} = 1 - \text{Dice}(P, G)$$
Transition	$$\mathcal{L}_{\mathrm{trans}} = (1-\lambda)\mathcal{L}_{\text{dice}} + \lambda (R-V(s))^2$$
Full RL	$$\mathcal{L}_{\mathrm{full}} = (1-\alpha-\beta)\mathcal{L}_{\text{dice}} + \alpha (R-V(s))^2 + \beta (-\log \pi(a|s)A)$$

Key hyperparameters include: SGD optimizer (lr=$1\!\times\!10^{-4}$); batch size 1; $1000$ epochs; discount factor $\gamma\approx 0.99$; RL weights $\alpha=0.2, \beta=0.2$; transition weight $\lambda=0.3$; temperature $\tau$ annealed to 0; n-step rollout length $K=5$; entropy regularization is not required as $\beta$ effectively controls policy loss (Fu et al., 7 Jan 2026).

5. Empirical Evaluation and Ablation

vA3C’s contribution was isolated via controlled ablations. On the LA dataset with 50% SFDA-Noise, vanilla baseline Dice is 83.37%, improved to 84.93% without the full RL stage ("w/o FRL"), and to 88.65% (+3.72%) with full vA3C. On Pancreas-CT with noise, baseline is 74.31%, "w/o FRL" is 76.58%, and vA3C achieves 78.64% (+2.06%). Ablations removing the entire RL mechanism, including warmup and transition ("w/o WAT"), cause further degradation, demonstrating the necessity of staged actor-critic optimization. This partitioned analysis attributes the majority of noise-robust performance gains specifically to the voxel-level asynchronous policy updates.

6. Robustness to Annotation Noise and Convergence

Three principled mechanisms underlie vA3C’s robustness:

• Locality: Each voxel-agent can self-correct in mislabeled regions, avoiding wholesale sample rejection.
• Asynchronous updates: Decorrelation across voxels in both spatial and temporal dimensions stabilizes training in high-dimensional image spaces.
• Composite reward: Immediate feedback from segmentation improvements (Dice delta) and anatomical priors enables focused, incremental denoising. Warmup and transition stages produce favorable initializations, protecting against policy collapse.

A plausible implication is that fine-grained, parallel corrective mechanisms at the voxel level confer both rapid convergence and an ability to resist global annotation errors prevalent in complex medical datasets. Compared with conventional supervised or sample-level denoising approaches, vA3C’s reinforcement assignment scheme is more aligned to the local structure of medical segmentation errors.

7. Summary and Broader Impact

vA3C integrates A3C-style actor-critic optimization at voxel granularity, running all $N$ voxel agents in parallel under a unified feature-encoder and jointly learned segmentation, value, and policy heads. It utilizes a reward signal that fuses segmentation improvement with anatomical constraints and employs a staged learning protocol to stabilize and accelerate convergence. Quantitative results across multiple datasets validate its state-of-the-art resilience to annotation noise, establishing voxel-level asynchronous policy optimization as a principal mechanism for robust medical image segmentation under real-world labeling uncertainties (Fu et al., 7 Jan 2026).