Information Filtering via Variational Regularization for Robot Manipulation

Technical Summary: Information Filtering via Variational Regularization for Robot Manipulation

Problem Statement
Diffusion-based visuomotor policies utilizing 3D visual representations (e.g., point clouds) have achieved state-of-the-art performance in robot manipulation. However, these systems typically employ a significant architectural asymmetry: a compact, information-rich scene encoder paired with an oversized denoising decoder (such as U-Net or DiT with ~250M parameters). While large decoders increase model capacity, empirical evidence suggests they introduce redundancy and task-irrelevant noise into intermediate feature blocks. The authors hypothesize that this "noisy" information degrades decision-making, yet the internal representations of these decoders remain under-explored as a source of performance limitation.

Methodology
The paper proposes Variational Regularization (VR), a plug-and-play module designed to adaptively filter noise from intermediate decoder features. The methodology proceeds in three stages:

1. Noise Identification: The authors first validate the presence of noise through inference-time ablation studies.

   • For U-Net architectures, randomly masking backbone features at inference time often improves success rates, indicating that these features contain more noise than signal for specific tasks.
   • For DiT architectures, scaling down or skipping intermediate layers (specifically layers 4–8 in a 12-layer model) similarly improves performance, suggesting these layers introduce noisy feature updates.

2. Variational Regularization Module: To address this, VR is inserted after the noisy intermediate layers (after the final downsampling stage in U-Net and after intermediate layers in DiT).

   • The module treats the noisy features $Z$ as a latent variable and imposes a context-conditioned Gaussian distribution $p_\theta(\hat{Z}|Z, C)$, where $C$ is the compact scene context.
   • It predicts a mean $\mu_\theta(Z, C)$ and standard deviation $\sigma_\theta(Z, C)$ using MLPs, modulated by the context via a FiLM layer.
   • The filtered representation $\hat{Z}$ is obtained via the reparameterization trick: $\hat{Z} = \mu_\theta + \sigma_\theta \odot \epsilon$.
   • A KL-divergence regularizer is added to the loss function to encourage the distribution of $\hat{Z}$ to match a standard normal prior, acting as an adaptive information bottleneck.

3. Theoretical Justification: The authors provide a theoretical grounding using the Variational Information Bottleneck (VIB) framework. They demonstrate that minimizing the proposed loss function (denoising loss + KL regularizer) is equivalent to maximizing the Evidence Lower Bound (ELBO) of the information bottleneck objective. This theoretically guarantees that the module filters out task-irrelevant information while preserving signals critical for prediction.

Key Contributions

• Empirical Discovery: The paper identifies and quantifies that intermediate features in large diffusion decoders (U-Net and DiT) often contain task-irrelevant noise, evidenced by performance gains when these features are masked or skipped at inference.
• Novel Module: The introduction of Variational Regularization (VR), a context-conditioned variational bottleneck that adaptively filters intermediate features without requiring architectural changes to the backbone.
• Theoretical Link: A principled derivation connecting the VR objective to the Variational Information Bottleneck, providing theoretical guarantees for information filtering.
• State-of-the-Art Performance: The method achieves new state-of-the-art results across multiple benchmarks, improving upon the baseline DP3 policy.

Experimental Results
The authors evaluated VR on three simulation benchmarks (RoboTwin2.0, Adroit, MetaWorld) and a real-world deployment:

• RoboTwin2.0: VR improved the average success rate of DP3-UNet from 56.6% to 63.8% (+7.2) and DP3-DiT from 62.7% to 67.0% (+4.3).
• Adroit: Using only 10 expert demonstrations, VR improved DP3-UNet average SR5 from 68.6% to 73.7% and DP3-DiT from 70.4% to 73.8%.
• MetaWorld: VR consistently improved average SR5 for both backbones (e.g., DP3-UNet from 63.1% to 68.1%).
• Real-World: In a cup-stacking task, VR increased the success rate from 73.3% to 86.7% (+13.4%).
• Ablation: Experiments confirmed that context injection is critical for VR's performance and that VR outperforms other regularization techniques like dropout, width reduction, and vanilla Information Bottleneck.

Significance
The paper claims that by treating the oversized decoder not as a black box but as a source of potential noise, performance can be significantly enhanced through adaptive filtering. The significance lies in demonstrating that "bigger is not always better" in diffusion-based manipulation; instead, explicitly modeling and filtering intermediate representations via a variational bottleneck leads to more robust and effective policies. The method is presented as a simple, plug-and-play solution that improves generalization and success rates across diverse manipulation tasks and decoder architectures.