Trace-Focused Diffusion Policy for Multi-Modal Action Disambiguation in Long-Horizon Robotic Manipulation

Technical Summary: Trace-Focused Diffusion Policy (TF-DP)

1. Problem: Multi-Modal Action Ambiguity (MA2)

The paper identifies a fundamental failure mode in generative imitation learning for long-horizon robotic manipulation: Multi-Modal Action Ambiguity (MA2).

In long-horizon tasks, visually similar observations often recur at different stages of execution (e.g., a robot returning to a "home" position after different sub-tasks). Because standard Diffusion Policies are typically conditioned only on instantaneous observations, they encounter a "one-to-many" mapping problem: the same visual input corresponds to multiple valid but distinct actions depending on the task stage. This leads to:

• Unstable action sampling: The policy assigns comparable likelihoods to incompatible actions.
• Temporal inconsistency: The robot fails to follow the correct sequence, causing task failure.
• Sensitivity to disturbances: Visual clutter or background changes exacerbate this ambiguity, leading to compounding errors.

2. Methodology: Trace-Focused Diffusion Policy (TF-DP)

To resolve MA2, the authors propose TF-DP, a framework that incorporates the robot's execution history as an explicit, spatially-grounded "trace" rather than just a sequence of past tokens.

Key components include:

• Motion Trace Acquisition: Instead of stacking high-dimensional past observations (which increases computational cost), the policy projects the 3D historical end-effector positions ($H_t$) into the 2D global camera image space ($I_{trace}$) using standard perspective projection.
• Trace-Focused Field (TFF) Rendering: To make the discrete trajectory points more useful for a neural network, the authors convert the 2D trace into a continuous "energy field." They apply a dense Gaussian response map to each trace point and sum them to create a field ($F_{t}^{2d}$). This field is used to modulate the global observation ($\tilde{I}_{global} = I_{global} \odot F_{t}^{2d}$), effectively "highlighting" the task-relevant regions where the robot has recently moved.
• Execution-Aware Conditioning: The diffusion denoising process is conditioned on an augmented observation $\tilde{o}_t$ that includes the enhanced global view, side views, wrist views, the projected trace image, and the current end-effector position. This provides the policy with the necessary "stage-aware" context to distinguish between visually identical scenes.

3. Key Contributions

• Identification of MA2: Formally defining the ambiguity problem in long-horizon generative policies.
• TF-DP Framework: A lightweight method to integrate historical motion into the visual observation space via TFF rendering.
• Efficiency-First Design: Unlike methods that stack past observations or actions (which increase latency and memory), TF-DP uses a single augmented image, keeping the neural architecture simple and the computational overhead minimal.
• Generality: Demonstrating that the Trace-Focused module is compatible with various generative formulations, including DDIM and Flow Matching.

4. Experimental Results

The authors evaluated TF-DP on three real-world tasks (Place Cube, Press Keyboard, and Pick & Place from Drawers) using a Franka Emika Research 3 arm.

• Resolving MA2: TF-DP significantly outperformed the vanilla Diffusion Policy (DP) and the history-action conditioned baseline (DP-HistAct). In Task 3 (the most complex), TF-DP achieved an 83.33% success rate, whereas DP and DP-HistAct failed to maintain temporal consistency.
• Robustness to Disturbances: When background clutter and distractors were introduced, TF-DP maintained high performance (88.89% average success), while the vanilla DP dropped to near 0%. This proves the TFF effectively suppresses irrelevant visual noise.
• Computational Efficiency:
  • Inference: Only a 6.4% increase in runtime (approx. 0.02s extra).
  • Memory: While training memory increased by ~18% due to image augmentation, it was far more efficient than DP-HistAct, which saw a 107% increase in memory usage.

5. Significance

This work is significant because it moves away from the heavy reliance on hierarchical planners (like LLMs) or complex temporal stacking to solve long-horizon problems. Instead, it proposes a principled, reactive, and single-policy approach. By grounding temporal context in the robot's own physical motion trace, TF-DP provides a scalable way to achieve robust, long-horizon manipulation in real-world, cluttered environments with minimal additional hardware or computational complexity.