Diffusion Policy with Bayesian Expert Selection for Active Multi-Target Tracking

Abstract: Active multi-target tracking requires a mobile robot to balance exploration for undetected targets with exploitation of uncertain tracked ones. Diffusion policies have emerged as a powerful approach for capturing diverse behavioral strategies by learning action sequences from expert demonstrations. However, existing methods implicitly select among strategies through the denoising process, without uncertainty quantification over which strategy to execute. We formulate expert selection for diffusion policies as an offline contextual bandit problem and propose a Bayesian framework for pessimistic, uncertainty-aware strategy selection. A multi-head Variational Bayesian Last Layer (VBLL) model predicts the expected tracking performance of each expert strategy given the current belief state, providing both a point estimate and predictive uncertainty. Following the pessimism principle for offline decision-making, a Lower Confidence Bound (LCB) criterion then selects the expert whose worst-case predicted performance is best, avoiding overcommitment to experts with unreliable predictions. The selected expert conditions a diffusion policy to generate corresponding action sequences. Experiments on simulated indoor tracking scenarios demonstrate that our approach outperforms both the base diffusion policy and standard gating methods, including Mixture-of-Experts selection and deterministic regression baselines.

• The paper presents a novel diffusion policy that decouples mode selection from action generation for robust multi-target tracking.
• It employs a Variational Bayesian Last Layer with a Lower Confidence Bound rule to select among multiple expert strategies under uncertainty.
• Experimental results demonstrate superior tracking performance, achieving lower RMSE, NLL, and belief entropy compared to existing methods.

Diffusion Policy with Bayesian Expert Selection for Active Multi-Target Tracking

Introduction and Motivation

The paper "Diffusion Policy with Bayesian Expert Selection for Active Multi-Target Tracking" (2604.03404) addresses the challenge of active multi-target tracking, where a mobile robot must optimize its navigation strategy to detect, track, and reacquire multiple dynamic targets under field-of-view constraints and limited target observability. Typical approaches struggle with the exploration–exploitation dilemma: choosing when to search unexplored areas for undetected targets versus when to reduce track uncertainty over those currently detected. Most existing policy learning or planning solutions collapse to unimodal strategies, exhibit mode averaging, or incur excessive computational cost due to online sampling or optimization.

Diffusion policies, leveraging conditional denoising generative modeling, have recently exhibited strong performance for multi-modal sequential decision-making by learning from expert demonstrations with diverse behaviors. However, current practice samples behavioral modes stochastically during denoising, lacking a mechanism for explicit, uncertainty-aware mode selection. This deficiency is especially acute in safety-critical applications, where overcommitting to poorly understood behaviors degrades reliability.

Methodology: Bayesian Expert Selection as Offline Contextual Bandits

The core methodological contribution is the explicit decoupling of mode selection from low-level action generation. Multiple expert planners—frontier-based exploration, uncertainty-driven hybrid, and time-based hybrid—generate a suite of multi-modal demonstration trajectories. A single diffusion policy, conditioned on an expert index embedding, is trained to imitate these action sequences given the current observation history and belief state.

At deployment, the selection of which expert's strategy to adopt is framed as an offline contextual bandit problem. Crucially, the agent must choose among $K$ experts under uncertainty regarding their efficacy in the current context. The authors leverage a Variational Bayesian Last Layer (VBLL) model—one head per expert—to regress each expert's expected tracking reward (negative log-likelihood of target state estimates) and quantify predictive uncertainty.

The selection rule applies a Lower Confidence Bound (LCB) criterion:

$$\mathrm{LCB}_*^k = \hat{r}_*^k - \lambda \cdot \sigma^k_*$$

where $\hat{r}_*^k$ and $\sigma^k_*$ are the predicted mean and epistemic uncertainty of the reward for expert $k$ in current context, and $\lambda$ tunes pessimism. The expert maximizing the LCB is chosen, mitigating risk of optimistic extrapolation in out-of-distribution states.

Diffusion Policy Architecture and Training

The architecture comprises:

• An observation encoder using a hybrid CNN and Performer transformer for egocentric spatial maps, with a separate attention-based module for target belief state vectors.
• Expert embeddings concatenated to context features, enabling the diffusion policy to condition action generation on expert identity.
• A UNet backbone for the denoising action diffusion process, trained on multi-expert data with action corruption under a controlled variance schedule.

All expert and context encoders are held fixed post-training, ensuring that improvements in behavior arise solely from selection quality, not architectural variance.

Inference proceeds as follows: For each receding-horizon step, context features are extracted, the VBLL module computes confidence-aware scores per expert, the best is selected, and the diffusion model generates the next action sequence conditioned on this choice.

Experimental Results

Tracking performance is evaluated in a simulated 2-D indoor environment with multiple mobile targets, using metrics of root-mean-squared error (RMSE), negative log-likelihood (NLL) of true target states, and entropy of belief distributions. Comparisons span:

• Rule-based expert planners (frontier explore, hybrid, fixed strategies)
• Unconditioned and fixed-expert diffusion baselines (including random selection)
• Standard mixture-of-experts (MoE) models with deterministic or softmax gating
• Several uncertainty-aware selection baselines, including MC Dropout and Neuralbandit LCB routing

Strong numerical results include:

• The proposed VBLL-LCB method achieves the lowest RMSE ($480.8 \pm 107.3$), NLL ($14.06 \pm 0.89$), and entropy ($13.01$) among all learning-based policies, outperforming both unconditioned stochastic sampling and deterministic MoE gating.
• Explicit, uncertainty-aware selection yields substantial improvement over implicit, stochastic mode selection by the diffusion model.
• Deterministic selectors (MLP Regression, MoE) offer negligible gains, supporting the necessity of uncertainty quantification for robust strategy routing.

  Figure 1: Trajectory comparison on a representative episode with three targets (marked by colored diamonds).

Figure 1 illustrates that the VBLL-guided policy efficiently navigates to cover all targets, while the stochastic baseline explores more erratically, and rule-based planning produces compact, goal-directed paths.

Ablation and Analysis

Ablation studies quantify the sources of improvement:

• The LCB selection rule outperforms posterior mean (greedy), UCB, and Thompson sampling approaches, corroborating the pessimism principle in the offline bandit setting.
• VBLL offers better calibration and computational efficiency (single pass per expert) than alternative UQ methods such as MC Dropout, especially critical for real-time deployment.
• The impact of the pessimism parameter $\lambda$ is minor within reasonable ranges but necessary—mean-based selection degrades performance.

Computation-wise, VBLL incurs negligible overhead compared to diffusion-only baselines and is far faster than traditional online planning (e.g., RRT*-based exploration).

Implications and Future Directions

This work demonstrates the practical and theoretical benefits of merging uncertainty-aware contextual bandit formulations with expert-conditioned diffusion policy generation for active perception and robot planning. The approach directly addresses mode collapse and overfitting to dominant behaviors, which pose acute risks in multi-modal, partially observable settings.

Practically, the Bayesian selector enables robust, conservative adaptation to belief state uncertainty, offering a methodology readily extensible to larger planner ensembles or other multimodal imitation scenarios. The computational efficiency and generality of the VBLL-LCB combination make it viable for real-time systems with strict latency budgets.

Open theoretical and algorithmic questions include:

• Generalizing to online contextual bandit settings, updating posterior reward models as fresh experience is gathered, thus adapting to nonstationary or evolving environments.
• Extending action generation fidelity using more expressive or hierarchical generative models to close the performance gap to rule-based planners in low-level policy execution.
• Integrating multi-objective selection or dynamic expert pools to further enhance adaptability in more complex or safety-critical tracking domains.

Conclusion

The paper establishes a robust framework for uncertainty-aware expert selection in diffusion policy learning for active multi-target tracking. By formulating selection as an offline contextual bandit regression problem with Bayesian last-layer uncertainty, and decoupling it from action generation, the method offers superior empirical performance, demonstrated reliability, and clear scalability pathways for future intelligent robotic systems engaging in complex, multimodal sequential decision-making tasks.