Distill-Belief: Closed-Loop Inverse Source Localization and Characterization in Physical Fields

1. Problem Definition

The paper addresses Closed-Loop Inverse Source Localization and Characterization (ISLC) in physical fields (e.g., gas dispersion, radiation, temperature). In this setting, a mobile agent (robot/drone) must sequentially select measurement locations to:

1. Localize a hidden source (location).
2. Characterize latent field parameters (e.g., source strength, wind speed, diffusion coefficients).
3. Terminate the mission when the uncertainty of these parameters falls below a specific tolerance.

Core Challenges:

• Belief-Space Objective: The goal is to contract the Bayesian posterior distribution (reduce epistemic uncertainty), not just find high-signal regions.
• Sparse/Undefined Rewards: Real-world scientific tasks often lack dense rewards; success is defined by uncertainty reduction, which is a sparse or late signal.
• Computational Constraints: Real-time deployment requires constant-time decision-making ($O(1)$), whereas accurate Bayesian inference (e.g., Particle Filters) scales linearly with the number of particles ($O(N)$), making it too slow for online control.
• Reward Hacking: If a learned belief model is used both to define the reward and to guide the policy, the agent may exploit approximation errors (e.g., artificially shrinking the belief spread) to maximize rewards without genuinely reducing uncertainty.

2. Methodology: Distill-Belief

The authors propose Distill-Belief, a Teacher-Student framework that decouples scientific correctness (Bayesian accuracy) from deployment efficiency (computational speed).

A. The Teacher (Bayes-Correct Inference)

• Mechanism: A Particle Filter (PF) maintains a statistically consistent posterior over the full parameter vector $\Theta$.
• Role:
  • It performs exact Bayesian updates (reweighting, resampling, Metropolis-Hastings rejuvenation) during training.
  • It computes a dense intrinsic reward based on the Kullback-Leibler (KL) divergence between consecutive posteriors ($D_{KL}(b_t || b_{t-1})$). This reward directly measures one-step information gain.
  • It provides the "ground truth" belief statistics for distillation.
• Constraint: The Teacher is never used during deployment. It is computationally expensive ($O(N)$) and serves only as a source of supervision.

B. The Student (Amortized Inference)

• Mechanism: A compact neural network (MLP) that approximates the Teacher's posterior as a factorized Gaussian: $q_\phi(\Theta) = \mathcal{N}(\mu, \text{diag}(\sigma^2))$.
• Role:
  • Distillation: During training, the Student minimizes the weighted negative log-likelihood against the Teacher's particle distribution.
  • Control Input: At test time, the Student provides constant-time ($O(1)$) belief features (mean and variance) to the policy.
  • Stopping Certificate: The Student calculates a Spread metric ($\sqrt{\text{tr}(\Sigma)}$), which represents the Root-Mean-Square (RMS) localization error. The episode terminates when this spread falls below a threshold $\zeta$.

C. Policy Learning (Actor-Critic)

• Algorithm: Proximal Policy Optimization (PPO).
• State Input: The policy receives the current observation, agent position, and Student belief features (mean and spread).
• Reward Signal: The policy is trained using the KL-based intrinsic reward computed exclusively from the Teacher. This prevents reward hacking because the policy cannot manipulate the Teacher's posterior to inflate rewards.

3. Key Contributions

1. Decoupled Architecture: Introduces a framework that separates Bayes-correct objectives from deployment-time computation, satisfying four critical requirements simultaneously: (R1) Belief-space optimization, (R2) Handling sparse rewards, (R3) Real-time deployability, and (R4) Robustness to reward hacking.
2. Dense Information-Gain Reward: Proposes a KL-divergence-based intrinsic reward derived from the Teacher. This provides a dense, continuous learning signal that aligns RL optimization directly with posterior contraction, overcoming the sparsity of traditional success/failure signals.
3. Principled Stopping: Develops a Spread-based stopping certificate derived from the Student's belief covariance. Theoretically proven to upper-bound the expected localization error, enabling an explicit accuracy-budget trade-off.
4. Constant-Time Deployment: By discarding the Particle Filter at test time and relying solely on the distilled Student, the system achieves $O(1)$ per-step complexity, making it viable for resource-constrained mobile robots.

4. Experimental Results

The method was evaluated on seven physical field modalities (Temperature, Concentration, Magnetic, Electric, Gas, Energy, Noise) and stress-tested under multi-source and obstacle-constrained scenarios.

• Performance vs. Baselines: Distill-Belief consistently outperformed strong baselines (including InfoTaxis, EntroTaxis, DCEE, and various RL methods like GMM-PFRL and PCDQN).
  • Success Rate (SR): Achieved the highest SR across all fields (e.g., 0.95 on Temperature, 0.96 on Gas).
  • Trajectory Efficiency (TE): Required significantly fewer steps to terminate (e.g., ~18 steps vs. 40–60 for planning baselines).
  • Uncertainty Quality: Achieved the lowest Local Posterior Spread (LPS), proving it genuinely contracts the belief rather than stopping prematurely.
• Multi-Source & Obstacles:
  • In multi-source scenarios (up to 4 sources), Distill-Belief maintained robust performance with minimal degradation, whereas planning-based methods suffered severe performance drops due to multi-modal posterior confusion.
  • In obstacle-rich environments, it successfully navigated non-convex spaces, outperforming greedy planners that often got trapped in local optima.
• Ablation Studies:
  • Removing the KL reward caused a significant drop in success rate, confirming the necessity of dense belief-space shaping.
  • Using the Student for reward calculation (instead of the Teacher) led to reward hacking and poor uncertainty calibration.
  • Removing distillation (using PF at test time) maintained performance but increased latency by 6.5x, highlighting the deployment cost benefit.
• Hyperparameter Sensitivity: The method is robust to variations in particle budget ($N$) and resampling thresholds, with $N=200$ identified as the optimal cost-performance trade-off.

5. Significance

• Bridging Theory and Practice: Distill-Belief solves the fundamental tension between the need for rigorous Bayesian inference in scientific sensing and the strict latency requirements of real-world robotics.
• Mitigating Reward Hacking: By strictly separating the reward source (Teacher) from the control input (Student), the paper provides a robust solution to the "reward hacking" problem common in RL-based active sensing.
• Generalizability: The framework is agnostic to the specific physics of the field (advection-diffusion, electrostatics, etc.), making it a versatile tool for diverse AI4Science applications.
• Scalability: The constant-time inference capability enables the deployment of sophisticated uncertainty-aware agents on edge devices where running a full Particle Filter online is impossible.

In summary, Distill-Belief represents a state-of-the-art approach for autonomous scientific sensing, successfully combining the statistical rigor of Bayesian inference with the efficiency of deep reinforcement learning to achieve reliable, fast, and accurate source localization in complex physical environments.