Information Filtering via Variational Regularization for Robot Manipulation

TL;DR

The paper addresses inefficiencies in diffusion-based visuomotor policies caused by oversized decoders that introduce task-irrelevant noise in intermediate features. It presents Variational Regularization (VR), a timestep-conditioned Gaussian bottleneck with a KL regularizer, theoretically linked to the variational information bottleneck, to adaptively filter backbone features. Empirically, VR yields notable improvements over DP3 on RoboTwin2.0 (+6.1 percentage points) and Adroit/MetaWorld (+4.1), plus positive real-world deployments like cup-stacking with a 13.4% gain, demonstrating enhanced SNR and generalization with minimal overhead. The method offers a practical, plug-in solution to tighten information flow in diffusion-based robotic manipulation, potentially benefiting a broad class of scene-conditioned generative policies.

Abstract

Diffusion-based visuomotor policies built on 3D visual representations have achieved strong performance in learning complex robotic skills. However, most existing methods employ an oversized denoising decoder. While increasing model capacity can improve denoising, empirical evidence suggests that it also introduces redundancy and noise in intermediate feature blocks. Crucially, we find that randomly masking backbone features at inference time (without changing training) can improve performance, confirming the presence of task-irrelevant noise in intermediate features. To this end, we propose Variational Regularization (VR), a lightweight module that imposes a timestep-conditioned Gaussian over backbone features and applies a KL-divergence regularizer, forming an adaptive information bottleneck. Extensive experiments on three simulation benchmarks (RoboTwin2.0, Adroit, and MetaWorld) show that, compared to the baseline DP3, our approach improves the success rate by 6.1% on RoboTwin2.0 and by 4.1% on Adroit and MetaWorld, achieving new state-of-the-art results. Real-world experiments further demonstrate that our method performs well in practical deployments. Code will released.

Figures (7)

• Figure 1: Our proposed method, Variational Regularization (VR), adaptively filters out noise and redundant information from features.(a) Our method, built on DP3, introduces a Variational Regularization module immediately after the last downsampling features in the U-Net decoder, where noise is most likely to accumulate, enabling effective information filtering. (b) Architecture of the Variational Regularization module: it modulates the features conditioned on the diffusion timestep, then predicts the feature-wise mean and standard deviation, and uses the reparameterization trick to obtain the filtered features.
• Figure 2: Random masking analysis of U-Net decoder intermediate features. (a) We apply masks separately to the backbone features and to skip connections at different decoder depths to study the role of each component. (b) We consider two masking schemes—point-wise mask and channel-wise mask—to probe how information is organized within the representations and whether intermediate features contain redundant or noisy signals.
• Figure 3: Noisiness analysis of U-Net decoder features on Adroit-Door, Adroit-Pen, MetaWorld Disassemble, and MetaWorld StickPull, where M.W. denotes an abbreviation for MetaWorld. (a--d) Channel-wise masking of backbone features; (e--h) point-wise masking of backbone features. Backbone features are likely noisy and redundant, yet they can also contain decision-relevant signal. Our variational regularization (VR) module substantially improves the signal-to-noise ratio of the backbone features. (i--l) Skip-connection feature mask applied to different skip depth (drop_idx).
• Figure 4: Simulation experiments visualization samples across RoboTwin2.0, Adroit and MetaWorld benchmarks. Our method consistently improves the baseline model across a wide range of tasks on multiple benchmarks.
• Figure 5: Impact of Different KL Weights ($\beta$)

Theorems & Definitions (4)

• Theorem 4.1: alemi2016deep
• Corollary 4.2
• proof
• proof