PRISM: Pointcloud Reintegrated Inference via Segmentation and Cross-attention for Manipulation

TL;DR

PRISM is an end-to-end framework that leverages segmentation, cross-attention, and diffusion modules to improve robot manipulation from raw point clouds, outperforming baseline policies in complex environments.

Contribution

It introduces a novel integrated approach that directly learns from raw point clouds and robot states without pretrained models, enhancing manipulation accuracy and robustness.

Findings

Outperforms baseline policies in accuracy and efficiency

Demonstrates robustness in complex, object-dense scenarios

Effective in simulated environments with limited demonstrations

Abstract

Robust imitation learning for robot manipulation requires comprehensive 3D perception, yet many existing methods struggle in cluttered environments. Fixed camera view approaches are vulnerable to perspective changes, and 3D point cloud techniques often limit themselves to keyframes predictions, reducing their efficacy in dynamic, contact-intensive tasks. To address these challenges, we propose PRISM, designed as an end-to-end framework that directly learns from raw point cloud observations and robot states, eliminating the need for pretrained models or external datasets. PRISM comprises three main components: a segmentation embedding unit that partitions the raw point cloud into distinct object clusters and encodes local geometric details; a cross-attention component that merges these visual features with processed robot joint states to highlight relevant targets; and a diffusion module that translates the fused representation into smooth robot actions. With training on 100 demonstrations per task, PRISM surpasses both 2D and 3D baseline policies in accuracy and efficiency within our simulated environments, demonstrating strong robustness in complex, object-dense scenarios. Code and some demos are available on https://github.com/czknuaa/PRISM.