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

1. Problem Statement

Active Multi-Target Tracking (AMTT) involves a mobile robot navigating an environment with limited Field-of-View (FoV) to track multiple moving targets. The core challenge is the exploration-exploitation dilemma:

• Exploration: The robot must search unobserved regions to find new or lost targets.
• Exploitation: The robot must follow currently detected targets to reduce state estimation uncertainty.

Limitations of Existing Approaches:

• Unimodal Policies: Traditional methods (e.g., standard Deep Reinforcement Learning or regression-based imitation learning) often average across different behavioral strategies, leading to suboptimal "mode collapse" where the robot fails to commit to either exploration or exploitation effectively.
• Implicit Selection in Diffusion Policies: While Diffusion Policies (like MATT-Diff) can naturally represent multi-modal action distributions, they rely on stochastic noise during the denoising process to implicitly select a behavioral mode. This lacks a principled mechanism for uncertainty-aware strategy selection based on the robot's current belief state.
• Deterministic Gating: Recent Mixture-of-Experts (MoE) approaches use softmax gating to route inputs to specific experts. However, these provide only point estimates without uncertainty quantification, making them prone to overcommitting to experts when the robot encounters states outside the training distribution.

2. Methodology

The authors propose a framework that decouples strategy selection from action generation, treating expert selection as an offline contextual bandit problem.

A. System Architecture

1. Expert Planners: Three distinct rule-based planners generate demonstration data covering different modes:
   • $\pi_1$ (Exploration): Frontier-based exploration.
   • $\pi_2$ (Hybrid): Switches between tracking and exploration based on target covariance uncertainty.
   • $\pi_3$ (Time-based): Tracks for a fixed duration before exploring.
2. Diffusion Policy (Action Generator): A Denoising Diffusion Probabilistic Model (DDPM) generates action sequences. It is conditioned on:
   • Observation Features: Encoded via a CNN and Transformer (Performer) from egocentric maps and target beliefs.
   • Expert Embedding: A learned vector representing the selected expert ($k$).
3. Bayesian Expert Selector (Strategy Router): Instead of stochastic sampling or softmax gating, the system uses a Variational Bayesian Last Layer (VBLL) model to predict the performance of each expert.

B. The Bayesian Framework

• Offline Contextual Bandit Formulation: The problem is framed as selecting an "arm" (expert) given a "context" (current belief state) to maximize a reward (negative tracking error). Since the model is trained offline, it cannot collect new feedback, necessitating a pessimistic approach.
• Multi-Head VBLL:
  • A shared feature encoder feeds into $K$ separate regression heads (one per expert).
  • Each head models the reward $r_k$ as a linear function of features with a Gaussian prior on weights.
  • The model is trained to maximize the Evidence Lower Bound (ELBO), providing both a point estimate (expected reward) and predictive uncertainty (variance) in a single forward pass.
• Pessimistic Selection (LCB):
  • To avoid overcommitting to experts with unreliable predictions (high uncertainty), the system uses the Lower Confidence Bound (LCB) criterion:
    $$LCB_k = \hat{r}_k - \lambda \cdot \sigma_k$$
    Where $\hat{r}_k$ is the predicted mean reward, $\sigma_k$ is the predictive standard deviation, and $\lambda$ is a pessimism coefficient.
  • The expert with the highest LCB score is selected. This naturally penalizes experts with high variance, favoring those with both high predicted performance and high confidence.

3. Key Contributions

1. Identification of a Gap: Highlighted that current diffusion policies lack mechanisms for uncertainty-aware strategy adaptation, relying instead on stochastic noise for mode selection.
2. Novel Formulation: Framed expert selection for diffusion policies as an offline contextual bandit problem, bridging the gap between offline bandit literature and diffusion-based control.
3. Practical Algorithm: Developed a method combining Multi-Head VBLL for single-pass uncertainty estimation and LCB for pessimistic selection. This avoids the computational cost of ensemble methods (like MC Dropout) while providing robust uncertainty quantification.
4. Performance Gains: Demonstrated that explicit, uncertainty-aware selection significantly outperforms both implicit stochastic selection and deterministic gating methods.

4. Experimental Results

Experiments were conducted in simulated 2D indoor environments with 3–6 moving targets.

• Metrics: Root Mean Squared Error (RMSE), Negative Log-Likelihood (NLL), and Entropy (lower is better).
• Key Findings:
  • vs. Implicit Selection: The proposed VBLL method significantly outperformed the unconditioned MATT-Diff baseline (which uses implicit stochastic selection) across all metrics.
  • vs. Deterministic Gating: VBLL outperformed both MoE Selection (softmax routing) and MLP Regression (deterministic reward prediction). This confirms that uncertainty quantification is critical; simply predicting rewards without uncertainty leads to suboptimal decisions.
  • vs. Fixed/Random Experts: The adaptive selection strategy outperformed fixed-expert policies and random selection, proving that context-dependent switching is superior to static strategies.
  • Ablation Studies:
    • Greedy vs. LCB: Using only the mean prediction (Greedy) performed similarly to deterministic baselines, confirming that uncertainty penalization (pessimism) is the key driver of improvement.
    • UQ Methods: VBLL outperformed MC Dropout (which requires 20 forward passes) and Neuralbandit, offering better accuracy with significantly lower computational latency (single pass).
  • Efficiency: The inference time for VBLL (16.25 ms) was comparable to the baseline diffusion policy (15.21 ms) and vastly faster than rule-based RRT* planners (184.12 ms).

5. Significance

This work addresses a critical limitation in learning-based robotics: the inability to reliably select between distinct behavioral strategies in uncertain environments. By integrating Bayesian inference into the expert selection loop of a diffusion policy, the authors provide a robust, computationally efficient solution for active tracking.

The approach ensures that the robot does not blindly follow a potentially unreliable strategy when its belief state is ambiguous. Instead, it conservatively selects the expert with the most reliable predicted performance. This framework is generalizable to other multi-modal robotic tasks where safety and reliability depend on correctly identifying the appropriate behavioral mode under uncertainty.