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

Tan Su; Haoxiang Yang; Ruxin Wang; Binghui Xie

cs.ROSep 6, 2026

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

Authors: Tan Su, Haoxiang Yang, Ruxin Wang, Binghui Xie

Organizations: Southern University of Science and Technology · Sun Yat-sen University · Knowin

Abstract

Behavior-cloned visuomotor policies can remain accurate near their training distribution yet fail when object position and camera viewpoint change together. A successful reference trajectory contains the geometry needed to transfer the same interaction, but the policy must align that geometry with the current scene and remain sensitive to it during denoising. To address these challenges, we present MemCorr-DP, a diffusion policy that lifts frozen RoMa v2 matches into explicit 3D relations between the current scene and the reference trajectory. A counterfactual paired objective assigns opposite behaviors the same physical state and noisy action while retaining reference-specific denoising targets. Mixed-condition fine-tuning then adapts the policy from ground-truth geometry to measured correspondence errors. Our strongest evaluation places the Door in the outermost position bands beyond the training support and changes the query camera by $\pm15^\circ$ . Under this combined shift, MemCorr-DP achieves 96.67% closed-loop success, compared with 88.00% for a visual Transformer with the same action architecture. Objective ablations and reference interventions show that behavior responds to the selected reference, while matched controls favor the complete relation set over future motion or centroid geometry alone. These results support explicit 3D reference relations as a robust conditioning interface when spatial and viewpoint changes are compounded in the evaluated task.

Figures & tables

Figure 1: End-to-end MemCorr-DP pipeline. Top: frozen RoMa v2 matches a fixed successful reference to the live query. Up to sixteen sampled reference points are matched and lifted into a shared 3D frame. Four future offsets per point produce a fixed $64\times 18$ relation-token array; invalid slots are masked. Green, orange, purple, and blue vectors denote the live gripper, reference anchor, future motion, and scene displacement. The Transformer uses these tokens throughout 20 DDIM steps. Bottom: paired training shares state and noisy input $z_{t}$ across opening and closing. Distinct references and clean action targets $x_{0}^{b}$ define the behavior-specific denoising targets. The policy predicts 16 actions, executes eight, and then rematches and replans.

Figure 2: Standard and counterfactual paired denoising. In standard conditional diffusion, the clean action for each behavior produces a different noisy input. Counterfactual paired denoising uses the same robot state $s$ and noisy action $z_{t}$ for both behaviors while changing the reference and denoising target, as shown explicitly in the two right-hand input tuples.

Figure 3: Factorized position and viewpoint evaluation. (a) Training uses $B\times B$ , with $B=[0.2,0.8]$ . Boundary and extreme tests place one coordinate in $N=[0.05,0.15]\cup[0.85,0.95]$ or $E=[0,0.05]\cup[0.95,1]$ , respectively, while keeping the other in $B$ . (b) Blue frames show representative training scenes; red frames show combined-shift scenes that pair an extreme position with a $\pm 15^{\circ}$ query-camera offset. (c) Camera tests fix the physical scene and reference camera and rotate only the query camera.

Condition	Position $(u_{x},u_{y})$	Query camera	Scenes
Training / validation	$(B,B)$	original	80 / 20
Boundary extrapolation	$(N,B)$ or $(B,N)$	original	50
Extreme position	$(E,B)$ or $(B,E)$	original	25
Camera shift	$(B,B)$	$\pm 15^{\circ}$	25
Combined shift	$(E,B)$ or $(B,E)$	$\pm 15^{\circ}$	25

Table 1: Training and test distributions.

B=[0.2,0.8]

is the training band,

N=[0.05,0.15]\cup[0.85,0.95]

is the boundary band, and

E=[0,0.05]\cup[0.95,1]

is the extreme band. The final column counts distinct physical scene setups; every test scene is evaluated for both tasks with three training seeds.

Figure 4: Closed-loop system comparison and matched relation controls. (a) Complete-system success under boundary extrapolation ( $n=300$ per method) and the combined shift ( $n=150$ ); bars pool three training seeds, circles show individual seeds, and labels give successes/total. (b) Paired difference between MemCorr-DP and Visual-DP-Transformer with 95% physical-configuration bootstrap intervals; positive values favor MemCorr-DP. (c) Combined-shift success by task ( $n=75$ per task and method). (d) Full relations versus separately trained motion-only and object-centroid controls with the matched Transformer architecture; marks follow panel (a).

Figure 5: Counterfactual paired training restores reference selectivity. (a,b) Stage-one closed-loop success with ground-truth relations from correct and opposite references after standard and paired training. Success is scored against the original task; lower opposite-reference success therefore indicates suppression of original-task behavior, not completion of the supplied opposite task. Each thin line connects the two reference conditions for one training seed (42–44; $n=100$ per seed and condition). Thick lines pool all 300 boundary-extrapolation rollouts per condition. Circles and diamonds denote correct and opposite references; badges report the pooled correct-minus-opposite difference.

Condition	Extreme position	Camera shift	Combined shift
Full RoMa v2	145/150 (96.67%)	148/150 (98.67%)	145/150 (96.67%)
Ground-truth correspondence	130/150 (86.67%)	150/150 (100.00%)	130/150 (86.67%)
Motion only	113/150 (75.33%)	119/150 (79.33%)	114/150 (76.00%)
Object centroid	–	–	117/150 (78.00%)

Table 4: Success under stronger shifts. Each entry reports count/150 (percentage). Full RoMa v2 and ground-truth correspondence use the same final checkpoints; motion and centroid conditions are trained separately. Bold marks the highest success in each column; an em dash denotes an unmeasured condition.

Figure 6: Sensitivity to reference phase under boundary extrapolation. (a) Fixed phase offsets; negative offsets lag the reference, whereas positive offsets lead it, and $+16$ frames is highlighted. (b) Phase-speed multipliers. Points pool 300 rollouts per condition; error bars are 95% physical-configuration cluster-bootstrap intervals. Axes begin at 80%.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Model	Component	Exact structure
MemCorr-DP	Relation/context encoder	Linear $18\!\to\!128$ ; learned embeddings for 16 points and 4 offsets; 2 pre-LN Transformer encoder layers, width 128, 4 heads, MLP $128\!\to\!512\!\to\!128$ , GELU, dropout 0.
	State/time/action inputs	State MLP $9\!\to\!128\!\to\!128$ with GELU; action Linear $4\!\to\!128$ plus 16 learned positional embeddings for actions; sinusoidal time embedding 128 followed by $128\!\to\!512\!\to\!128$ GELU MLP.
	Action denoiser/output	4 pre-LN blocks, each self-attention (4 heads), cross-attention to one state token plus 64 context tokens, and GELU MLP $128\!\to\!512\!\to\!128$ ; dropout 0; final LayerNorm and Linear $128\!\to\!4$ . Total 1,612,804.
Visual-DP-Transformer	Image/task tokenizer	RGB, depth, and mask form $5\times 64\times 64$ . One Conv2d $5\!\to\!18$ , kernel/stride 8, yields $8\times 8=64$ tokens. A Linear $2\!\to\!18$ one-hot task embedding is added to every token. Tokenizer total 5,832; no pretrained or frozen visual encoder.
	Shared action core	The complete MemCorr-DP core above with 1,612,804 parameters, including its $18\!\to\!128$ projection, context encoder, state encoder, time encoder, and action blocks. Total 1,618,636, or $+0.36\%$ ; the visual encoder is included.
Visual-DP-U-Net	Observation encoders	Five stride-2 Conv2d–GroupNorm(4)–Mish blocks: $5\!\to\!24$ with kernel 5, then $24\!\to\!32\!\to\!48\!\to\!64\!\to\!96$ with kernel 3; global average pooling and $96\!\to\!96$ Linear–LayerNorm–Mish. State: $9\!\to\!32\!\to\!24$ ; task: $2\!\to\!8$ . Concatenated global condition is 128D.

Appendix

Table 5: Trainable policy architectures. “Pre-LN” denotes layer normalization before the named block. Parameter totals include all listed modules and exclude frozen RoMa v2.

Goal-Conditioned Robotic Manipulation

See all 18 papers

Referring-Aware Visuomotor Policy Learning for Closed-Loop Manipulation
Apr 7, 2026Jiahua Ma, Yiran Qin, Xin Wen +5Visuomotor Policy Learning Robot Policy Adaptation
Sun Yat-sen University · Oxford University.
ReMAP-DP: Reprojected Multi-view Aligned PointMaps for Diffusion Policy
Mar 16, 2026Xinzhang Yang, Renjun Wu, Jinyan Liu +1Vision-Based Robot Control Diffusion Policy
School of Computer Science, Beijing Institute of Technology, China.
Latent Diffusion Policy: Shaping Latent Spaces for Diffusion-Based Robotic Manipulation
Jun 7, 2026Zhexuan Zhou, Yichen Lai, Jinhao Zhang +3Visuomotor Policy Learning Diffusion Policy

Apr 7, 2026J/K move · Enter open · S save

Referring-Aware Visuomotor Policy Learning for Closed-Loop Manipulation

Jiahua Ma, Yiran Qin, Xin Wen, Yixiong Li+4

Sun Yat-sen University · Oxford University.

Visuomotor Policy Learning Robot Policy Adaptation Robotic Manipulation Robot Policy Generalization Diffusion Policy

Open paper AskPDF arXiv

This paper addresses a fundamental problem of visuomotor policy learning for robotic manipulation: how to enhance robustness in out-of-distribution execution errors or dynamically re-routing trajectories, where the model relies solely on the original expert demonstrations for training. We introduce the Referring-Aware Visuomotor Policy (ReV), a closed-loop framework that can adapt to unforeseen circumstances by instantly incorporating sparse referring points provided by a human or a high-level reasoning planner. Specifically, ReV leverages the coupled diffusion heads to preserve standard task execution patterns while seamlessly integrating sparse referring via a trajectory-steering strategy. Upon receiving a specific referring point, the global diffusion head firstly generates a sequence of globally consistent yet temporally sparse action anchors, while identifies the precise temporal position for the referring point within this sequence. Subsequently, the local diffusion head adaptively interpolates adjacent anchors based on the current temporal position for specific tasks. This closed-loop process repeats at every execution step, enabling real-time trajectory replanning in response to dynamic changes in the scene. In practice, rather than relying on elaborate annotations, ReV is trained only by applying targeted perturbations to expert demonstrations. Without any additional data or fine-tuning scheme, ReV achieve higher success rates across challenging simulated and real-world tasks.

Abstract

Goal-Conditioned Robotic Manipulation

Explore similar work

Referring-Aware Visuomotor Policy Learning for Closed-Loop Manipulation

ReMAP-DP: Reprojected Multi-view Aligned PointMaps for Diffusion Policy

Latent Diffusion Policy: Shaping Latent Spaces for Diffusion-Based Robotic Manipulation