PRISM: Performer RS-IMLE for Single-pass Multisensory Imitation Learning

TL;DR

PRISM presents a single-pass multisensory imitation policy that combines a temporal multi-sensory encoder with a bidirectional linear-attention generator and a batch-global RS-IMLE training objective. By employing receding-horizon inference, it delivers real-time control (30–50 Hz) while preserving multimodal action coverage without iterative sampling. The approach yields significant improvements over diffusion and flow baselines across MetaWorld, CALVIN, Robomimic, and real hardware, including substantial reductions in jerk and robust performance under sensor dropouts. The work demonstrates practical impact for fast, safe imitation in real-world robotics and lays groundwork for language-conditioned manipulation and adaptive rejection schemes.

Abstract

Robotic imitation learning typically requires models that capture multimodal action distributions while operating at real-time control rates and accommodating multiple sensing modalities. Although recent generative approaches such as diffusion models, flow matching, and Implicit Maximum Likelihood Estimation (IMLE) have achieved promising results, they often satisfy only a subset of these requirements. To address this, we introduce PRISM, a single-pass policy based on a batch-global rejection-sampling variant of IMLE. PRISM couples a temporal multisensory encoder (integrating RGB, depth, tactile, audio, and proprioception) with a linear-attention generator using a Performer architecture. We demonstrate the efficacy of PRISM on a diverse real-world hardware suite, including loco-manipulation using a Unitree Go2 with a 7-DoF arm D1 and tabletop manipulation with a UR5 manipulator. Across challenging physical tasks such as pre-manipulation parking, high-precision insertion, and multi-object pick-and-place, PRISM outperforms state-of-the-art diffusion policies by 10-25% in success rate while maintaining high-frequency (30-50 Hz) closed-loop control. We further validate our approach on large-scale simulation benchmarks, including CALVIN, MetaWorld, and Robomimic. In CALVIN (10% data split), PRISM improves success rates by approximately 25% over diffusion and approximately 20% over flow matching, while simultaneously reducing trajectory jerk by 20x-50x. These results position PRISM as a fast, accurate, and multisensory imitation policy that retains multimodal action coverage without the latency of iterative sampling.

Figures (11)

• Figure 1: Overview of Performer RS--IMLE for Multisensory Imitation (PRISM). Per-timestep features from the available sensors are fused into temporal context tokens. Across all benchmarks we include one or more RGB cameras; depth camera, tactile, proprioception, text tokens, and audio are used when present in a dataset. A bidirectional FAVOR$^{+}$ (Performer) generator with learned query tokens outputs a full sequence of actions in a single pass. Training uses a batch-global RS-IMLE objective (robust Charbonnier distance, $\varepsilon$-rejection with EMA calibration, optional small coverage term) to preserve action multimodality without iterative sampling. We train separate models per benchmark (MetaWorld, CALVIN, Robomimic, real hardware) using only the modalities available in that benchmark (see Fig. \ref{['fig:dataset_summary']} for list of modalities used per dataset type); unused sensors are shown in gray. At inference we sample k latent-conditioned trajectories in one batched pass and select one for receding-horizon control for hardware experiment.
• Figure 2: Benchmark Datasets. On the right, we show the sensor modalities used per dataset. R is for RGB camera, D is for depth camera, Tact is for tactile, P is for proprioception, Text is for text tokens, and A is for audio.
• Figure 3: Modality robustness on CALVIN.Left: Single-modality dropouts during evaluation. Middle: Pairwise modality dropouts. Right: Performance change relative to the full-modality baseline. Wrist RGB and proprioception are the most critical modalities; removing both leads to near-complete failure, while PRISM degrades gracefully under single-modality removal.
• Figure 4: Real-World Hardware Setup and Performance.Top Left: Loco-manipulation platform (Unitree Go2+D1) and tasks including parking, pick-and-place, and insertion. Top Right: Loco-manipulation success rates (50 trials/method); PRISM maintains 30 Hz while achieving higher success. Bottom Left: Tabletop manipulation tasks showing nested containers and precise can placement. Bottom Right: Tabletop success rates and inference speed ablation for tasks (1) and (2).
• Figure 5: Push-T multimodality visualization.Columns sweep the end-effector start position along the top edge of the T-block from left to right. Each cell overlays multiple trajectories from a given policy. PRISM maintains diverse yet smooth trajectory families, especially in the ambiguous central region. The IMLE baseline shows fewer distinct modes, Diffusion Policy produces more noisy branches, and Flow-Matching remains effectively unimodal across all start states.

Theorems & Definitions (8)

• Lemma E.1: Implicit Likelihood Maximization and Mode Coverage
• proof
• Proposition E.1: Distance as Likelihood Lower Bound
• proof
• Lemma E.2: Consistency of Batch-Global Quantile Estimator
• proof
• Lemma E.3: Soft Coverage as Entropy Regularization
• proof