Latent Theory of Mind: A Decentralized Diffusion Architecture for Cooperative Manipulation

TL;DR

LatentToM tackles cooperative manipulation by decentralizing diffusion policies across arms with partial observations. It learns two latent representations per arm—a robot-specific ego embedding and a shared consensus embedding—and enforces global consistency via a first-order cohomology loss from sheaf theory, while incorporating Theory of Mind-inspired predictions and a directional consensus mechanism to preserve expressiveness. Training is centralized to supervise the consensus encoders, while inference can be fully decentralized with optional Sheaf Laplacian-based refinement or limited inter-arm communication. Hardware experiments show LatentToM achieves competitive performance with state-of-the-art centralized diffusion policies and outperforms naive decentralized approaches on two cooperative tasks, highlighting its potential for scalable multi-arm collaboration.

Abstract

We present Latent Theory of Mind (LatentToM), a decentralized diffusion policy architecture for collaborative robot manipulation. Our policy allows multiple manipulators with their own perception and computation to collaborate with each other towards a common task goal with or without explicit communication. Our key innovation lies in allowing each agent to maintain two latent representations: an ego embedding specific to the robot, and a consensus embedding trained to be common to both robots, despite their different sensor streams and poses. We further let each robot train a decoder to infer the other robot's ego embedding from their consensus embedding, akin to theory of mind in latent space. Training occurs centrally, with all the policies' consensus encoders supervised by a loss inspired by sheaf theory, a mathematical theory for clustering data on a topological manifold. Specifically, we introduce a first-order cohomology loss to enforce sheaf-consistent alignment of the consensus embeddings. To preserve the expressiveness of the consensus embedding, we further propose structural constraints based on theory of mind and a directional consensus mechanism. Execution can be fully distributed, requiring no explicit communication between policies. In which case, the information is exchanged implicitly through each robot's sensor stream by observing the actions of the other robots and their effects on the scene. Alternatively, execution can leverage direct communication to share the robots' consensus embeddings, where the embeddings are shared once during each inference step and are aligned using the sheaf Laplacian. In our hardware experiments, LatentToM outperforms a naive decentralized diffusion baseline, and shows comparable performance with a state-of-the-art centralized diffusion policy for bi-manual manipulation. Project website: https://stanfordmsl.github.io/LatentToM/.

Figures (8)

• Figure 1: Multi-arm robotic system. In our setup, the system consists of two robotic arms, each equipped with an end-effector camera, represented by the red and green areas indicating their respective fields of view. Additionally, a third-person camera observes the overlapping workspace between the two arms, shown in gray. The bottom part illustrates the consensus embeddings generated using sheaf theory from our collected data.
• Figure 2: Cooperative Manipulation Tasks. In (a), the yellow dots and lines represent the ideal trajectory of the T-shaped block, with the expectation that its orientation remains largely unchanged throughout the motion. In (b), the gray dots and solid lines indicate the past trajectories of the two arms, while the green dots and dashed lines depict their future planned trajectories.
• Figure 3:
• Figure 4:
• Figure 5:

Theorems & Definitions (2)

• Theorem 1
• proof