Source-linked AI summary
MemCorr-DP: Counterfactual Correspondence Conditioning for a Diffusion Policy Guided by a Reference
Tan Su, Haoxiang Yang, Ruxin Wang, Binghui Xie
TL;DR
Visuomotor policies can fail under combined object-position and viewpoint changes, motivating a method that transfers reference geometry while keeping behavior tied to the selected reference. MemCorr-DP uses explicit 3D reference relations, counterfactual paired denoising, and mixed-condition training, achieving higher success than a matched visual Transformer in the evaluated combined shift.
Problem
Behavior-cloned visuomotor policies can fail when object position and camera viewpoint change, while a reference trajectory must be aligned with the current scene and retained during denoising.
Method
MemCorr-DP lifts frozen RoMa v2 matches into explicit 3D reference–scene relations and uses counterfactual paired denoising with mixed-condition fine-tuning for measured correspondence errors.
Results
96.67% closed-loop success versus 88.00% for a visual Transformer with the same action architecture under the combined position and viewpoint shift.
Takeaways & Limitations
Matched controls and interventions support explicit 3D reference relations as a robust conditioning interface for the evaluated compounded spatial and viewpoint shift.
Takeaways & Limitations
Evaluation is limited to one simulated Door instance, specified position and camera shifts, calibrated RGB-D lifting, fixed references, and known task identity; broader transfer remains untested.
Abstract
from arXiv · showhide
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.
1 Introduction
MemCorr-DP transfers a successful reference trajectory through explicit 3D relations while training the policy to preserve reference-specific behavior under correspondence errors. Its pipeline combines correspondence-based relation tokens, counterfactual paired denoising, and mixed-condition fine-tuning.
- Behavior cloning can fail when object position and camera viewpoint change, even though the required interaction remains unchanged.
- A successful reference provides gripper–object interaction geometry, but the policy must align it with the current scene and retain its behavioral information during denoising.
- MemCorr-DP lifts frozen RoMa v2 matches into explicit 3D relations linking live geometry, reference gripper state, future motion, and scene displacement.
- The end-to-end pipeline samples up to sixteen matched reference points, creates 64 × 18 relation tokens, and attends to them throughout 20 DDIM steps.
- Counterfactual paired denoising gives opposite behaviors the same state and noisy action while changing their references and denoising targets.
- The two-stage training strategy first learns exact geometric relations, then adapts to measured correspondence errors using mixed visual conditions and empirical residual corruption.
2 Related Work
Prior work improves diffusion-policy generalization through 3D, semantic, equivariant, perceptual, and correspondence-based representations. MemCorr-DP builds on correspondence transfer while combining live–reference 3D relations with future motion from the selected reference.
- Diffusion-policy research adds sparse 3D geometry, semantic fields, perceptual cues, image-space projections, or transformation structure to address scene-geometry shifts.
- Pixel-correspondence methods estimate cross-view matches that support geometric transfer, while manipulation systems use descriptors, relative placement, keypoint graphs, or correspondence features.
- Figure 2 contrasts standard diffusion, where behaviors receive different noisy inputs, with paired denoising, which shares state and noisy action while changing references and targets.
- The cited correspondence formulations leave open jointly using live–reference 3D relations and future gripper motion from the selected reference at every replanning step.
3 Method
MemCorr-DP builds relational conditions by matching reference and live RGB-D observations, lifting valid correspondences into 3D, and encoding geometry and future reference motion. A diffusion Transformer denoises actions from these relations, while paired objectives and mixed-condition fine-tuning make behavior reference-dependent and robust to correspondence errors.
- 3.2 Object Centered Matching and 3D Lifting: At each replanning step, frozen RoMa v2 matches object-centered reference and query crops, then lifts valid depth correspondences into a shared simulator world frame.Invalid matches and lifts outside the query object's visible mask are excluded.
- 3.3 Relational Condition Tokens: The relational condition uses four 3D blocks to encode live gripper geometry, demonstrated gripper state, future reference motion, and query-reference scene displacement.Sixteen points and four future offsets produce 64 tokens, with invalid tokens masked.
- 3.4 Diffusion Policy: The policy encodes 64 relation tokens and a 9D robot state, then uses four Transformer action blocks with self-attention over actions and cross-attention to relational context.It denoises x0 ∈ R16×4, where each action contains a 3D end-effector displacement and one scalar gripper command.
- 3.5 Counterfactual Paired Denoising: The counterfactual paired objective gives opposite behaviors identical noisy inputs and physical initial conditions while preserving reference-specific denoising targets.The paired term is auxiliary; standard diffusion updates and the same DDIM sampler remain in use.
- 3.6 Mixed-Condition Fine-Tuning: Mixed-condition fine-tuning combines ground-truth, cached RoMa v2, and corrupted relation conditions to adapt the policy to measured correspondence errors.Each batch uses 25% ground-truth, 50% cached RoMa v2, and 25% corrupted relations; paired directions share the same corruption.
4 Experimental Protocol
The evaluation uses paired open/close Door tasks, controlled reference protocols, matched visual baselines, representation controls, and position/viewpoint interventions beyond the training distribution.
- Task, Data, and Reference Protocol: The dataset contains 100 paired physical configurations with opening and closing trajectories; pairs 0–79 train the policy and pairs 80–99 provide validation.Training uses 160 trajectories and 29,600 action windows, with positions sampled from [0.2, 0.8]^2.
- Task, Data, and Reference Protocol: Evaluation uses fixed successful references at test time, while training selects reference frames by nearest Door joint value.The reference advances with elapsed execution steps, and the policy replans after executing eight of sixteen predicted actions.
- Baselines and Controls: Visual-DP-Transformer matches the action denoiser, data, optimization, diffusion schedules, horizons, seeds, and checkpoint rule while replacing relations with visual inputs.Visual-DP-U-Net provides a secondary visual comparison but also changes the denoiser architecture and objective.
- Baselines and Controls: Motion-only and centroid controls preserve the 64×18 interface, Transformer, objective, budget, horizons, and seeds while removing or aggregating point geometry.The controls isolate future motion and object-centroid information against the complete relation representation.
- Fixed Position and Viewpoint Tests: Training positions occupy B = [0.2, 0.8]^2, whereas boundary and extreme tests place one coordinate in N or E while retaining the other in B.Camera interventions rotate only the query view, and combined-shift scenes pair extreme positions with ±15° query-camera offsets.
- Metrics and Statistical Unit: Success is the official Meta-World flag within 200 environment-step calls, with task-specific handle-distance thresholds and no added contact or gripper-hold requirement.Rollouts stop at first success; timeout or termination without success counts as failure.
- Metrics and Statistical Unit: Complete-system comparisons use 300 rollouts for boundary extrapolation and 150 for the combined shift, pooled across three training seeds.Matched comparisons share physical configuration, task, and seed identifiers, while Figure 4 reports system and representation controls.
5 Results
MemCorr-DP improves performance over the visual Transformer under the combined position-and-viewpoint shift, while matched controls and interventions support complete spatial relations and reference-selective behavior.
- Main System Comparison: Under the combined shift, MemCorr-DP solves 70/75 opening rollouts versus 57/75 for the visual Transformer, while both solve 75/75 closing rollouts.The conventional Visual-DP-U-Net is lower in both evaluated conditions.
- Spatial Relations Exceed Simpler Reference Information: Full relations outperform motion-only and object-centroid conditions in both tests with the matched Transformer architecture, and every paired 95% interval is above zero.The controls compare complete representations rather than individual token fields.
- Position, Viewpoint, and Matching Robustness: Permuting query-to-reference point assignments raises combined-shift success from 130/150 with correct pairing to 148/150 under the tested permutation.This fixed-permutation intervention indicates that correct point pairing was not necessary for high success in that condition.
- Paired Denoising Training Reduces Direction Bypass: Counterfactual paired training improves correct-reference success and suppresses original-task success under opposite references across all three seeds.The first action chunk follows the supplied reference direction in every paired-model rollout.
- Behavior Responds to Reference Content: With an opposite reference, original-task success nearly vanishes, while the first action chunk follows the supplied direction in every rollout and final displacement does so in 292/300.Empty conditioning serves as a control for dependence on relation input rather than reference content.
- Position, Viewpoint, and Matching Robustness: The mixed checkpoint performs better with RoMa v2 than with ground-truth correspondence, consistent with calibration to the distribution of visual conditions used during fine-tuning.Fine-tuning weights raw or corrupted visual conditions more heavily than clean geometric conditions.
- Position, Viewpoint, and Matching Robustness: Full RoMa v2 achieves similar success across extreme-position, camera, and combined shifts, while mean correspondence error rises from 11.78 mm to 43.66 mm.Invalid-condition fallback affects 14/3,399 boundary decisions and 39/1,731 combined-shift decisions; ground-truth geometry is used only after rollout collection for diagnosis.
6 Discussion
MemCorr-DP is most advantageous when object position and camera viewpoint shift together, where explicit 3D relations connect live geometry, gripper state, and demonstrated motion. Matched controls and interventions support reference-sensitive behavior and favor the complete relation set.
- 96.67% success versus 88.00% for the visual Transformer under combined position and viewpoint shifts.The comparison uses the same action denoiser architecture.
- Figure 6 evaluates fixed phase offsets and phase-speed multipliers using 300 rollouts per condition under boundary extrapolation.Error bars are 95% physical-configuration cluster-bootstrap intervals, and axes begin at 80%.
- Explicit 3D relations tie control conditioning to interaction geometry rather than training-view pixel layout.
- Full relations outperform motion and centroid controls, while reference interventions make actions respond to the selected reference.
- High success persists as matching error grows, and point permutation results rule out correct point pairing alone as the explanation.
7 Limitations and Broader Impact
The evaluation is limited to a simulated Door setting with specific translational and camera shifts, assumptions, and sensing conditions. Physical deployment would additionally require robustness to perception, calibration, contact, actuator, and safety errors.
- The conclusions are limited to one simulated Door instance, two directions, translational position shifts, and camera azimuth changes of ±15°.
- Transfer across instances, categories, rotations, sensing domains, and physical robots remains untested.
- The evaluation assumes known task identity, one fixed reference per task, simulator-provided object masks, calibrated RGB-D lifting, and elapsed-time reference phase.
- Parameter matching excludes the frozen RoMa v2 backbone, while the systems differ in observation bandwidth and inference cost; three seeds provide limited optimization-variability evidence.
- Physical deployment would require robustness to segmentation, depth, calibration, contact, and actuator errors, plus a validated safety mechanism.
8 Conclusion
MemCorr-DP conditions action diffusion on 3D relations between a reference trajectory and the live scene, using correspondence-derived geometry and validity handling during replanning. On the evaluated Door instance, it achieves high success under combined position and viewpoint shifts and outperforms the matched visual policy.
- MemCorr-DP conditions action diffusion on 3D relations between a reference trajectory and the live scene.
- Counterfactual paired training makes action direction responsive to the selected reference, while mixed conditions accommodate measured matching error.
- RoMa v2 correspondences are lifted into a shared world frame after validity checks for depth, masks, warps, and coordinate finiteness.
- Ground-truth relation diagnostics transport reference points through Door transforms and compare estimated query points by Euclidean world-coordinate distance.
- If fewer than four valid raw points remain, relation tokens and their validity mask are zeroed while the policy continues sampling from robot state and diffusion noise.
B Token Controls and Network Architectures
The token controls compare motion and centroid substitutions against the full relation interface, while the visual policies use standardized five-channel observations and a shared robot-state/action architecture.
- All relation values use meters in the world frame, with confidence and phase in [0, 1], and fields are not separately normalized.
- Motion control zeroes fields 1–12 except gripper fields 7–9, sets confidence to one, and repeats offset tokens across nominal point slots.
- Centroid control substitutes separate valid reference and query centroids into every valid point slot, while invalid slots remain masked.
- Both visual policies resize rendered RGB, depth, and mask observations to 64 × 64, concatenate five channels, and use standardized 9D robot state with a 2D open/close code.
C Data, Optimization, and Diffusion Details
The implementation constructs bounded, normalized action windows; trains diffusion policies with paired counterfactual batches and mixed correspondence conditions; and specifies fixed architecture, optimization, and sampling details.
- Window construction and normalization: Each 200-step trajectory contributes 185 query anchors, producing 29,600 windows across 80 physical pairs and two directions.The 16-step action window and largest future offset of 15 remain in bounds.
- Window construction and normalization: Actions and 9D states are standardized from training windows, while predicted clean actions are clipped in normalized and original units during inference.State standard deviations are floored at 10^-3 and action standard deviations at 5×10^-2.
- Diffusion implementation: The diffusion process uses 100 noise levels, Gaussian initialization, 20 descending inference indices, deterministic DDIM sampling, and condition-only guidance.Guidance uses ˆϵ = ˆϵu + 1.5(ˆϵc − ˆϵu), with robot state retained in both predictions.
- Counterfactual batches: Counterfactual updates pair opening and closing states with shared normalized midpoint noise and one timestep repeated across both directions.The loss averages over 128 direction samples, horizons, and action coordinates.
- Mixed cache and corruption: Mixed-condition batches contain exactly 25% ground-truth, 50% raw RoMa cache, and 25% residual-bootstrap examples.The cache has 10,000 rows, and residuals are sampled from an empirical pool after discarding vectors above 20.034 mm.
- Optimization: Training uses AdamW with batch size 128, staged learning rates, gradient clipping, EMA, alternating standard and paired updates, and condition dropout.The learning rate is 3 × 10^-4 for 20,000 updates and 5 × 10^-5 for the final 10,000.
D Evaluation Statistics
Evaluation uses deterministic matched rollout randomness and bootstrap intervals over configuration means, with paired comparisons formed within matched rollouts.
- Rollout randomness: Each method is evaluated once for every condition, configuration, task, and training-seed combination using deterministic action noise for matched comparisons.Matched rollouts share the same initial configuration.
- Confidence intervals: The reported intervals use 100,000 bootstrap resamples of configuration means and quantify tested-configuration variation conditional on three trained models.Paired comparisons subtract outcomes within matched rollouts before averaging and resampling.