🤖 AI

MemCorr-DP: Counterfactual Correspondence Conditioning for a Diffusion Policy Guided by a Reference

MemCorr-DP is a diffusion policy that enhances visuomotor robustness under compounded spatial and viewpoint shifts by lifting 2D feature matches into explicit 3D relations with a reference trajectory and employing counterfactual conditioning to align geometry with the current scene.

Original authors: Tan Su, Haoxiang Yang, Ruxin Wang, Binghui Xie

Published 2026-09-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tan Su, Haoxiang Yang, Ruxin Wang, Binghui Xie

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Robots that learn by watching humans perform a task are becoming increasingly common, but they face a stubborn limitation: they often fail the moment the world changes slightly. If a robot learns to open a door while standing in one spot, it may become confused if the door is moved even a few inches or if the camera watching the scene shifts its angle. This happens because many robot controllers rely on memorizing specific visual patterns, like the exact shape of a door handle in a specific corner of the image, rather than understanding the underlying geometry of the interaction. When the scene changes, the familiar pattern disappears, and the robot's plan collapses. To build machines that can truly adapt, researchers are exploring ways to teach robots to focus on the spatial relationship between objects and their own movements, rather than just the pixels on a screen. This approach aims to create a form of "muscle memory" that understands how a gripper should move relative to an object, regardless of where that object happens to be sitting or how the camera is looking at it.

In a new study, researchers have developed a system called MemCorr-DP that tackles this problem by teaching a robot to align its actions with a recorded example of success, even when the scene looks very different. The core idea is simple but powerful: instead of trying to memorize the look of a successful attempt, the robot learns to match the physical geometry of its current situation to the geometry of a saved, successful trajectory. Imagine a robot trying to open a door. The researchers provide it with a video recording of a successful opening. As the robot attempts the task in a new location with a different camera angle, it uses a visual matching system to find corresponding points on the door and the robot's hand in both the live view and the recorded video. It then lifts these matching points into a shared three-dimensional space, creating a set of relationships that describe where the robot's hand is relative to the door, and where the hand in the recording was relative to the door at that same moment. This allows the robot to see that, even though the door is in a different spot, the physical relationship between the hand and the handle is the same, and it can proceed with the correct movement.

The researchers tested this system on a simulated task where a robot must open and close a door. They deliberately made the task difficult by moving the door to positions far outside the range the robot saw during training and by shifting the camera angle by fifteen degrees. In these challenging conditions, the new system succeeded in 96.67% of attempts. By comparison, a standard robot controller that relied on visual images alone, without this geometric matching, succeeded only 88% of the time. The difference might seem small, but in the world of robotics, a failure rate of 12% versus 3% represents a massive gap in reliability. The study showed that the system's success depended heavily on using the full three-dimensional relationships between points. When the researchers stripped away the detailed point-to-point matching and replaced it with simpler information, such as just the center of the door or the general direction of movement, the success rate dropped significantly. This proved that the robot needed the precise spatial map provided by the matching system to navigate the difficult shifts.

To ensure the robot was truly following the instructions in the reference video and not just guessing, the researchers introduced a clever training method. They taught the robot to distinguish between opening and closing a door by giving it the exact same starting position and the same noisy, uncertain input, but asking it to predict the action for two different outcomes based on two different reference videos. If the robot could not tell the difference between opening and closing when the reference changed, it was not learning the right lesson. This "counterfactual" training forced the system to pay close attention to the specific details of the reference trajectory. The results showed that when the robot was given the wrong reference video, it would attempt the wrong action, proving that it was genuinely responding to the guidance provided by the reference rather than relying on a pre-learned habit.

The study also examined how well the system handled errors in the visual matching process. In the real world, matching points between two different images is never perfect; there is always some small error. The researchers found that their system remained robust even when the matching was imperfect, maintaining high success rates even when the calculated positions were off by several centimeters. This resilience suggests that the system does not need pixel-perfect alignment to function; it only needs a good enough approximation of the three-dimensional relationships to guide the robot's actions. The researchers noted that while the system worked well in their simulation, it has not yet been tested on physical robots or with different types of tasks. The evaluation was limited to a single door instance and specific types of movement and camera shifts. However, the results provide strong evidence that explicitly modeling the three-dimensional relationships between a robot, an object, and a reference example is a promising path toward creating robots that can generalize their skills to new and unpredictable environments.

Technical Summary: MemCorr-DP

Problem Statement

Behavior-cloned visuomotor policies often exhibit high accuracy within their training distribution but fail when object positions and camera viewpoints change simultaneously. While a successful reference trajectory contains the necessary geometric information to transfer an interaction to a new scene, standard policies struggle to align this geometry with the current observation and remain sensitive to it during the action generation process. Specifically, diffusion policies may rely too heavily on the noisy action input to infer behavior, weakening their dependence on the reference. Furthermore, a gap exists between the perfect geometric conditions available during training and the imperfect visual matching conditions encountered during execution.

Methodology

The authors propose MemCorr-DP, a diffusion policy designed to bridge this gap by conditioning action generation on explicit 3D relations between a reference trajectory and the live scene. The method operates through the following components:

1. 3D Relation Construction

At every replanning step, the system utilizes a frozen RoMa v2 matcher to establish dense pixel correspondences between the reference and the current query RGB-D observations.

Lifting: Valid matches are lifted into a shared 3D simulator world frame.
Relation Tokens: For each matched point, the system constructs a set of 18-dimensional relation tokens. These tokens encode:
- The live gripper's position relative to the live scene geometry.
- The reference gripper's position relative to the reference geometry.
- Future motion offsets of the reference gripper.
- The displacement between the query and reference scenes.
Input: Up to 16 reference points are matched, generating a fixed $64 \times 18$ relation-token array (including future offsets) that serves as the conditioning input for the policy.

2. Action Diffusion Architecture

The policy is a Transformer-based action denoiser (1.61M parameters) that processes the 64 relation tokens and the current 9D robot state. It employs four Transformer action blocks to denoise a 16-step action sequence, where each action consists of a 3D end-effector displacement and a scalar gripper command.

3. Counterfactual Paired Denoising

To ensure the policy's behavior is strictly dependent on the selected reference rather than leaking from the noisy action input, the authors introduce a counterfactual paired objective.

Mechanism: During training, opposite behaviors (e.g., opening vs. closing a door) are paired. They share the exact same physical robot state and the same noisy action input ( $z_t$ ).
Differentiation: The only differences are the reference trajectory and the clean denoising target ( $\epsilon_b$ ).
Effect: This forces the model to resolve the distinct action targets solely based on the reference input, preventing the noisy action from acting as a shortcut for behavior identification.

4. Two-Stage Training Strategy

The training process consists of two stages to handle the discrepancy between training and execution conditions:

Stage 1: The model is trained on ground-truth geometric relations to learn the ideal mapping.
Stage 2 (Fine-tuning): The model is fine-tuned with a mixture of conditions: 25% ground-truth, 50% cached RoMa v2 matches, and 25% corrupted relations (simulating matching errors and validity patterns observed during execution). This adapts the policy to the noise and errors inherent in real-world visual matching.

Key Contributions

3D Relation Representation: A novel conditioning interface for diffusion control that connects current scene geometry with demonstrated gripper states and future reference motion in a shared 3D frame.
Counterfactual Paired Objective: A training mechanism that assigns opposite behaviors the same state and noisy input, compelling the reference trajectory to resolve the specific action target.
Robust Training Strategy: A two-stage approach that learns from exact geometric relations and subsequently adapts to measured correspondence errors through mixed-condition fine-tuning.

Experimental Results

The method was evaluated on a Meta-World Door task involving opening and closing actions. The evaluation focused on combined shifts where the door position was moved to the outermost training bands (extrapolation) and the query camera viewpoint was shifted by $\pm 15^\circ$ .

Performance: Under the combined shift condition, MemCorr-DP achieved a 96.67% closed-loop success rate. This significantly outperformed a Visual Transformer baseline with the same action architecture (88.00%) and a standard U-Net visual policy (74.67%).
Ablation Studies:
- Representation: Full 3D relations significantly outperformed controls using only motion vectors or object centroids, indicating that the complete geometric relation set is critical.
- Paired Training: Counterfactual paired training successfully suppressed behavior when an "opposite" reference was provided, whereas standard training often failed to distinguish the reference direction.
- Robustness: The system maintained high success rates even when 3D correspondence errors increased significantly (from ~11mm to ~43mm) under extreme shifts.

Significance and Claims

The paper claims that explicit 3D reference relations provide a more robust conditioning interface for visuomotor policies when spatial and viewpoint changes are compounded. By lifting visual matches into a shared 3D frame and enforcing reference-dependency through counterfactual training, MemCorr-DP preserves the geometry of the demonstrated interaction while adapting it to the current scene.

The authors note that their conclusions are specific to the evaluated Door instance and the tested distribution shifts (translational position and azimuth camera changes). They do not claim generalizability across different object categories, rotations, or physical robot deployments without further testing, acknowledging limitations regarding segmentation, depth calibration, and actuator errors in real-world settings.