Act to See, See to Act: Diffusion-Driven Perception-Action Interplay for Adaptive Policies
Extracted source text · arXiv 2509.25822
\documentclass{article}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2025
% ready for submission
% \usepackage{neurips_2025}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
% \usepackage[preprint]{neurips_2025}
% to compile a camera-ready version, add the [final] option, e.g.:
\usepackage[final]{neurips_2025}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2025}
\usepackage{natbib}
\bibliographystyle{abbrvnat}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
% \usepackage{xcolor} % colors
\usepackage{algpseudocode}
% \usepackage[vlined, ruled, boxed, linesnumbered]{algorithm2e}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{geometry}
\usepackage{graphicx}
\usepackage{upgreek}
% \usepackage{cite}
\usepackage{longtable}
\usepackage{cellspace, makecell, multirow}
\usepackage{subcaption}
\usepackage{amsmath,amsfonts,amssymb,bm}
\usepackage{amsthm}
\usepackage{multirow}
\newtheorem{theorem}{Theorem}
\newtheorem{corollary}{Corollary}
\newtheorem{lemma}{Lemma}
\newtheorem{proposition}{Proposition}
\usepackage{soul}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator*{\argmax}{arg\,max}
\usepackage{enumitem}
\newtheorem*{theorem*}{Theorem}
\newtheorem*{lemma*}{Lemma}
\newtheorem*{proposition*}{Proposition}
\usepackage[normalem]{ulem}
\newcommand{\bl}[1]{\textcolor{blue}{#1}}
% \usepackage[table,xcdraw]{xcolor} % Load xcolor with the table option
% --- fix xcolor option clash + table colors ---
\PassOptionsToPackage{table}{xcolor}% must come before xcolor is loaded anywhere
% Load color-for-tables support
\usepackage{colortbl}
\usepackage{xcolor}
\newcommand{\veryshortarrow}[1][3pt]{\mathrel{%
\vcenter{\hbox{\rule[-.5\fontdimen8\textfont3]{#1}{\fontdimen8\textfont3}}}%
\mkern-4mu\hbox{\usefont{U}{lasy}{m}{n}\symbol{41}}}}
% \title{Future-Aware Diffusion Policies: Enhancing Robotic Control with Variational Spatiotemporal Embeddings}
\title{
Act to See, See to Act: Diffusion-Driven Perception-Action Interplay for Adaptive Policies
}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
% \author{%
% Jing Wang\thanks{Use footnote for providing further information
% about author (webpage, alternative address)---\emph{not} for acknowledging
% funding agencies.} \\
% Department of Computer Science\\
% Cranberry-Lemon University\\
% Pittsburgh, PA 15213 \\
% \texttt{hippo@cs.cranberry-lemon.edu} \\
% % examples of more authors
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \AND
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% }
\author{
Jing Wang$^{1}$ \quad
Weiting Peng$^{2}$ \quad
Jing Tang$^{2}$ \quad
Zeyu Gong$^{2}$ \quad
Xihua Wang$^{1}$ \quad\\
\textbf{Bo Tao}$^{2}$ \quad
\textbf{Li Cheng}$^{1}$ \\
$^{1}$University of Alberta \qquad
$^{2}$Huazhong University of Science and Technology\\
\{jing39, xihua, lcheng5\}@ualberta.ca \\
\{u202210565, j\_tang, gongzeyu, taobo\}@hust.edu.cn
}
\begin{document}
\maketitle
\begin{abstract}
Existing imitation learning methods decouple perception and action, which overlooks the causal reciprocity between sensory representations and action execution that humans naturally leverage for adaptive behaviors.
To bridge this gap, we introduce Action-Guided Diffusion Policy (DP-AG), a unified representation learning that explicitly models a dynamic interplay between perception and action through probabilistic latent dynamics.
DP-AG encodes latent observations into a Gaussian posterior via variational inference and evolves them using an action-guided SDE, where the Vector–Jacobian Product (VJP) of the diffusion policy's noise predictions serves as a structured stochastic force driving latent updates.
To promote bidirectional learning between perception and action, we introduce a cycle-consistent contrastive loss that organizes the gradient flow of the noise predictor into a coherent perception–action loop, enforcing mutually consistent transitions in both latent updates and action refinements.
Theoretically, we derive a variational lower bound for the action-guided SDE, and prove that the contrastive objective enhances continuity in both latent and action trajectories.
Empirically, DP-AG significantly outperforms state-of-the-art methods across simulation benchmarks and real-world UR5 manipulation tasks.
As a result, our DP-AG offers a promising step toward bridging biological adaptability and artificial policy learning.
Code is available on our project website: \href{https://jingwang18.github.io/dp-ag.github.io/}{https://jingwang18.github.io/dp-ag.github.io/}.
\end{abstract}
\section{Introduction}
\begin{figure}[t]
\centering
\includegraphics[width=0.9\textwidth]{imgs/overview_gap.pdf}
\vspace{-1mm}
\caption{{\bf Use of Observation Features.} {\bf (a)} Conventional methods map observation features directly to actions. {\bf (b)} DP models action distributions through incremental denoising from white noise, conditioned on observation features. {\bf (c)} DP-AG refines observation features via noise predictions, establishing a mutually reinforcing cycle between perception and action.}
\label{fig:gap}
\vspace{-3mm}
\end{figure}
Imitation learning (IL) enables agents to replicate expert behavior from demonstrations.
Direct mapping methods, such as~\citep{codevilla2018end, mandlekar22a, florence2022implicit}, learn a direct mapping between observations and actions.
In contrast, generative models such as Diffusion Policy (DP)~\citep{chi2023diffusion} and flow-matching methods~\citep{hu2024adaflow, zhang2025flowpolicy} model action distributions to improve continuity across time steps.
Vision-Language-Action (VLA)~\citep{kim2024openvla, black2024pi_0} improves perception by leveraging Vision-Language Models (VLMs) to interpret environmental cues and high-level instructions.
Despite progress, existing methods treat observation features as static during each action sequence generation (extracting them from a single time-point observation and holding them fixed while generating a short sequence of actions), thus overlooking the opportunity for the intermediate action feedback to refine perception.
Robust decision-making relies on a continuous interplay between perception and action~\citep{o2001sensorimotor}.
Humans naturally embody this principle by dynamically refining their environmental understanding through feedback from their own actions~\citep{brooks1991intelligence}.
Motivated by this, we propose {\em Action-Guided Diffusion Policy} (DP-AG), a representation learning framework for IL that explicitly models the perception–action interplay through probabilistic latent dynamics.
We build upon DP because it models the continuity within each short-horizon action sequence, where intermediate action feedback can be derived from its noise predictions at each diffusion step.
DP-AG first grounds observation features in a Gaussian posterior via variational inference, capturing uncertainty from observation inputs.
To enable dynamic perception, we introduce an action-guided stochastic differential equation (SDE) in which latent features evolve across diffusion steps, driven by action-conditioned noise predictions.
Here, the Vector–Jacobian Product (VJP) from the diffusion model acts as a structured stochastic force that shape feature evolution.
% To enable dynamic perception, we introduce an action-guided stochastic differential equation (SDE), where latent features evolve over diffusion steps.
% This evolution is driven by action-conditioned noise predictions, where the Vector–Jacobian Product (VJP) from the diffusion model acts as a structured stochastic force.
Although raw observations contain all available information, their usefulness for policy learning depends on how they are internally represented.
Unlike static encoders that keep features fixed over an entire action sequence generation, DP-AG continuously refines them using action feedback.
This process mirrors active perception in biological agents, where even if the external scene remains unchanged, sensory inputs are reinterpreted in the context of ongoing actions.
This captures the essence of {\em Act to See, See to Act}: the same observation is continually reinterpreted as actions unfold, where {\em see} refers to the evolving latent representation refined through action feedback rather than new external sensing. As shown in Figure~\ref{fig:gap}, DP-AG conditions latent evolution on action-guided noise, using the continuity of action diffusion to keep perceptual dynamics aligned with the denoising of actions.
A cycle-consistent contrastive loss further aligns latent evolution with action diffusion, preventing drift and reinforcing the perception–action loop.
Combined with variational inference, these components yield a principled and adaptive representation that produces smoother, more context-aware trajectories, as confirmed by both synthetic and real-robot experiments.
% Although raw observations contain all available information, the way they are internally represented is important for policy learning.
% Unlike static encoders, which lock features for the duration of an action sequence, DP-AG continually refines them in response to action feedback.
% This process parallels active perception in biological agents: even when the external scene is unchanged, agents reinterpret sensory inputs in light of their ongoing actions.
% Here, {\em Act to See, See to Act} denotes this evolving internal representation: seeing the same observation input differently as actions unfold.
% By aligning perception with behavior in this dynamic manner, DP-AG produces smoother and more context-aware trajectories, as confirmed by both synthetic and real-robot experiments.
% Figure~\ref{fig:gap} illustrates this adaptive perception–action interplay within each diffusion-based action sequence generation.
% Importantly, by conditioning latent evolution on action-guided noise, DP-AG leverages the continuity of the DP’s action refinement to guide perceptual dynamics.
% This ensures that observation features evolve coherently alongside the action denoising process.
% To reinforce this bidirectional learning, we introduce a cycle-consistent contrastive loss that aligns latent evolution with action denoising across diffusion steps.
% By tightly integrating variational inference, VJP-guided latent evolution, and cycle-consistent learning, DP-AG offers a principled and adaptive observation representation for policy learning in dynamic environments.
% Unlike existing methods that fix observation features during each action sequence generation, DP-AG continuously refines them alongside diffusion-based action denoising, aligning perception with evolving behavior.
% Figure~\ref{fig:gap} illustrates the adaptive perception–action interplay that occurs within each diffusion-based action sequence generation.
Theoretically, DP-AG introduces an action-guided latent SDE, derives a principled variational lower bound, and rigorously proves that the proposed cycle-consistent contrastive loss enforces continuous and coherent trajectories in both perception and action spaces.
Empirically, DP-AG consistently outperforms state-of-the-art methods in success rate, convergence speed, and action smoothness, achieving gains of 6\% in Push-T and 13\% in Dynamic Push-T benchmarks, and delivering at least 23\% higher manipulation success and approximately 60\% smoother actions on real-world UR5 robot tasks compared to the baseline DP.
Our contributions are summarized below:
\begin{itemize}[leftmargin=*]
\vspace{-2mm}
\item We propose a novel observation representation learning that establishes a closed perception–action loop by refining latent observation features via VJP-guided noise predictions derived from DP.
\vspace{-0.5mm}
\item We formulate an action-guided latent SDE, derive a variational lower bound, and prove that cycle-consistent InfoNCE enforces mutual smoothness in both latent and action trajectories.
\vspace{-0.5mm}
\item We validate the effectiveness of DP-AG in both simulation and real-world scenarios, demonstrating consistent and significant improvements in task success rate, convergence speed, and smoothness of generated actions compared to state-of-the-art methods.
\end{itemize}
\section{Related Work}
\vspace{-1mm}
\subsection{Imitation Learning}
\vspace{-1mm}
IL provides an alternative to reinforcement learning (RL) by removing the need for explicit reward signals~\citep{sutton1999policy, osa2018algorithmic}.
Existing methods include Behavioral Cloning (BC)~\citep{bain1995framework, mandlekar22a, florence2022implicit}, which learns the mapping from observations to actions, and Inverse RL~\citep{arora2021survey} that infers implicit rewards from expert demonstrations.
Recent work~\citep{kim2024openvla, black2024pi_0} extends BC by incorporating VLMs into the perception pipeline, pushing robotic manipulation toward near-human capabilities.
However, existing methods decouple perception and action: the observation feature is frozen during each action sequence generation, without adapting to the evolving action refinements.
This constraint often leads to abrupt or discontinuous motions.
In contrast, we propose a perception–action interplay where latent continuity and action smoothness reinforce each other, yielding more coherent behavior.
% However, existing IL methods treat observations as fixed inputs throughout the generation of each action sequence, which often leads to abrupt and discontinuous motions.
% In contrast, we propose modeling a perception-action interplay where latent continuity and action smoothness reinforce each other, which enables more coherent understandings and behaviors.
\vspace{-2mm}
\subsection{Generative Models in Policy Learning}
\vspace{-1mm}
Generative models have shown promise in modeling continuous action trajectories.
DP~\citep{chi2023diffusion} leverages Denoising Diffusion Probabilistic Models (DDPMs)~\citep{ho2020denoising} to iteratively refine action predictions using diffusion models.
Policy-guided diffusion~\citep{jackson2024policy} generates synthetic trajectories aligned with a target policy for offline RL, while Diffusion-QL~\citep{wang2023diffusion} uses diffusion models to learn expressive Q-functions.
Recently, flow-matching methods~\citep{lipman2023flow, hu2024adaflow, zhang2025flowpolicy} model actions as deterministic flow ordinary differential equations (ODEs), achieving faster inference without comprising generation quality.
Building on these, DP-AG extends action-space continuity of DP into the latent observation space, forming a dynamic perception-action loop that further enhances action smoothness.
% \subsection{Diffusion Models in Policy Learning}
% Diffusion models are widely used in policy learning.
% DP~\citep{chi2023diffusion} refines discrete actions via DDPMs~\citep{ho2020denoising} to enable smooth execution.
% Policy-guided diffusion~\citep{jackson2024policy} generates synthetic trajectories aligned with a target policy to improve offline reinforcement learning, while Diffusion-QL~\citep{wang2023diffusion} leverages diffusion models to represent flexible action-value distributions for Q-learning.
% Building on these, DP-AG extends diffusion continuity from action space to latent observation space, which allows perception to evolve dynamically with interaction.
\vspace{-2mm}
\subsection{Latent Observation Modeling in Agent Control}
\vspace{-1mm}
Modeling latent observations for control has been explored through planning and generative methods. PlaNet~\citep{hafner2019learning} and Dreamer~\citep{wu2023daydreamer} use variational autoencoders (VAEs)~\citep{kingma2014auto} to simulate future states for RL.
Embed to Control~\citep{watter2015embed} uses VAEs to learn latent states from images for control policy design.
In contrast, DP-AG closes the perception-action loop by reparameterizing latent observations with VJP-guided noise from DP, extending variational inference to capture action-guided evolution.
\vspace{-2mm}
\section{Preliminaries}
\vspace{-2mm}
% In imitation learning, the goal is to learn a policy $\pi_\theta(a_t|o_t)$ that mimics expert demonstrations, where $o_t$ is the observation and $a_t$ the action at static time $t$.
In imitation learning, the goal is to learn a policy $\pi_\theta(a_t|o_t)$ that replicates expert behavior, where $o_t$ and $a_t$ denote the observation and action at a static time point $t$.
An encoder $f_\psi$ extracts static features $z_t = f_\psi(o_t)$, which are held fixed while the policy generates the corresponding $a_t$, meaning there is no feedback from action generation to perceptual features within the same short-horizon.
Given an expert demonstration dataset $\mathcal{D} = \{(o_t, a_t)\}$, standard BC minimizes the supervised loss:
\vspace{-1mm}
\begin{equation}
\mathcal{L}_{\text{BC}}(\theta, \psi) = \mathbb{E}_{(o_t, a_t) \sim \mathcal{D}} \left[ \| \pi_\theta(f_\psi(o_t)) - a_t \|^2_2 \right].
\end{equation}
Rather than the direct mapping, DP models the conditional distribution $p_{\theta}(a_t|z_t)$ as a denoising diffusion process.
Starting from white noise $\epsilon \sim \mathcal{N}(0, I)$, actions are generated over $K$ denoising steps using a learned noise predictor $\epsilon_\theta$, conditioned on observation features.
\vspace{-1mm}
\begin{equation}\label{eqn:dp_denoise}
\hat{a}_t=\hat{a}_t^{0} = \hat{a}_t^{K} - \sum_{k=1}^K g(k) \cdot \epsilon_\theta(\hat{a}_t^{k}, z_t, k),
\end{equation}
where $a_t^{K}$ is the white noise, $g(k)$ is a noise schedule, and $\epsilon_\theta$ predicts the noise at diffusion step $k$.
The diffusion model is trained by minimizing the noise matching objective~\citep{ho2020denoising}:
\vspace{-1mm}
\begin{equation}\label{eqn:dp_train}
\mathcal{L}_{\text{DP}}(\theta, \psi) = \mathbb{E}_{(o_t, a_t) \sim \mathcal{D}, k \sim \mathcal{U}(1, K)} \left[ \| \epsilon_\theta(a_t^{k}, f_\psi(o_t), k) - \epsilon \|^2_2 \right].
\end{equation}
Through the progressive refinement of noisy actions, DP estimates the continuity underlying discrete expert actions, which allows the policy to generate smooth transitions based on discrete observations.
To leverage this continuity for perception, we introduce DP-AG, which extends DP by propagating action refinement into the latent observations through Vector–Jacobian Products (VJPs) of noise predictions, enabling features to evolve dynamically alongside actions throughout diffusion.
Although we instantiate DP-AG with DDPMs, the method generalizes to other generative families such as score-based diffusion~\citep{song2021scorebased} and flow matching~\citep{lipman2023flow}, provided that intermediate action refinements are available (see Appendix~\ref{sec:flow-matching} for extending DP-AG to flow matching).
\vspace{-2mm}
\section{Our Approach}
\vspace{-2mm
}
\begin{figure}[t]
\centering
\includegraphics[width=0.99\textwidth]{imgs/NeurIPS_Interplay_Overview.pdf}
\vspace{-1mm}
\caption{{\bf Method Overview}.
While Diffusion Policy (DP) generates actions from static observation features, our DP-AG establishes a dynamic perception–action loop by guiding feature evolution via the VJP of DP’s predicted noise.
To reinforce interplay, a cycle-consistent contrastive loss aligns noise predictions from static and evolving features, enabling mutual perception–action influence.
}
\label{fig:overview}
\vspace{-2mm}
\end{figure}
We propose Action-Guided Diffusion Policy (DP-AG), a latent generative model that explicitly captures the dynamic interplay between perception and action.
As illustrated in Figure~\ref{fig:overview}, DP-AG extends static variational inference by introducing an SDE over observation features, where latent representations evolve under action-guided VJPs of the diffusion noise predictions rather than being driven by static white noise.
This replaces random perturbations with structured stochastic forces that align uncertainty in the latent dynamics with the task’s perception–action structure.
To further reinforce this interplay, we introduce a cycle-consistent contrastive loss that enforces agreement between action noise predictions conditioned on static and evolving latents, thereby closing the perception–action loop and enabling coherent updates throughout the diffusion process.
Note that, this interplay is modeled within each action sequence generation $a_t$, occurring between time points $t$ and $t+1$ across the $K$ diffusion steps that connect them.
\vspace{-2mm}
\subsection{Latent Observation Modeling via Variational Inference}
\vspace{-1mm}
In real-world settings, observations are high-dimensional and noisy, which makes it impractical to rely on deterministic representations.
To model this inherent uncertainty, we begin by constructing a variational posterior over observation features:
\vspace{-1mm}
\begin{equation}
q_\phi(z_t|o_t) = \mathcal{N}(\mu_\phi(z_t), \sigma_\phi^2(z_t)),
\end{equation}
where $\mu_\phi(z_t)$ and $\sigma^2_\phi(z_t)$ capture the mean and variance of observation features.
Sampling latent observations via the reparameterization trick~\citep{kingma2015variational}:
\vspace{-1mm}
\begin{equation}\label{eqn:vae}
z_t = \mu_\phi(z_t) + \sigma_\phi(z_t) \odot \epsilon; \quad \epsilon \sim \mathcal{N}(0, I),
\end{equation}
provides a stochastic encoding of these features, which ensures that the policy can reason about uncertainty in sensory inputs.
However, it still treats perception as static: once extracted, observation features remain fixed during the corresponding action generation and cannot adapt to refinements.
% \vspace{-2mm}
% \subsection{Diffusion-Based Action Generation}
% \vspace{-1mm}
% Smooth and coherent actions are important for manipulation, where discontinuities often cause instability.
% To leverage this continuity for perception, we adopt DP, which generates actions through gradual denoising from white noise.
% Our DP-AG extends the action continuity learned in DP to the latent observation space by leveraging Vector–Jacobian Products (VJP) to extract the gradient flow of noise predictions with respect to the latent variables.
% This enables dynamic perception that evolves with action refinement across diffusion steps.
% Though demonstrated with DDPMs, DP-AG generalizes to other generative models like score-based diffusion~\citep{song2021scorebased} and flow matching~\citep{lipman2023flow}, as long as action increments are accessible (see Appendix~\ref{sec:flow-matching} for flow matching).
\vspace{-2mm}
\subsection{Action-Guided Latent Evolution through Stochastic Differential Equations}
\vspace{-1mm}
To capture continuity from action diffusion, we model latent observation evolution using a SDE, where the drift is guided by action refinement.
At each step $k$, the latent feature $z_t$ evolves as:
\vspace{-1mm}
\begin{equation}
d\tilde{z}_t^{k} = \text{VJP}(\hat{a}_t^{k}, z_t) \, dt + \sigma_\phi(z_t) \, dW_t,
\end{equation}
where the drift term $\text{VJP}(\hat{a}_t^{k}, z_t)$ is computed from the action noise prediction:
% where $\sigma_\phi(z_t)dW_t$ models observation uncertainty, while the drift term $\text{VJP}(\hat{a}_t^{k}, z_t)$ is computed from the action noise prediction:
\vspace{-1mm}
\begin{equation}
\text{VJP}(\hat{a}_t^{k}, z_t) = \left( \frac{\partial \epsilon_\theta(\hat{a}_t^{k}, z_t, k)}{\partial z_t} \right)^\top \epsilon_\theta(\hat{a}_t^{k}, z_t, k).
\end{equation}
% This follows stochastic adjoint sensitivity methods~\citep{li2019scalable}, which show that VJP-based updates propagate gradients through SDEs efficiently without computing full Jacobians.
% Modern autodiff frameworks make these computations practical with negligible overhead.
% Intuitively, the VJP keeps latent SDE evolution directionally synchronized with action diffusion: pushing features forward during action noise addition and reversing them during action denoising.
% This phase alignment ensures that perceptual features evolve consistently with action refinement, making the VJP not only efficient but also a principled bridge between perception and action dynamics.
This follows stochastic adjoint sensitivity methods~\citep{li2019scalable}, which show that VJP-based updates efficiently propagate gradients through SDEs without computing full Jacobians; modern autodiff frameworks make these computations practical with negligible overhead.
Beyond efficiency, the VJP serves as a task-driven attentional force, guiding latent features toward the parts of the observation that most effectively reduce uncertainty in predicting the next-step action noise.
Because its drift is computed from the DP’s instantaneous noise prediction, latent evolution is phase-aligned with action diffusion: moving forward during noise addition and reversing during action denoising.
Unlike stochastic adjoint methods that separate forward evolution from backward updates, leaving latent features fixed during action inference, DP-AG’s VJP delivers immediate feedback-driven updates at every diffusion step.
This mechanism preserves the full Jacobian structure in real time, keeps latent perception features phase-aligned with ongoing action diffusion, and captures the instantaneous sensitivity between actions and observation features.
Thus, instead of a static variational posterior, we reparameterize with an action-guided noise term:
\vspace{-1mm}
\begin{equation}\label{eqn:latent_update}
\tilde{z}_t^{k} = \mu_\phi(z_t) + \gamma\sigma_\phi(z_t) \odot \text{VJP}(\hat{a}_t^{k}, z_t),
\end{equation}
where $\gamma$ controls VJP strength (see ablation studies in Appendix~\ref{sec:ablation_vjp} for its effect).
This formulation allows the latent observations to evolve with the action refinements; the variational lower bound for the corresponding optimization is detailed in Section~\ref{sec:elbo}.
\vspace{-2mm}
\subsection{Cycle-Consistent Contrastive Learning}
\vspace{-1mm}
While VJP-guided SDEs enable dynamic latent evolution, they do not guarantee coherent alignment with the action diffusion, which is significant for achieving perception–action interplay.
To bridge this gap, we introduce a cycle-consistent contrastive loss that aligns noise predictions with evolving latent dynamics at each diffusion step $k$:
\begin{itemize} [leftmargin=*]
\vspace{-2mm}
\item $\varepsilon_k=\epsilon_{\theta}(\hat{a}_t^k, z_t, k)$ is the noise prediction conditioned on the static latent $z_t$ at the step $k$;
\vspace{-1mm}
\item $\tilde{\varepsilon}_k=\epsilon_{\theta}(\hat{a}_t^k, \tilde{z}_t^k, k)$ is the noise prediction conditioned on the VJP-guided latent $\tilde{z}_t^k$ at the step $k$.
\vspace{-1mm}
\end{itemize}
An InfoNCE~\citep{chen2020simple} loss promotes the similarity of matched pairs while pushing apart mismatched ones within the same mini-batch (see Appendix~\ref{sec:ablation_cont} for temperature parameter tuning):
\vspace{-1mm}
\begin{equation}\label{eqn:contrastive}
\mathcal{L}_{\text{cont}} = -\frac{1}{B}\sum_{i=1}^B \log \frac{\exp\left( \text{sim}\left( \varepsilon^i_k, \tilde{\varepsilon}^i_k \right) / \tau \right)}{\sum_{j\neq i} \exp\left( \text{sim}\left( \varepsilon^i_k, \tilde{\varepsilon}^j_k \right) / \tau \right)},
\end{equation}
where $\text{sim}(u, v) = \frac{u^\top v}{||u| ||v||}$ is cosine similarity, with $\tau$ as the temperature and $B$ the mini-batch size.
Our contrastive loss enforces cycle consistency between perception and action: noise predictions conditioned on static latents guide continuous latent updates via VJP, and the refined latents improve subsequent noise predictions.
Unlike MSE-based objectives that rigidly match features, the contrastive loss promotes a clustering effect: it pulls evolved features toward their static counterparts to preserve semantic grounding, while pushing them away from features of other observations to maintain task-relevant distinctions.
This bounded adaptation keeps updates semantically meaningful, enabling smooth refinement without arbitrary drift and reinforcing coherence in both latent and action trajectories (see Section~\ref{sec:cont_theory} for theory and Section~\ref{sec:exp_irregular} for empirical validation).
\vspace{-2mm}
\subsection{Training Objective}
\vspace{-1mm}
The overall training objective combines three components: noise matching from DPs, variational optimization for latent evolution, and contrastive loss for the perception-action interplay:
\vspace{-1mm}
\begin{equation}
\mathcal{L}_{\text{DP-AG}} = \mathcal{L}_{\text{DP}} + \lambda_{\text{cont}} \mathcal{L}_{\text{cont}} + \lambda_{\text{KL}} \mathcal{L}_{\text{KL}}.
\end{equation}
% Minimizing $\mathcal{L}_{\text{DP}}$ and $\mathcal{L}_{\text{KL}}$ corresponds to maximizing the Evidence Lower Bound (ELBO) of our variational objective as discussed in Section~\ref{sec:elbo}.
Jointly optimizing these terms enables DP-AG to co-evolve perception and action, resulting in more adaptive policies (see Appendix~\ref{sec:ablation_kl} for ablation studies on hyperparameter tuning).
\vspace{-2mm}
\subsection{Intuitions Behind DP-AG}
\vspace{-1mm}
We now explain why DP-AG works in practice.
A static encoder yields a one-shot representation of the observation that stays fixed during action generation, forcing the policy to commit to an interpretation before actions unfold.
DP-AG instead refines the latent representation step by step, guided by feedback from the evolving action sequence.
Think of driving a car: the view through the windshield may remain unchanged, yet the driver’s focus shifts as they act, paying more attention to the curve’s edge while turning, or to nearby cars when accelerating.
The VJP supplies both the {\em direction} and {\em magnitude} of this latent drift: it points features toward the adjustments that most reduce action uncertainty, with larger updates when the model is less confident and smaller ones when it is certain.
Meanwhile, the cycle-consistent contrastive loss anchors these adaptations, keeping updated features aligned with their static counterparts and preventing excessive drift.
Together, this {\bf action-driven refinement} and {\bf bounded consistency} yield smoother, more stable, and context-aware trajectories, as confirmed in our experiments.
\vspace{-1mm}
\section{Theoretical Insights}
\vspace{-2mm}
\subsection{Action-Guided Variational Lower Bound}\label{sec:elbo}
\vspace{-1mm}
To learn a dynamic latent observation that evolves with the action diffusion, we define a posterior $q_{\phi}(\tilde{z}_t^k|z_t,\hat{a}_t^k)$ conditioned on both the static latent $z_t$ and the denoised action $\hat{a}_t^k$ at step $k$.
\textbf{Variational Objective.}
The marginal likelihood of the observed noise throughout the diffusion is:
\vspace{-1mm}
\begin{equation}
\log p(\varepsilon_k|z_t) = \log \int p(\tilde{\varepsilon}_k | \tilde{z}_t^k) \, p(\tilde{z}_t^k | z_t) \, d\tilde{z}_t^k.
\end{equation}
Applying Jensen’s inequality with the posterior, we obtain the ELBO:
\begin{equation}\label{eqn:elbo}
\log p(\varepsilon_k | z_t) \geq \mathbb{E}_{q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k)} \left[ \log p(\tilde{\varepsilon}_k | \tilde{z}_t^k) \right] - \text{KL}\left( q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k | z_t) \right).
\end{equation}
Maximizing the ELBO encourages $\tilde{z}_t^k$ to capture action-driven updates while staying close to the prior conditioned on $z_t$.
Full derivation of the ELBO is provided in Appendix~\ref{sec:elbo_proof}.
\textbf{Likelihood Term.}
The likelihood term $p(\tilde{\varepsilon}_k|\tilde{z}_t^k)$ models the distribution of action noise conditioned on the evolving latent $\tilde{z}_t^k$, and can be optimized using a noise matching loss similar to DP:
\begin{equation}\label{eqn:lh}
\mathcal{L}_{\text{LH}} = \mathbb{E}_{(o_t, a_t) \sim \mathcal{D}, k \sim \mathcal{U}(1, K)} \left[ \| \epsilon_\theta(\hat{a}_t^{k}, \tilde{z}_t^{k}, k) - \epsilon \|^2_2 \right].
\end{equation}
% This likelihood-based term promotes accurate noise prediction from the VJP-guided latent but overlaps with our contrastive loss, while the former enforces absolute noise regression, the latter enforces relative similarity.
This likelihood term encourages absolute accuracy in noise prediction but overlaps with the contrastive loss that emphasizes relative similarity (see the detailed ablation in Appendix~\ref{sec:ablation_likelihood}).
To avoid redundancy and instability, we use only the contrastive loss.
\textbf{Kullback–Leibler Divergence Term.}
The variational posterior is modeled as:
\begin{equation}
q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) = \mathcal{N}(\mu_\phi(z_t, \hat{a}_t^k), \sigma_\phi^2(z_t, \hat{a}_t^k)),
\end{equation}
while the conditional prior can be modeled as a Gaussian distribution:
\begin{equation}
p(\tilde{z}_t^k | z_t) = \mathcal{N}(z_t, I),
\end{equation}
which encourages the evolved latent to stay near the static latent initially encoded from the observation.
The Kullback–Leibler (KL) divergence thus becomes (derivation of Equation~\ref{eqn:kl} is in Appendix~\ref{sec:kl_app}):
% \vspace{-1mm}
% \begin{equation}
% \label{eqn:kl}
% \text{KL}( q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k | z_t) )
% = \frac{1}{2} \sum_{i=1}^d \left( \sigma_\phi^2(z_t, \hat{a}_t^k)_i + (\mu_\phi(z_t, \hat{a}_t^k)_i - z_{t,i})^2 - 1 - \log \sigma_\phi^2(z_t, \hat{a}_t^k)_i \right),
% \end{equation}
\vspace{-1mm}
\begin{equation}
\label{eqn:kl}
\text{KL}( q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k | z_t) )
= \frac{1}{2} \sum_{i=1}^d \left( \sigma_{\phi,i}^2 + (\mu_{\phi,i} - z_{t,i})^2 - 1 - \log \sigma_{\phi,i}^2 \right),
\end{equation}
where $d$ is the latent dimension, and $\mu_{\phi}(\cdot)$ and $\log \sigma_{\phi}^2(\cdot)$ are parameterized by a separate linear layer.
\vspace{-2mm}
\subsection{Contrastive Similarity Promotes Latent-Action Continuity}\label{sec:cont_theory}
\vspace{-1mm}
In this section, we theoretically demonstrate that minimizing our cycle-consistent InfoNCE loss leads to smooth and coherent transitions in both latent and action spaces.
The key insight is that minimizing the contrastive loss enforces strong similarity between static and dynamic noise predictions; under Lipschitz continuity, this similarity tightly bounds the drift in latent dynamics.
\textbf{Noise Simlarity Lower Bound.}
We first show that minimizing the contrastive loss imposes a lower bound on the similarity between noise predictions $\varepsilon^i_k$ and $\tilde{\varepsilon}^i_k$.
Proof of Lemma~\ref{lem:cluster} is in Appendix~\ref{sec:cluster_proof}.
\begin{lemma}[Noise Similarity Lower Bound]
\label{lem:cluster}
For unit-normalized vectors $\varepsilon^i_k$ and $\tilde{\varepsilon}^i_k$, and a temperature $\tau > 0$, if the InfoNCE loss satisfies $\mathcal{L}_{\text{cont}} \leq \alpha$ for some small constant $\alpha$, then for each $i \in \{1, \dots, B\}$, the similarity between corresponding pairs is bounded accordingly:
\begin{equation}
\underbrace{\mathrm{sim}\left(\varepsilon^i_k, \tilde{\varepsilon}^i_k\right)}_{\text{positive pair similarity}} \geq \tau (\ln(B-1) - \alpha) - 1.
\end{equation}
\vspace{-5mm}
\end{lemma}
\textbf{Continuity Induced by Noise Alignment.}
We then demonstrate that the lower bound on the noise similarity provides an upper bound on the latent drift, meaning that the contrastive alignment also leads to smooth transitions in the latent space.
Proof of Theorem~\ref{prop:mutual} is in Appendix~\ref{sec:mutual_proof}.
\begin{theorem}[Continuity Upper Bound]
\label{prop:mutual}
Suppose $\epsilon_\theta(\cdot)$ is L-Lipschitz with respect to $z$, and that $\tilde{z}_t^{k+1}$ evolves via the VJP of $\epsilon_\theta$.
Then under the contrastive constraint $\mathcal{L}_{\text{cont}} \leq \alpha$, we have:
\begin{equation}
\|\tilde{z}_t^{k+1} - \tilde{z}_t^k\|^2_2
\leq L^2\|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\|^2_2
\leq 2L^2(2 - \tau \ln(B-1) + \tau \alpha).
\end{equation}
\end{theorem}
\vspace{-2mm}
Together, Lemma~\ref{lem:cluster} and Theorem~\ref{prop:mutual} establish a theoretical connection between minimizing our cycle-consistent contrastive loss and promoting temporal continuity in both latent and action trajectories.
\vspace{-1mm}
\section{Empirical Evaluations}
\vspace{-2mm}
\subsection{Latent and Action Continuity on an Irregular Time-Series Regression}\label{sec:exp_irregular}
\vspace{-1mm}
\begin{figure}[t]
\centering
\includegraphics[width=1.0\textwidth]{imgs/synthetic.pdf}
\vspace{-5mm}
\caption{{\bf Regression results on irregular spirals.} {\bf Left}: Trajectories and latent dynamics predicted by the Base Flow. {\bf Right}: The VJP-Guided Flow continuously refines latents through output-guided corrections, which results in smoother and more coherent trajectories in both output and latent spaces.}
\label{fig:synthetic}
\vspace{-3mm}
\end{figure}
To empirically validate DP-AG’s ability to improve latent and action continuities, we conducted experiments on the irregular spiral dataset, originally introduced in the neural ODE evaluations~\citep{chen2018neural}.
This dataset is designed to evaluate continuous-time dynamics under irregular sampling, making it well-suited for evaluating the impact of the VJP-guided evolution on latent dynamics.
\textbf{Dataset.}
Following~\citep{chen2018neural}, the dataset contains 1,000 samples of 2D spiral trajectories, with half following a clockwise pattern and the other half counterclockwise, each sampled at 100 irregular time points.
The spiral exhibits continuous radial growth with Gaussian noise ($\sigma \in [0.02, 0.1]$) added to simulate real-world variability with randomly assigned angular velocities.
\textbf{Method.}
We compare two models:
\vspace{-2mm}
\begin{itemize}[leftmargin=*]
\item \textbf{Base Flow:} An LSTM with 128 hidden units and 2 layers is used to extract discrete latent features $z_t$, followed by a Multi-Layer Perceptron (MLP) to predict the regression output $y_\theta(z_t)$.
\item \textbf{VJP-Guided Flow:} Augmenting the base flow with a VJP-guided SDE, where the latent feature $\tilde{z}_t$ evolves according to: $d\tilde{z}_t = \left( \frac{\partial y_{\theta}(z_t)}{\partial z_t} \right)^\top y_\theta(z_t) \, dt + \sigma_{\phi}(z_t) \, dW_t$ with the base drift $\mu_\phi(z_t)$ and base log-variance $\log \sigma^2_\phi(z_t)$ are each predicted by a separate linear layer.
\end{itemize}
\vspace{-2.5mm}
Both models are trained for 100 epochs with Adam optimizer (learning rate $10^{-3}$ and batch size $32$).
\textbf{Results.}
As shown in Figure~\ref{fig:synthetic}, VJP-Guided Flow generates smoother and more coherent trajectories than the Base Flow, reducing the MSE from 0.0095 to 0.0052 (a 45.3\% improvement).
Importantly, latent visualizations show that while the Base Flow yields scattered latent states, VJP guidance shapes a structured latent manifold aligned with output predictions.
This illustrates our perception-action interplay: predictions provide gradient feedback that dynamically refines latent embeddings through the VJP.
As a result, latent embeddings and predictions evolve in synchrony, enhancing continuity in both representation and regression, which is consistent with our theoretical findings in Theorem~\ref{prop:mutual}.
% As shown in Figure~\ref{fig:synthetic}, VJP-Guided Flow generates smoother and more extrapolable trajectories, which reduces the MSE from 0.0095 (Base Flow) to 0.0052 (45.3\% improvement).
% Latent visualizations show that VJP-Guided Flow shapes coherent manifolds by refining the static latents toward smooth dynamics, while Base Flow latents remain scattered. This supports our claim that VJP-guided evolution improves both latent continuity and prediction smoothness.
\begin{figure}
\centering
\includegraphics[width=0.8\textwidth]{imgs/robot_datasets.pdf}
\vspace{-1mm}
\caption{Benchmark simulation environments include Robomimic, Franka Kitchen, Push-T, and Dynamic Push-T, with task and dataset details provided in Appendix~\ref{sec:dataset}.}
\label{fig:simulation}
\vspace{-3mm}
\end{figure}
\begin{table}
\centering
\fontsize{8}{10}\selectfont
\setlength{\tabcolsep}{8pt}
\caption{{\bf Target coverage score on Push-T and Dynamic Push-T tasks.}
\textbf{img} and \textbf{kp} refer to the observation modalities: RGB images or 2D keypoints.
}
\label{tab:pushT}
\begin{tabular}{llcc|c}
\toprule
Type & Method & \multicolumn{2}{c|}{Push-T} & Dynamic Push-T \\
\cmidrule(lr){3-4}
& & img & kp & img \\
\midrule
& LSTM-GMM [Mandlekar, 2022] & 0.69$\pm$0.02 & 0.67$\pm$0.03 & 0.34$\pm$1.24 \\
Mapping & IBC [Florence, 2022] & 0.75$\pm$0.02 & 0.90$\pm$0.02 & 0.52$\pm$0.98 \\
& BET [Shafiullah, 2022] & 0.80$\pm$0.02 & 0.79 $\pm$0.02& 0.58$\pm$1.35 \\
\hline
Flow & FlowPolicy [Zhang, 2025] & 0.85$\pm$0.01 & 0.88$\pm$0.01 & 0.53$\pm$0.88 \\
& AdaFlow [Hu, 2024] & 0.87$\pm$0.02 & 0.91$\pm$0.01 & 0.67$\pm$0.79 \\
% & Flow-AG (ours) & 0.91$\pm$0.01 & \textbf{1.00}$\pm$0.00 & 0.78$\pm$0.35 \\
\hline
Diffusion & DP [Chi, 2023] & 0.87$\pm$0.04 & 0.95$\pm$0.03 & 0.65$\pm$0.85 \\
\rowcolor{gray!20}
& DP-AG (ours) & \textbf{0.93}$\pm$0.02 & {\bf 0.99}$\pm$0.01 & {\bf 0.80}$\pm$0.53 \\
\hline
\end{tabular}
\vspace{-3mm}
\end{table}
\begin{table}[t]
\centering
\fontsize{7.5}{9}\selectfont
\setlength{\tabcolsep}{2.7pt}
\caption{
Success rates (ph/mh) across Robomimic, Transport, ToolHang, and Franka Kitchen tasks.
\textbf{ph}: proficient human demos; \textbf{mh}: mixed-quality demos; {\bf t1} to {\bf t4} denote task IDs in Franka Kitchen.
}
\label{tab:combined}
\begin{tabular}{llccccccccccccccc}
\toprule
Type & Method & \multicolumn{2}{c}{Lift} & \multicolumn{2}{c}{Can} & \multicolumn{2}{c}{Square} & \multicolumn{2}{c}{Transport} & ToolHang & \multicolumn{4}{c}{Franka Kitchen} \\
\cmidrule(lr){3-4} \cmidrule(lr){5-6} \cmidrule(lr){7-8} \cmidrule(lr){9-10} \cmidrule(lr){12-15}
& & ph & mh & ph & mh & ph & mh & ph & mh & ph & t1 & t2 & t3 & t4 \\
\midrule
& LSTM-GMM [Mandlekar, 2022] & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & 0.98 & 0.82 & 0.64 & 0.88 & 0.44 & 0.68 & 1.00 & 0.90 & 0.74 & 0.34 \\
Mapping & IBC [Florence, 2022] & 0.94 & 0.39 & 0.08 & 0.00 & 0.03 & 0.00 & 0.00 & 0.00 & 0.00 & 0.99 & 0.87 & 0.61 & 0.24 \\
& BET [Shafiullah, 2022] & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & 0.76 & 0.68 & 0.38 & 0.21 & 0.58 & 0.99 & 0.93 & 0.71 & 0.44 \\
\hline
Flow & FlowPolicy [Zhang, 2025] & 0.98 & 0.95 & 0.98 & 0.98 & 0.86 & 0.90 & 0.88 & 0.82 & 0.85 & 0.96 & 0.86 & 0.95 & 0.87 \\
& AdaFlow [Hu, 2024] & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & 0.96 & 0.98 & 0.96 & 0.92 & 0.80 & 0.88 & 0.99 & 0.89 & 0.92 & 0.83 \\
% \rowcolor{gray!20}
% & Flow-AG (ours) & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & 0.96 & 0.92 & 0.95 & \textbf{1.00} & 0.97 & \textbf{1.00} & 0.91 \\
\hline
Diffusion & DP [Chi, 2023] & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & 0.98 & 0.98 & \textbf{1.00} & 0.89 & 0.95 & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & 0.99 \\
\rowcolor{gray!20}
& DP-AG (ours) & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{0.94} & \textbf{0.98} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} & \textbf{1.00} \\
\hline
\end{tabular}
\vspace{-3mm}
\end{table}
\vspace{-2mm}
\subsection{Experiments on Simulation Benchmarks}
\vspace{-1mm}
We then evaluate our DP-AG on the simulation benchmarks, illustrated in Figure~\ref{fig:simulation}.
\textbf{Robomimic.}
Robomimic~\citep{mandlekar22a} is a large-scale benchmark for robotic manipulation consisting of five tasks with nine datasets: Can (ph/mh), Square (ph/mh), Transport (ph/mh), Tool Hang (ph), and Lift (ph/mh).
Here, {\bf ph} denotes proficient human demonstrations, and {\bf mh} indicates a mix of proficient and non-proficient demonstrations.
% Robomimic~\citep{mandlekar22a} is a large-scale benchmark for robotic manipulation, which consists of five tasks: Can, Square, Transport, Tool Hang, and Lift.
% Each task includes expert demonstrations collected via proficient human ({\bf ph}) teleoperation. Additionally, Lift, Can, Square, and Transport provide mixed human ({\bf mh}) variants that combine proficient and non-proficient demonstrations, yielding a total of nine datasets.
\textbf{Franka Kitchen.}
Franka Kitchen~\citep{gupta2020} is a simulation benchmark with a 9-DoF Franka arm performing four household tasks ({\bf t1}, {\bf t2}, {\bf t3}, and {\bf t4}) per trajectory, using 566 human demonstrations across 7 interactive objects.
Available through platforms like Gymnasium-Robotics~\footnote{\url{https://robotics.farama.org/envs/franka_kitchen/franka_kitchen/}}.
% it serves as an important simulation platform for studying robotic manipulation.
\textbf{Push-T.}
Push-T~\citep{chi2023diffusion} is a manipulation task adapted from IBC~\citep{florence2022implicit}, where a circular end-effector pushes a T-shaped block to a target location.
Both the block and end-effector start at random positions.
Observations consist of either RGB images ({\bf img}) or a set of nine 2D keypoints ({\bf kp}) outlining the T block’s shape, along with the position of the end-effector.
\textbf{Dynamic Push-T (Ours).}
Many IL benchmarks are nearly saturated: they are deterministic, scripted, and lack diversity, often solvable without closed-loop feedback~\citep{jia2024towards}.
To evaluate real-time adaptability, we propose \emph{Dynamic Push-T}, which augments Push-T with a moving ball that bounces unpredictably and intermittently interferes with manipulation.
Since the ball’s trajectory varies in each episode, the agent must adapt online, blocking or avoiding the ball while pursuing the main objective.
This develops a dynamic and unscripted challenge that evaluates true adaptability beyond replaying demonstrations.
Details of this benchmark are provided in Appendix~\ref{sec:dataset}.
\begin{figure}
\centering
\includegraphics[width=1.0\textwidth]{imgs/convergence.pdf}
\caption{{\bf Convergence Plots.} Training action MSE over epochs on Robomimic Push-T and Can.}
\label{fig:convergence}
\vspace{-4mm}
\end{figure}
\textbf{Baselines and Evaluation Metrics.}
We compare DP-AG against IL baselines that include direct mapping (LSTM-GMM~\citep{mandlekar22a}, IBC~\citep{florence2022implicit}, BET~\citep{shafiullah2022behavior}), diffusion-based (DP~\citep{chi2023diffusion}), and flow-matching methods (AdaFlow~\citep{hu2024adaflow}, FlowPolicy~\citep{zhang2025flowpolicy}).
Results are from either our reimplementation or the original papers, averaged over 5 training seeds and multiple evaluation seeds (50 for Push-T/Kitchen, 22 for Robomimic, 50 for Dynamic Push-T).
We use success rate as the main metric, with target area coverage for Push-T variants.
Additional implementation details are in Appendix~\ref{sec:implementation}.
\textbf{Results.}
Tables~\ref{tab:pushT} and~\ref{tab:combined} show that DP-AG consistently outperforms baselines across Robomimic, Franka Kitchen, Push-T, and Dynamic Push-T, achieving near-perfect success on static tasks and the highest coverage on dynamic ones.
These gains come from enforcing smooth latent and action trajectories, which are especially important in Dynamic Push-T.
DP-AG also converges faster than DP (Figure~\ref{fig:convergence} with additional plots in Appendix~\ref{sec:more_convergence}).
Ablation results are provided in Appendix~\ref{sec:ab}, and computational analysis is detailed in Appendix~\ref{sec:ablation_computation}.
\vspace{-2mm}
\subsection{Real-World Evaluation on UR5 Robot Arm}
\vspace{-1mm}
\begin{figure}
\centering
\includegraphics[width=1.0\textwidth]{imgs/ur5-real.pdf}
\vspace{-5mm}
\caption{Real-world evaluation on a UR5 robot arm across three manipulation tasks.}
\label{fig:realworld}
\vspace{-5mm}
\end{figure}
To evaluate the real-world performance of our DP-AG, we deploy DP and DP-AG policies onto a UR5 robotic arm in three visuomotor manipulation tasks: \textbf{Painting}, \textbf{Candy Push}, and \textbf{Peg-In-Hole Insertion} (Figure~\ref{fig:realworld}).
These tasks are designed to evaluate different aspects of perception–action interplay and the policy’s ability to generalize from partially observable inputs.
The details (dataset, training, etc.) of the real-world manipulation experiments are provided in Appendix~\ref{sec:real-world}.
\textbf{Painting (Planar Precision).} The robot is tasked with tracing heart-shaped and circular-shaped paths using a paintbrush.
This task requires smooth and continuous motion to avoid over-painting or streaking.
Our DP-AG improves both path fidelity and smoothness compared to DP.
\textbf{Candy Push (Object-aware Adaptation).}
The end-effector pushes small candies into a designated goal area.
Object positions are randomized across trials.
DP-AG’s latent evolution allows it to adapt to variations in candy layouts and locations, which results in fewer collisions and smoother trajectories.
\textbf{Peg-in-Hole (3D Visual Reasoning).}
The robot must insert a circular peg into a vertical hole using only RGB inputs from scene and wrist cameras, without explicit depth sensing, requiring the policy to infer 3D geometry from indirect cues.
Baseline DP offers no mechanism to adapt when the hole is slightly misaligned.
Upon contact, it repeatedly executes the same ineffective motion, failing in all trials.
In contrast, DP-AG leverages action feedback: each blocked step triggers a VJP-guided latent update that sharpens sensitivity to contact-region cues ({\em e.g.,} rim alignment).
Over a few refinements, the same RGB inputs are reinterpreted to expose geometry cues sufficient for correction, enabling successful insertions under occlusion and unseen perturbations.
% The robot should insert a circular peg into a vertical hole on a block using only RGB images from front and wrist-mounted cameras, without any explicit depth sensing.
% This requires the policy to infer 3D observations from indirect visual cues and proprioceptive inputs.
% DP-AG leverages action feedback to infer depth over time, which enables successful insertions despite occlusion or ambiguity.
% In contrast, DP fails to complete this task in all trials.
\begin{table}
\centering
% \footnotesize
\fontsize{8.55}{10}\selectfont
\setlength{\tabcolsep}{6pt}
\caption{{\bf Performance on Real-World UR5 Tasks.} Mean and standard deviation are reported.}
\label{tab:realworld}
\begin{tabular}{llcccc}
\toprule
Task & Method & Success Rate (\%) & Smoothness (Avg. Jerk) & IoU (\%) & Time to Complete (s) \\
\hline
Painting & DP & – & $0.083 \pm 0.014$ & $68.9 \pm 5.2$ & $49.5 \pm 4.1$ \\
\rowcolor{gray!20}
& DP-AG & – & $\mathbf{0.032 \pm 0.009}$ & $\mathbf{92.1 \pm 3.4}$ & $\mathbf{18.0 \pm 3.2}$ \\
\hline
Candy Push & DP & $65.0 \pm 8.4$ & $0.107 \pm 0.016$ & – & $24.0 \pm 3.9$ \\
\rowcolor{gray!20}
& DP-AG & $\mathbf{90.0 \pm 5.5}$ & $\mathbf{0.039 \pm 0.011}$ & – & $\mathbf{9.5 \pm 2.6}$ \\
\hline
Peg-in-Hole & DP & $0.0 \pm 0.0$ & $0.096 \pm 0.017$ & – & – \\
\rowcolor{gray!20}
& DP-AG & $\mathbf{85.0 \pm 6.0}$ & $\mathbf{0.036 \pm 0.008}$ & – & $\mathbf{13.0 \pm 2.1}$ \\
\hline
\end{tabular}
\vspace{-5mm}
\end{table}
\textbf{Evaluation.}
We evaluate our DP-AG against DP on the three tasks, each repeated over 20 trials.
For Painting, we report IoU between the painted and target shapes;
for Candy Push and Peg-In-Hole, we report success rate.
Moreover, we measure trajectory smoothness via average jerk and task completion time.
As shown in Table~\ref{tab:realworld}, DP-AG consistently outperforms DP across all metrics.
Notably, the largest gain occurs in the peg-in-hole task, where DP-AG’s dynamic latent updates help infer 3D information from 2D observations, which is an ability DP lacks due to its static features.
\vspace{-1mm}
\section{Conclusion}
\vspace{-2mm}
We introduce DP-AG, a representation learning framework that closes the perception–action loop in diffusion policies.
By evolving latent observations through an action-guided SDE, DP-AG transforms diffusion noise gradients into structured perceptual updates via VJPs.
A cycle-consistent contrastive loss aligns static and evolving latents, enabling continuous and bidirectional coupling between perception and action throughout diffusion.
We derive a principled ELBO and prove that contrastive alignment enforces mutual continuity in latent and action trajectories.
Empirically, DP-AG achieves state-of-the-art performance, especially under partial observability and dynamic conditions, demonstrating that perception–action interplay is significant for effective and adaptive policy learning.
\textbf{Broader Impacts.}
DP-AG advances imitation learning by modeling dynamic perception-action interplay, enabling smoother and more context-aware robotic manipulation and reducing the risk of failure or unintended motions.
This capability is particularly valuable in safety-critical domains such as automation, manufacturing, and assistive robotics, where robust decision-making is significant for effective human-robot collaboration.
However, relying on expert demonstrations can introduce biases or suboptimal behaviors into the policy. To mitigate these risks, it is important to rigorously validate the quality and diversity of demonstration data to minimize unintended actions.
\begin{ack}
This work was partly supported by NSERC Discovery, CFI-JELF, NSERC Alliance, Alberta Innovates and PrairiesCan grants.
\end{ack}
%%%%%%%%%%%%%%%%%%%%%%%%%% Reference %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\bibliography{bibtex}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section*{NeurIPS Paper Checklist}
%%% BEGIN INSTRUCTIONS %%%
% The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: {\bf The papers not including the checklist will be desk rejected.} The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page
% limit.
% Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
% \begin{itemize}
% \item You should answer \answerYes{}, \answerNo{}, or \answerNA{}.
% \item \answerNA{} means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
% \item Please provide a short (1–2 sentence) justification right after your answer (even for NA).
% % \item {\bf The papers not including the checklist will be desk rejected.}
% \end{itemize}
% {\bf The checklist answers are an integral part of your paper submission.} They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
% The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "\answerYes{}" is generally preferable to "\answerNo{}", it is perfectly acceptable to answer "\answerNo{}" provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "\answerNo{}" or "\answerNA{}" is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer \answerYes{} to a question, in the justification please point to the section(s) where related material for the question can be found.
% IMPORTANT, please:
% \begin{itemize}
% \item {\bf Delete this instruction block, but keep the section heading ``NeurIPS Paper Checklist"},
% \item {\bf Keep the checklist subsection headings, questions/answers and guidelines below.}
% \item {\bf Do not modify the questions and only use the provided macros for your answers}.
% \end{itemize}
%%% END INSTRUCTIONS %%%
\begin{enumerate}
\item {\bf Claims}
\item[] Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the abstract and introduction do not include the claims made in the paper.
\item The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
\item The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
\item It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
\end{itemize}
\item {\bf Limitations}
\item[] Question: Does the paper discuss the limitations of the work performed by the authors?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper discusses the limitations of the proposed work in the Broader Impacts section.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
\item The authors are encouraged to create a separate "Limitations" section in their paper.
\item The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
\item The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
\item The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
\item The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
\item If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
\item While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
\end{itemize}
\item {\bf Theory assumptions and proofs}
\item[] Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper provides the full set of assumptions and a complete and correct proof for each theoretical result in the Appendix.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include theoretical results.
\item All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
\item All assumptions should be clearly stated or referenced in the statement of any theorems.
\item The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
\item Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
\item Theorems and Lemmas that the proof relies upon should be properly referenced.
\end{itemize}
\item {\bf Experimental result reproducibility}
\item[] Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The paper fully discloses all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper. More importantly, the paper provides the sufficient details of the conducted real-world experiments including dataset collections and model implementation. The code is also provided in the supplement materials.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
\item If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
\item Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
\item While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
\begin{enumerate}
\item If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
\item If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
\item If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
\item We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
\end{enumerate}
\end{itemize}
\item {\bf Open access to data and code}
\item[] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper provides open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material.
The code and data are available on our project website: \href{https://jingwang18.github.io/dp-ag.github.io/}{https://jingwang18.github.io/dp-ag.github.io/}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that paper does not include experiments requiring code.
\item Please see the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
\item While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
\item The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
\item The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
\item The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
\item At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
\item Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
\end{itemize}
\item {\bf Experimental setting/details}
\item[] Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper specifies all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results for all experiments.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
\item The full details can be provided either with the code, in appendix, or as supplemental material.
\end{itemize}
\item {\bf Experiment statistical significance}
\item[] Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper reports error bars suitably and correctly defined or other appropriate information about the statistical significance for all experiments including the real-world experiments.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
\item The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
\item The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
\item The assumptions made should be given (e.g., Normally distributed errors).
\item It should be clear whether the error bar is the standard deviation or the standard error of the mean.
\item It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96\% CI, if the hypothesis of Normality of errors is not verified.
\item For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
\item If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
\end{itemize}
\item {\bf Experiments compute resources}
\item[] Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper provides sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce each experiment.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
\item The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
\item The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
\end{itemize}
\item {\bf Code of ethics}
\item[] Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics \url{https://neurips.cc/public/EthicsGuidelines}?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
\item If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
\item The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
\end{itemize}
\item {\bf Broader impacts}
\item[] Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, the paper discusses both potential positive societal impacts and negative societal impacts of the work performed in the Broader Impacts section.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that there is no societal impact of the work performed.
\item If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
\item Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
\item The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
\item The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
\item If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
\end{itemize}
\item {\bf Safeguards}
\item[] Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The paper poses no such risks.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper poses no such risks.
\item Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
\item Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
\item We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
\end{itemize}
\item {\bf Licenses for existing assets}
\item[] Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The paper does not use existing assets.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not use existing assets.
\item The authors should cite the original paper that produced the code package or dataset.
\item The authors should state which version of the asset is used and, if possible, include a URL.
\item The name of the license (e.g., CC-BY 4.0) should be included for each asset.
\item For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
\item If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, \url{paperswithcode.com/datasets} has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
\item For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
\item If this information is not available online, the authors are encouraged to reach out to the asset's creators.
\end{itemize}
\item {\bf New assets}
\item[] Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The paper does not release new assets.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not release new assets.
\item Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
\item The paper should discuss whether and how consent was obtained from people whose asset is used.
\item At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
\end{itemize}
\item {\bf Crowdsourcing and research with human subjects}
\item[] Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The paper does not involve crowdsourcing nor research with human subjects.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
\item Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
\item According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
\end{itemize}
\item {\bf Institutional review board (IRB) approvals or equivalent for research with human subjects}
\item[] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The paper does not involve crowdsourcing nor research with human subjects.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
\item Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
\item We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
\item For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
\end{itemize}
\item {\bf Declaration of LLM usage}
\item[] Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
%this research?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components. LLMs are used solely for language refinement.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
\item Please refer to our LLM policy (\url{https://neurips.cc/Conferences/2025/LLM}) for what should or should not be described.
\end{itemize}
\end{enumerate}
\newpage
\appendix
\section*{Appendix}
\addcontentsline{toc}{section}{Appendix}
\section{Extension to Flow Matching Models}\label{sec:flow-matching}
While we demonstrate DP-AG built upon DDPMs (as in DP), the proposed perception-action interplay can be extended to alternative generative models that produce smooth action trajectories, such as flow matching models~\citep{lipman2023flow, black2024pi_0, zhang2025flowpolicy}.
\textbf{Background on Flow Matching.}
Flow matching formulates action generation as solving an ordinary differential equation (ODE) that deterministically transforms a sample from a known base distribution (e.g., Gaussian) to the target data distribution.
Given a continuous time variable $t\in[0,1]$, the action $a(t)$ evolves according to:
\begin{equation}
\frac{d a(t)}{dt} = v_\theta(a(t), z_t, t),
\end{equation}
where $v_\theta$ denotes a learned velocity field conditioned on observation features $z_t$.
The goal is to match the score of the intermediate marginal distribution $p_t(a)$ at time $t$ by minimizing a regression loss.
Unlike diffusion models, flow matching generates continuous evolution directly without stochastic perturbations.
\textbf{Action-Guided Latent Evolution with Flow Matching.}
To extend DP-AG to flow matching models, we leverage the structure of the learned velocity field $v_\theta$ to guide the latent dynamics.
To be specific, at each time $t$, we define the evolution of the observation feature $\tilde{z}_t$ through a VJP-guided ODE:
\begin{equation}
\frac{d \tilde{z}_t}{dt} = \left( \frac{\partial v_\theta(a(t), z_t, t)}{\partial z_t} \right)^\top v_\theta(a(t), z_t, t),
\end{equation}
where $v_\theta(a(t), z_t, t)$ serves as the action-conditioned driving force, and the VJP propagates this force back to update the observation feature $\tilde{z}_t$.
This setup follows our DP-AG formulation for DDPMs but replaces the stochastic increments from noise predictions with deterministic updates guided by flow matching dynamics.
\textbf{Training Objective.}
Similar to the DP-AG framework, we introduce a variational posterior $q_\phi(\tilde{z}_t | z_t, a(t))$ over the evolved latents and formulate a ELBO:
\begin{equation}
\mathcal{L}_{\text{ELBO}} = \mathbb{E}{q_\phi}\left[ \log p_\theta(v_\theta(a(t), \tilde{z}_t, t)) \right] - \text{KL}\left(q_\phi(\tilde{z}_t | z_t, a(t)) || p(\tilde{z}_t|z_t)\right),
\end{equation}
where $p(\tilde{z}_t|z_t)$ is the base distribution, e.g., isotropic Gaussian centered at $z_t$.
Moreover, the cycle-consistent contrastive loss can be formulated between the velocity fields $v_\theta(a(t), z_t, t)$ and $v_\theta(a(t), \tilde{z}_t, t)$ to promote the consistent evolution between the action and the latent trajectories.
\textbf{Discussion.}
The flow matching extension has several benefits:
\begin{itemize}[leftmargin=*]
\item Deterministic evolution improves sample efficiency compared to stochastic DDPM training.
\item The continuous latent trajectories can be directly controlled via $v_\theta$, which simplifies the interpretation.
\item The perception-action interplay remains: smoother latent evolution produces more coherent actions, and better action flow, in turn, improves the latent dynamics.
\end{itemize}
Thus, DP-AG extends beyond DDPMs, which provide a flexible framework for integrating with generative-model-based policy learning under both stochastic and deterministic dynamics.
\section{Derivation of DP-AG Variational Lower Bound}\label{sec:elbo_proof}
In this section, we show the full derivation of our ELBO using Jensen's inequality:
\begin{equation}
\log p(\varepsilon_k | z_t) \geq \mathbb{E}_{q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k)} \left[ \log p(\tilde{\varepsilon}_k | \tilde{z}_t^k) \right] - \text{KL}\left( q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k | z_t) \right),
\end{equation}
where $q_\phi(\tilde{z}_t^k|z_t, \hat{a}_t^k)$ is the variational posterior approximating the true posterior.
We start from the exact marginal:
\begin{equation}
p(\varepsilon_k | z_t) = \int p(\tilde{\varepsilon}_k | \tilde{z}_t^k) p(\tilde{z}_t^k | z_t) \, d\tilde{z}_t^k.
\end{equation}
Inserting the variational posterior $q_\phi(\tilde{z}_t^k|z_t, \hat{a}_t^k)$ via importance sampling:
\begin{equation}
p(\varepsilon_k | z_t) = \int q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \frac{p(\tilde{\varepsilon}_k|\tilde{z}_t^k)p(\tilde{z}_t^k|z_t)}{q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k)} \, d\tilde{z}_t^k = \mathbb{E}_{q_\phi}\left[ \frac{p(\tilde{\varepsilon}_k|\tilde{z}_t^k)p(\tilde{z}_t^k|z_t)}{q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k)} \right].
\end{equation}
Thus, we have:
\begin{equation}
\log p(\varepsilon_k | z_t) = \log \mathbb{E}_{q_\phi}\left[ \frac{p(\tilde{\varepsilon}_k|\tilde{z}_t^k)p(\tilde{z}_t^k|z_t)}{q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k)} \right].
\end{equation}
Since $\log(\cdot)$ is a concave function, by Jensen’s inequality:
\begin{equation}
\log \mathbb{E}_{q_\phi}[X] \geq \mathbb{E}_{q_\phi}[\log X],
\end{equation}
for any non-negative random variable X.
Here, $X = \frac{p(\tilde{\varepsilon}_k|\tilde{z}_t^k)p(\tilde{z}_t^k|z_t)}{q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k)}$ is clearly non-negative because all densities are non-negative.
Thus, applying Jensen’s inequality to our case gives:
\begin{equation}
\begin{aligned}
\log \mathbb{E}_{q_\phi}\left[ \frac{p(\tilde{\varepsilon}_k|\tilde{z}_t^k)p(\tilde{z}_t^k|z_t)}{q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k)} \right]
&\geq \mathbb{E}_{q_\phi}\left[ \log \left( \frac{p(\tilde{\varepsilon}_k|\tilde{z}_t^k)p(\tilde{z}_t^k|z_t)}{q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k)} \right) \right] \\
&= \mathbb{E}_{q_\phi}\left[ \log p(\tilde{\varepsilon}_k|\tilde{z}_t^k) + \log p(\tilde{z}_t^k|z_t) - \log q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k) \right] \\
&= \mathbb{E}_{q_\phi}\left[ \log p(\tilde{\varepsilon}_k|\tilde{z}_t^k) \right] + \mathbb{E}_{q_\phi}\left[ \log p(\tilde{z}_t^k|z_t) \right] - \mathbb{E}_{q_\phi}\left[ \log q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k) \right], \\
&= \mathbb{E}_{q_\phi}\left[ \log p(\tilde{\varepsilon}_k|\tilde{z}_t^k) \right] - \text{KL}\left( q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k | z_t) \right).
\end{aligned}
\end{equation}
Thus, we have proven:
\begin{equation}
\log p(\varepsilon_k | z_t) \geq \mathbb{E}_{q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k)}\left[ \log p(\tilde{\varepsilon}_k|\tilde{z}_t^k) \right] - \text{KL}\left( q_\phi(\tilde{z}_t^k | z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k | z_t) \right).
\end{equation}
$\square$
\section{Derivation and Intuition for the KL Divergence Term in Equation~\ref{eqn:kl}}\label{sec:kl_app}
To regularize the latent evolution $\tilde{z}_t^k$ and preserve semantic consistency with the observation-encoded latent $z_t$, we introduce a KL divergence between a variational posterior and a conditional prior.
This section provides a step-by-step derivation and intuition for the expression in Equation~\ref{eqn:kl}.
\textbf{Formulation.} The variational posterior is defined as a Gaussian distribution:
\begin{equation}
q_\phi(\tilde{z}_t^k|z_t,\hat{a}_t^k) = \mathcal{N}(\mu_\phi(z_t, \hat{a}_t^k), \operatorname{diag}(\sigma^2_\phi(z_t, \hat{a}_t^k))),
\end{equation}
where both the mean and log-variance are predicted by linear layers.
The conditional prior is set to:
\begin{equation}
p(\tilde{z}_t^k|z_t) = \mathcal{N}(z_t, I),
\end{equation}
which enforces that the evolved latent remains close to its observation-driven anchor $z_t$ unless strongly influenced by the action.
\textbf{KL Between Diagonal Gaussians.}
The KL divergence between two diagonal Gaussians $\mathcal{N}(\mu_q, \Sigma_q)$ and $\mathcal{N}(\mu_p, \Sigma_p)$ is given by:
\begin{equation}
\text{KL}(q \| p) = \frac{1}{2} \sum_{i=1}^d \left[
\frac{\sigma_{q,i}^2}{\sigma_{p,i}^2}
+ \frac{(\mu_{q,i} - \mu_{p,i})^2}{\sigma_{p,i}^2}
- 1 - \log \frac{\sigma_{q,i}^2}{\sigma_{p,i}^2}
\right],
\end{equation}
where $d$ is the dimensionality of samples drawn from Guassians.
Substituting $\mu_q = \mu_\phi(z_t, \hat{a}_t^k)$, $\sigma_q=\sigma_{\phi}(z_t,\hat{a}_t^k)$, $\mu_p = z_t$, and $\Sigma_p = I$, we obtain:
\begin{equation}
\text{KL}( q_\phi(\tilde{z}_t^k|z_t, \hat{a}_t^k) \, \| \, p(\tilde{z}_t^k|z_t) ) =
\frac{1}{2} \sum_{i=1}^d \left( \sigma_{\phi,i}^2 + (\mu_{\phi,i} - z_{t,i})^2 - 1 - \log \sigma_{\phi,i}^2 \right),
\end{equation}
where $d$ is now the latent dimension, and all quantities are computed element-wise.
\textbf{Intuitions.} Each term in the KL expression serves a distinct regularization role:
\begin{itemize}[leftmargin=*]
\item $(\mu_{\phi,i} - z_{t,i})^2$: penalizes deviation of the evolved latent mean from the observation-encoded static latent.
\item $\sigma_{\phi,i}^2$: penalizes over-dispersion and encourages confident latent representations.
\item $\log \sigma_{\phi,i}^2$: rewards expressive uncertainty when needed, balancing the previous term.
\end{itemize}
This KL divergence term acts as a geometric constraint to preserve alignment with the observation-anchored latent space, while still allowing dynamic evolution based on the action.
Combined with our contrastive loss (Equation~\ref{eqn:contrastive}), it stabilizes training by maintaining local continuity and global discriminability.
\section{Proof of Lemma~\ref{lem:cluster}}\label{sec:cluster_proof}
Following the theoretical work on the InfoNCE loss~\citep{parulekar2023infonce,wang2024what}, we derive the lower bound for the positive key alignment:
\begin{lemma*}[Noise Similarity Lower Bound]
For unit-normalized vectors $\varepsilon^i_k$ and $\tilde{\varepsilon}^i_k$, and a temperature $\tau > 0$, if the InfoNCE loss satisfies $\mathcal{L}_{\text{cont}} \leq \alpha$ for some small constant $\alpha$, then for each $i \in \{1, \dots, B\}$, the similarity between corresponding pairs is bounded accordingly:
\begin{equation}
\underbrace{\mathrm{sim}\left(\varepsilon^i_k,\tilde{\varepsilon}^i_k\right)}_{\text{positive pair similarity}} \geq \tau (\ln(B-1) - \alpha) - 1.
\end{equation}
\end{lemma*}
\begin{proof}
Recall that the InfoNCE loss is defined as:
\begin{equation}
\mathcal{L}_{\text{cont}} = -\frac{1}{B} \sum_{i=1}^B \log \frac{\exp\left( \mathrm{sim}\left(\varepsilon^i_k, \tilde{\varepsilon}^i_k\right)/\tau \right)}{\sum_{j \neq i} \exp\left( \mathrm{sim}\left(\varepsilon^i_k, \tilde{\varepsilon}^j_k\right)/\tau \right)}.
\end{equation}
To simplify the notations, let $s_{ii} = \mathrm{sim}(\varepsilon^i_k, \tilde{\varepsilon}^i_k)$ and $s_{ij} = \mathrm{sim}(\varepsilon^i_k, \tilde{\varepsilon}^j_k)$ for $j \neq i$.
We define:
\begin{equation}
p_i = \frac{\exp(s_{ii}/\tau)}{\sum_{j \neq i} \exp(s_{ij}/\tau)}.
\end{equation}
The loss constraint $\mathcal{L}_{\text{cont}} \leq \alpha$ implies:
\begin{equation}
-\frac{1}{B} \sum_{i=1}^B \log p_i \leq \alpha \implies \sum_{i=1}^B \log p_i \geq -B \alpha.
\end{equation}
Taking exponential at both sides, we have:
\begin{equation}
\prod_{i=1}^B p_i \geq \exp(-B\alpha).
\end{equation}
To derive a lower bound on $s_{ii}$, we need the denominator $\sum_{j \neq i} \exp(s_{ij}/\tau)$ to be as small as possible. Since $s_{ij} \geq -1$ (unit vectors), the denominator is minimized when all $s_{ij} = -1$. Thus, we have:
\begin{equation}
\sum_{j \neq i} \exp(s_{ij}/\tau) \geq (B-1)\exp(-1/\tau),
\end{equation}
and
\begin{equation}\label{eqn:inequal}
p_i \leq \frac{\exp(s_{ii}/\tau)}{(B-1)\exp(-1/\tau)}=\frac{1}{B-1} \exp\left(\frac{s_{ii} + 1}{\tau}\right),
\end{equation}
Assume that $p_i$ is equal across all $i$ to maximize the product under the constraint by Arithmetic Mean-Geometric Mean (AM-GM) inequality~\citep{tan2020extension}:
\begin{equation}
p_i = \left( \prod_{i=1}^B p_i \right)^{1/B} \geq \exp(-B\alpha)^{1/B} = \exp(-\alpha).
\end{equation}
Substituting into Inequality~\ref{eqn:inequal}, we have:
\begin{equation}
\exp(-\alpha) \leq p_i \leq \frac{1}{B-1} \exp\left(\frac{s_{ii} + 1}{\tau}\right).
\end{equation}
Focusing on the inequality:
\begin{equation}
\exp(-\alpha) \leq \frac{1}{B-1} \exp\left(\frac{s_{ii} + 1}{\tau}\right).
\end{equation}
Multiply both sides by \( B-1 \):
\begin{equation}
(B-1) \exp(-\alpha) \leq \exp\left(\frac{s_{ii} + 1}{\tau}\right).
\end{equation}
Take the natural logarithm of both sides:
\begin{equation}
\ln\left((B-1) \exp(-\alpha)\right) \leq \frac{s_{ii} + 1}{\tau}.
\end{equation}
Simplify the left-hand side:
\begin{equation}
\ln(B-1) + \ln(\exp(-\alpha)) = \ln(B-1) - \alpha.
\end{equation}
Thus:
\begin{equation}
\ln(B-1) - \alpha \leq \frac{s_{ii} + 1}{\tau}.
\end{equation}
Multiply through by \( \tau \):
\begin{equation}
\tau (\ln(B-1) - \alpha) \leq s_{ii} + 1.
\end{equation}
Finally, we have the lower bound for \( s_{ii} \):
\begin{equation}
s_{ii} \geq \tau (\ln(B-1) - \alpha) - 1.
\end{equation}
Thus, we have proved:
\begin{equation}
\mathrm{sim}\left(\varepsilon^i_k, \tilde{\varepsilon}^i_k\right) \geq \tau (\ln(B-1) - \alpha) - 1.
\end{equation}
\end{proof}
\section{Proof of Theorem~\ref{prop:mutual}}\label{sec:mutual_proof}
\begin{theorem*}[Continuity Upper Bound]
Suppose $\epsilon_\theta(\cdot)$ is L-Lipschitz with respect to $z$, and that $\tilde{z}_t^{k+1}$ evolves via the VJP of $\epsilon_\theta$.
Then under the contrastive constraint $\mathcal{L}_{\text{cont}} \leq \alpha$, we have:
\begin{equation}
\|\tilde{z}_t^{k+1} - \tilde{z}_t^k\|^2_2
\leq L^2\|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\|^2_2
\leq 2L^2(2 - \tau \ln(B-1) + \tau \alpha).
\end{equation}
\end{theorem*}
\begin{proof}
We split the proof into two steps:
\textbf{Step 1: Lipschitz continuity of \(\epsilon_\theta\) relates latent updates and noise difference.}
Taking the vector-Jacobian product (VJP) approximation for the drift of the SDE evolution, the update of $\tilde{z}_t^k$ at each step satisfies:
\begin{equation}
\|\tilde{z}_t^{k+1} - \tilde{z}_t^k\|_2 \;\propto\; \text{VJP}(\hat{a}_t^k, z_t),
\end{equation}
where, under the Lipschitz continuity assumption, the VJP depends linearly on the local noise gradient structure.
Assuming proper normalization of the step size in the SDE discretization (absorbed into the $L$-Lipschitz constant), we have:
\begin{equation}
\|\tilde{z}_t^{k+1} - \tilde{z}_t^k\|_2 \;\leq\; L \|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\|_2.
\end{equation}
Squaring both sides yields:
\begin{equation}
\|\tilde{z}_t^{k+1} - \tilde{z}_t^k\|^2_2 \;\leq\; L^2\,\|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\|^2_2.
\end{equation}
\textbf{Step 2: Contrastive loss bounds the noise difference.}
From Lemma~\ref{lem:cluster}, when $\mathcal{L}_{\mathrm{cont}} \leq \alpha$, the cosine similarity between $\varepsilon^i_k$ and $\tilde{\varepsilon}^i_k$ is bounded below:
\begin{equation}
\text{sim}\left(\varepsilon^i_k, \tilde{\varepsilon}^i_k\right) \geq \tau (\ln(B-1) - \alpha) - 1.
\end{equation}
Recall that cosine similarity between two unit-norm vectors \(u, v\) satisfies:
\begin{equation}
\|u-v\|^2_2 = 2(1 - \text{sim}(u,v)).
\end{equation}
Thus, applying this relation to $\varepsilon_k$ and $\tilde{\varepsilon}_k$, we obtain:
\begin{equation}
\|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^k, k)\|^2_2 = 2\left(1 - \text{sim}\left(\epsilon_\theta(\hat{a}_t^k, z_t, k), \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\right)\right).
\end{equation}
By plugging in the lower bound on similarity from Lemma~\ref{lem:cluster}, we get:
\begin{equation}
\|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\|^2_2 \leq 2(2 - \tau\ln(B-1) + \tau\alpha).
\end{equation}
Thus, combining Step 1 and Step 2, we can conclude that:
\begin{equation}
\|\tilde{z}_t^{k+1} - \tilde{z}_t^k\|^2_2
\;\leq\; L^2\,\|\epsilon_\theta(\hat{a}_t^k, z_t, k) - \epsilon_\theta(\hat{a}_t^k, \tilde{z}_t^{k}, k)\|^2_2
\;\leq\; 2L^2(2 - \tau\ln(B-1) + \tau\alpha).
\end{equation}
This completes the proof.
\end{proof}
\section{Robotic Manipulation Benchmarks}\label{sec:dataset}
This ablation section details the tasks and corresponding datasets for four robotic manipulation benchmarks for simulation: Robomimic, Franka Kitchen, Push-T, and Dynamic Push-T.
Each benchmark contributes unique tasks and datasets to evaluate various aspects of robotic manipulation, from precision and coordination to adaptability in dynamic environments.
Below, we describe each task, its objectives, and the associated datasets, highlighting their roles in our empirical evaluations.
\textbf{Robomimic Benchmark.}
Robomimic is a comprehensive benchmark for robotic manipulation with five tasks: {\bf Can}, {\bf Square}, {\bf Transport}, {\bf Tool Hang}, and {\bf Lift}.
These tasks are designed to evaluate a range of manipulation skills that include pick-and-place, precision assembly, multi-arm coordination, and basic object manipulation.
Each task includes expert demonstrations collected via proficient human (PH) teleoperation, with additional mixed human (MH) datasets, resulting in nine datasets total.
The datasets are structured in HDF5 format, containing observations, actions, rewards, and other metadata, and are available online (\url{https://github.com/ARISE-Initiative/robomimic}).
\begin{table}[h]
\centering
\caption{Summary of Robomimic Tasks and Datasets}
\label{tab:robomimic_summary}
\begin{tabular}{lccp{4cm}}
\toprule
\textbf{Task} & \textbf{PH Demos} & \textbf{MH Demos} & \textbf{Key Skills Tested} \\
\midrule
Can & 200 & 300 & Pick-and-place \\
Square & 200 & 300 & Precision manipulation \\
Transport & 200 (2 operators) & 300 & Multi-arm coordination \\
Tool Hang & 200 & 300 & Precision grasping, insertion \\
Lift & 200 & 300 & Basic object manipulation \\
\bottomrule
\end{tabular}
\end{table}
\textbf{Can.} The Can task involves picking up a can and placing it in the correct bin, testing pick-and-place skills. The dataset includes two variants:
\begin{itemize}[leftmargin=*]
\item \textbf{Proficient Human (PH)}: 200 demonstrations collected by a single proficient operator using the RoboTurk platform (\url{https://roboturk.stanford.edu/}).
\item \textbf{Mixed Human (MH)}: 300 demonstrations from six operators of varying proficiency, introducing variability in demonstration quality.
\end{itemize}
The PH dataset provides high-quality demonstrations, while MH introduces real-world variability.
\textbf{Square.} The Square task, also known as Square Nut Assembly, requires fitting a square nut onto a square peg (and potentially a round nut onto a round peg).
The scene includes two colored pegs (square and round) and two nuts on a tabletop, with randomized nut locations at episode start (\url{https://robosuite.ai/docs/modules/environments.html}).
This task evaluates precision manipulation.
The dataset includes two variants:
\begin{itemize}[leftmargin=*]
\item \textbf{Proficient Human (PH)}: 200 demonstrations collected by a proficient operator via RoboTurk.
\item \textbf{Mixed Human (MH)}: 300 demonstrations from operators of varying proficiency.
\end{itemize}
These datasets test algorithms on precise manipulation with both expert and varied human inputs.
\textbf{Transport.} The Transport task involves transporting an object using two robotic arms, which requires multi-arm coordination.
It is noted for its complexity, with observation spaces including shoulder and wrist views per arm, indicating a dual-arm setup (\url{https://robomimic.github.io/study/}).
The dataset includes two variants:
\begin{itemize}[leftmargin=*]
\item \textbf{Proficient Human (PH)}: 200 demonstrations collected by two proficient operators working together (\url{https://www.tensorflow.org/datasets/catalog/robomimic_ph}).
\item \textbf{Mixed Human (MH)}: 300 demonstrations from varied operators.
\end{itemize}
The datasets evaluate coordination and robustness in multi-arm cooperation tasks.
\textbf{Tool Hang.} The Tool Hang task requires hanging a tool (e.g., a hook), which involves precise grasping and insertion.
The dataset includes two variants:
\begin{itemize}[leftmargin=*]
\item \textbf{Proficient Human (PH)}: 200 demonstrations collected by a proficient operator via RoboTurk.
\item \textbf{Mixed Human (MH)}: 300 demonstrations from operators of varying proficiency.
\end{itemize}
These datasets evaluate on high-precision tasks.
\textbf{Lift.} The Lift task involves lifting an object, a simple manipulation task that benefits less from large datasets compared to complex tasks.
The dataset includes two variants:
\begin{itemize}[leftmargin=*]
\item \textbf{Proficient Human (PH)}: 200 demonstrations collected by a proficient operator via RoboTurk.
\item \textbf{Mixed Human (MH)}: 300 demonstrations from operators of varying proficiency.
\end{itemize}
The datasets provide a baseline for evaluating basic manipulation skills.
Table~\ref{tab:robomimic_summary} provides a summary of the Robomimic tasks and their corresponding datasets.
\textbf{Franka Kitchen Benchmark.}
Franka Kitchen is a simulation benchmark where a 9-DoF Franka arm operates in a kitchen environment, performing four household tasks per trajectory.
The tasks involve manipulating objects to achieve a desired goal configuration, which include:
\begin{itemize}[leftmargin=*]
\item \textbf{Open the Microwave}: The robot opens the microwave door.
\item \textbf{Move the Kettle}: The robot repositions the kettle to a target location.
\item \textbf{Flip the Light Switch}: The robot toggles a light switch to turn on a light.
\item \textbf{Slide Open the Cabinet Door}: The robot slides open a cabinet door.
\end{itemize}
The benchmark is hosted on platforms like Gymnasium-Robotics (\url{https://robotics.farama.org/envs/franka_kitchen/}) and uses datasets from D4RL (\url{https://github.com/Farama-Foundation/D4RL}).
\begin{table}[h]
\centering
\caption{Summary of Franka Kitchen Tasks and Datasets}
\label{tab:kitchen_dataset}
\begin{tabular}{lccp{3.5cm}}
\toprule
\textbf{Task} & \textbf{Description} & \textbf{Dataset} & \textbf{Key Skills Tested} \\
\midrule
Open Microwave & Open the microwave door & 566 trajectories & Object interaction, door manipulation \\
Move Kettle & Reposition the kettle & 566 trajectories & Object repositioning \\
Flip Light Switch & Toggle the light switch & 566 trajectories & Switch manipulation \\
Slide Cabinet Door & Slide open the cabinet door & 566 trajectories & Sliding mechanism interaction \\
\bottomrule
\end{tabular}
\end{table}
\textbf{Datasets.} The Franka Kitchen dataset consists of 566 trajectories, each completing all four tasks, collected via human teleoperation across seven interactive objects (a microwave, a kettle, an overhead light switch, a sliding cabinet, a hinged cabinet, a top burner, and a bottom burner) (\url{https://minari.farama.org/datasets/D4RL/kitchen/index.html}). The dataset evaluates the ability to sequence and execute multiple subtasks in a realistic kitchen environment.
Table~\ref{tab:kitchen_dataset} provides a summary of the Franka Kitchen tasks and their corresponding datasets.
\textbf{Push-T.}
Push-T is a manipulation task adapted from Implicit Broadcast Communication (IBC), where a circular end-effector pushes a T-shaped block to a target location.
The block and end-effector start at random positions, which can be defined by different evaluation seeds.
Observations are either RGB images ({\bf img}) or nine 2D keypoints ({\bf kp}) outlining the T block’s shape and the target location, as well as the end-effector’s position (\url{https://paperswithcode.com/task/robot-manipulation}).
{\bf Dataset}: The Push-T dataset includes 62,500 push interactions across 200 evaluation seeds, designed to evaluate the pushing dynamics and perception in a controlled setting.
\textbf{Dynamic Push-T.}
We extend the Push-T task by introducing a moving orange ball that travels at 100 units per second within a 500-unit square.
The ball bounces off both the walls and the end-effector, occasionally disrupting the T-shaped block.
The agent must simultaneously block these disturbances while pushing the block to the target, requiring it to maintain stable control under dynamic interference.
This setting explicitly evaluates latent continuity and adaptability: abrupt latent shifts can destabilize interactions, whereas smooth latent evolution enables the agent to adapt effectively to changing dynamics.
{\bf Dataset}: Similar to the Push-T dataset, we include 62,500 push interactions across 200 evaluation seeds, specifically designed to evaluate robustness and adaptability in dynamic and unpredictable environments.
Table~\ref{tab:pusht_dataset} provides a summary of the Push-T and Dynamic Push-T tasks and their corresponding datasets.
\begin{table}[h]
\centering
\caption{Summary of Push-T and Dynamic Push-T Tasks and Datasets}
\label{tab:pusht_dataset}
\begin{tabular}{lp{4cm}p{2cm}p{4cm}}
\toprule
\textbf{Task} & \textbf{Description} & \textbf{Dataset} & \textbf{Key Skills Tested} \\
\midrule
Push-T & Push T-shaped block to target & 62,500 pushes & Pushing dynamics and perception \\
Dynamic Push-T & Push T-shaped block with moving ball disturbance & 62,500 pushes & Adaptability and disturbance handling \\
\bottomrule
\end{tabular}
\end{table}
\section{Implementation Details}\label{sec:implementation}
Our DP-AG builds on the Diffusion Policy (DP)\citep{chi2023diffusion}, with two key modifications: (1) we add two linear heads for predicting the base drift $\mu_\phi(z_t)$ and log-variance $\log \sigma_\phi^2(z_t)$ from the static latent features, and (2) we incorporate VJP-guided SDE for latent evolution without introducing extra learnable parameters.
For observation encoding, we use a ResNet-18~\citep{he2016deep} backbone for image-based tasks.
We adopt the same conditional U-Net architecture from DP~\citep{chi2023diffusion} to predict diffusion noise.
During training, we apply random cropping with task-specific sizes as in DP, while a center crop is used at inference.
To maintain dynamic consistency, no color jitter or random flipping is applied.
An MLP is used to encode the agent’s proprioceptive inputs.
For key-point-based tasks, we follow the original DP setup and use fully connected networks.
We also apply action normalization scales each action dimension independently to $[-1, 1]$ to ensure compatibility with the DDPM denoising process, where predictions are clipped within this range at each step.
For positional control tasks, actions use 6D rotation representations.
Velocity control tasks use 3D axis-angle representations, which follow standard practice.
We train the model using the iDDPM algorithm~\citep{nichol2021improved} with 100 diffusion steps. All models are trained for 300 epochs on vision-based tasks and 200 epochs for key-point-based tasks.
For learning rate scheduling, we use a cosine annealing schedule with a linear warmup of 500 steps.
Batch sizes are set to 64 for image-based tasks and 256 for key-point-based tasks.
Optimization uses AdamW with a learning rate of $1\times10^{-4}$ in all experiments.
All hyperparameters not directly related to DP-AG extensions (e.g., diffusion step count, augmentation, action normalization) are kept identical to DP for a controlled and fair comparison.
For inference, we maintain the same number of diffusion denoising steps as training to avoid introducing a distributional shift.
\section{Ablation Studies}\label{sec:ab}
\subsection{Latency and Computational Cost Analysis}\label{sec:ablation_computation}
We evaluate the computational overhead introduced by our DP-AG compared to the original DP by measuring both training time per epoch, the inference latency, and the total training time to convergence.
The computation analysis is conducted on the Push-T benchmark.
Experiments are conducted on an Nvidia RTX 4090 GPU with 24GB VRAM.
Table~\ref{tab:latency} summarizes the detailed computational cost comparison between the DP baseline and our DP-AG.
\textbf{Training Time per Epoch.}
On the Push-T benchmark, the average training time per epoch increases slightly from 114.2 seconds for DP to 119.5 seconds for DP-AG, corresponding to an overhead of 4.6\%.
This increase is expected because VJP computation requires additional backward-mode automatic differentiation through the noise predictor.
However, the extra cost remains moderate, largely due to the relatively small dimensionality of the latent observation space.
\textbf{Inference Latency.}
At inference time, the VJP computations are omitted entirely.
Our DP-AG simply operates with the latent features extracted from the observation encoder at inference.
As a result, inference latency remains virtually unchanged: 145.3 milliseconds per action sequence generation for DP versus 146.5 milliseconds for our DP-AG, a marginal 0.8\% difference.
This negligible overhead ensures that DP-AG maintains real-time responsiveness for high-frequency robotic control while improving the smoothness and consistency of action generation.
\textbf{Training Efficiency and Total Time to Convergence.}
Although our DP-AG introduces a minor increase in per-epoch training time, it significantly accelerates convergence.
On the Push-T benchmark, DP requires approximately 200 epochs to converge, whereas DP-AG achieves comparable performance within only around 100 epochs.
This effectively reduces the number of required training epochs and the total training time by nearly 50\%.
To be specific, DP completes training in about 6.2 hours, while DP-AG completes training in approximately 3.5 hours, which results in a net saving of 2.7 hours.
Thus, despite the slight per-epoch overhead, our DP-AG achieves faster overall training and improved sample efficiency.
\textbf{Summary.}
Therefore, while DP-AG introduces additional VJP computations, modern autodiff frameworks ({\em e.g.}, PyTorch) execute them efficiently, leading to negligible runtime overhead. As shown in Table~\ref{tab:latency}, DP-AG incurs only a minor 4.6\% increase in per-epoch training time but delivers nearly 50\% higher training efficiency.
Inference latency remains virtually unchanged: DP-AG sustains real-time control on the UR5 robot while producing smoother trajectories, lower jerk, and faster task completion.
These results demonstrate that the added computation does not hinder deployment or responsiveness, and that our perception–action interplay both improves DP and substantially accelerates training, making DP-AG well suited for real-world applications where rapid retraining and real-time control are important.
\begin{table}[h]
\centering
\small
\setlength{\tabcolsep}{3pt}
\caption{Computational cost comparison between DP and DP-AG on Nvidia RTX 4090 (Push-T).}
\begin{tabular}{lcccc}
\toprule
Model & One Epoch Time (s) & Epochs to Converge & Total Training Time (h) & Inference Latency (ms) \\
\midrule
DP & 114.2 & $\sim$200 & $\sim$6.2 & 145.3 \\
\rowcolor{gray!20}
DP-AG (Ours) & 119.5 & $\sim$100 & $\sim$3.5 & 146.5 \\
\hline
\end{tabular}
\label{tab:latency}
\end{table}
\subsection{Effect of Cycle-Consistent Contrastive Loss}
\begin{figure}
\centering
\includegraphics[width=0.8\textwidth]{imgs/withwithout_contrastive.png}
\caption{Comparison of training convergence on the Push-T benchmark for DP-AG with and without cycle-consistent contrastive loss.}
\label{fig:withwithout_contrastive}
\end{figure}
To evaluate the importance of the cycle-consistent contrastive loss in our DP-AG, we conduct an ablation study by removing this component while keeping the rest of the architecture unchanged.
This allows us to isolate the role of cycle consistency in enforcing mutual smoothness between latent evolution and action refinement during training.
\textbf{Setup.}
We compare two variants on the Push-T benchmark:
\begin{itemize}[leftmargin=*]
\item DP-AG (full model): Includes the cycle-consistent contrastive loss between static and VJP-guided noise predictions.
\item DP-AG w/o Contrastive: Removes the contrastive term from the training objective, relying only on the diffusion loss and KL regularization.
\end{itemize}
Both variants are trained under the same settings; and the action MSE is evaluated across training epochs.
We repeat the experiments with 5 different initialization seeds for training.
\textbf{Results.}
As shown in Figure~\ref{fig:withwithout_contrastive} and Table~\ref{tab:contrastive}, removing the cycle-consistent loss leads to slower convergence, higher final training action MSE, and reduced target coverage scores.
To be specific, the model without contrastive loss requires approximately 3 times more epochs to converge compared to the full DP-AG; the final success rate drops from 93\% to 85\% on the Push-T benchmark; and the training action MSE remains consistently higher throughout training, which indicates less accurate action generation.
\textbf{Analysis.}
Without the contrastive alignment between evolving and static latents, VJP-guided perturbations can drift away from the optimal action refinement trajectory, which degrades both latent continuity and action smoothness.
The cycle-consistent loss is important in closing the perception-action loop, which can ensure that latent evolution remains tightly coupled with action denoising across diffusion steps.
\begin{table}[h]
\centering
\small
\setlength{\tabcolsep}{8pt}
\caption{Effect of cycle-consistent contrastive loss on Push-T benchmark.}
\begin{tabular}{lccc}
\toprule
Model & Epochs to Converge & Converged Train Action MSE & Score \\
\midrule
DP-AG (full) & 100 & 65.8 & 0.93 \\
DP-AG w/o Contrastive & 300 & 183.5 & 0.85 \\
\hline
\end{tabular}
\label{tab:contrastive}
\end{table}
\subsection{Likelihood Supervision vs. Contrastive Loss}\label{sec:ablation_likelihood}
\begin{figure}
\centering
\includegraphics[width=0.9\textwidth]{imgs/contrastive_v.s_mse.png}
\caption{Validation accuracy during training with contrastive loss vs. MSE noise regression on the Push-T benchmark.
Contrastive loss achieves faster convergence and better performance but with slightly more variance during training.}
\label{fig:contrastive_mse}
\vspace{-2mm}
\end{figure}
In Section~\ref{sec:elbo}, we derive a noise regression objective from the variational lower bound, which serves a similar role to our cycle-consistent contrastive loss.
To evaluate the necessity of the contrastive loss, we replace it with the likelihood-based noise regression objective defined in Equation~\ref{eqn:lh}:
\begin{equation}
\mathcal{L}_{\text{LH}} = \mathbb{E}_{(o_t, a_t) \sim \mathcal{D},\; k \sim \mathcal{U}(1, K)} \left[ \| \epsilon_\theta(\hat{a}_t^{k}, \tilde{z}_t^{k}, k) - \epsilon \|^2_2 \right].
\end{equation}
\textbf{Experimental Setup.}
We compare three variants:
\begin{itemize}[leftmargin=*]
\item \textbf{Cycle-Consisteny Contrastive}: Uses the proposed cycle-consistent contrastive loss $\mathcal{L}_{\text{cont}}$ only.
\item \textbf{MSE-Only}: Replaces $\mathcal{L}_{\text{cont}}$ with the noise regression loss $\mathcal{L}_{\text{LH}}$.
\item \textbf{Combined}: Combines both objectives: $\mathcal{L} = \mathcal{L}_{\text{cont}} + \mathcal{L}_{\text{LH}}$.
\end{itemize}
All models are trained with 5 random seeds on the Push-T benchmark. We report the mean and standard deviation of validation accuracy over 100 epochs.
\textbf{Results.}
As shown in Figure~\ref{fig:contrastive_mse}, the contrastive-only variant not only converges faster but also achieves the best validation accuracy.
Although MSE yields smoother convergence, it lacks the structural benefits of noise alignment that contrastive learning offers.
The combined objective does not improve performance over the contrastive-only variant and sometimes results in unstable training, likely due to conflicting optimization signals between absolute and relative supervision.
These results validate that the contrastive alignment provides stronger inductive bias for learning dynamic consistency between static and VJP-guided latents, while the likelihood term, though informative, introduces redundancy. We thus omit the likelihood objective from our final model.
\subsection{Effect of VJP Strength}~\label{sec:ablation_vjp}
\vspace{-4mm}
In this section, we study how the strength of the VJP-guided perturbations affects the performance of our DP-AG.
Recall that the VJP acts as a stochastic ``force'' that shapes the evolution of latent observation features based on the diffusion process that refines the action generation.
While moderate VJP guidance helps structure latent trajectories coherently, overly weak or strong guidance may destabilize training or restrict flexibility.
\textbf{Setup.}
We introduce a scaling factor $\gamma$ applied to the VJP term during latent updates in Equation~\ref{eqn:latent_update}:
\begin{equation}
dz_t = \gamma \cdot \text{VJP}(\hat{a}_t^k, z_t)\,dt + \sigma_\phi(z_t)\,dW_t,
\end{equation}
We evaluate the following settings:
$\gamma \in \{0.0,\ 0.5,\ 1.0,\ 2.0,\ 5.0\}$, where $\gamma=1.0$ is our default setting.
We repeat the experiments with 5 different initialization seeds for training.
\begin{figure}
\centering
\includegraphics[width=0.75\textwidth]{imgs/vjp_gamma_score.png}
\caption{Effect of the VJP strength $\gamma$ on the Push-T benchmark.}
\label{fig:vjp_gamma}
\end{figure}
\textbf{Results.}
The results are presented in Figure~\ref{fig:vjp_gamma} and Table~\ref{tab:vjp_ablation}.
We observe that:
\begin{itemize}[leftmargin=*]
\item Without VJP guidance, latent features remain static during diffusion, effectively reducing our DP-AG to a standard DP.
\item Moderate VJP strength ($\gamma=0.5, 1.0$) achieves the best results, which tradeoff a balance between encouraging latent evolution and maintaining stability.
\item High VJP strength ($\gamma=2.0, 5.0$) leads to unstable latent evolution, where the updated trajectories deviate excessively from the static latent features, ultimately degrading policy performance.
\end{itemize}
\begin{table}[h]
\centering
\small
\setlength{\tabcolsep}{8pt}
\caption{Effect of VJP strength $\gamma$ on Push-T benchmark.}
\begin{tabular}{lccc}
\toprule
$\gamma$ & Final Train Action MSE & Score & Latent Behavior \\
\midrule
0.0 & 75.3 & 0.87 & static \\
0.5 & 58.3 & 0.93 & smooth \\
1.0 & 65.8 & 0.93 & smooth \\
2.0 & 71.2 & 0.90 & over-reactive \\
5.0 & 241.5 & 0.85 & unstable \\
\hline
\end{tabular}
\label{tab:vjp_ablation}
\end{table}
\subsection{Effect of KL Loss Hyperparameter}\label{sec:ablation_kl}
\begin{figure}
\centering
\includegraphics[width=0.75\textwidth]{imgs/kl_ablation_score_uncertainty.png}
\caption{
Effect of KL loss weighting $\lambda_{\text{KL}}$ on the target coverage score for the Push-T benchmark.}
\label{fig:kl_weight}
\end{figure}
In this section, we investigate the effect of varying the KL divergence coefficient $\lambda_{\text{KL}}$ on the performance of our DP-AG using the Push-T benchmark.
While $\beta$-VAE~\citep{higgins2017betavae} tunes KL regularization to promote disentanglement, in our case, the KL term stabilizes the action-guided latent evolution without affecting its ability to adaptively track action refinements.
\textbf{Setup.}
We train DP-AG on the Push-T benchmark with the following $\lambda_{\text{KL}}$ values:
\{0,\ 0.5,\ 1.0,\ 2.0,\ 4.0,\ 5.0,\ 10.0\}.
All other training configurations remain unchanged across experiments.
We repeat the experiments with 5 different initialization seeds for training.
\textbf{Results.}
The results are presented in Figure~\ref{fig:kl_weight} and summarized in Table~\ref{tab:kl_ablation}.
We observe the following trends:
\begin{itemize}[leftmargin=*]
\item Without KL regularization ($\lambda_{\text{KL}} = 0$), latent evolution becomes less stable, causing updates to drift excessively from the static latent and resulting in lower target coverage scores.
\item Small KL values ($\lambda_{\text{KL}}=0.5, 1.0$) achieve the best trade-off between flexibility and stability, which leads to the best performance.
\item Moderate to large KL values ($\lambda_{\text{KL}}=2.0, 4.0, 5.0$) begin to over-regularize the latent space, which limits its ability to adapt to action refinements.
\item Strong KL regularization ($\lambda_{\text{KL}}=10.0$) severely restricts latent evolution, which causes underfitting and significantly lower target coverage score.
\end{itemize}
\textbf{Analysis.}
These results suggest that the KL regularization in our DP-AG should not be viewed through the lens of promoting disentanglement, as in $\beta$-VAE.
Instead, it acts to anchor the VJP-guided latent dynamics toward stable yet adaptive evolution.
\begin{table}
\centering
\small
\setlength{\tabcolsep}{8pt}
\caption{Effect of KL loss hyperparameter $\lambda_{\text{KL}}$ on Push-T benchmark.}
\begin{tabular}{lccc}
\toprule
$\lambda_{\text{KL}}$ & Converged Train Action MSE & Score & Latent Stability \\
\midrule
0 & 123.1 & 0.89 & slightly unstable \\
0.5 & 86.6 & 0.91 & stable \\
1.0 & 65.8 & 0.93 & stable \\
2.0 & 91.3 & 0.91 & slightly over-constrained \\
4.0 & 78.3 & 0.92 & over-constrained \\
5.0 & 168.8 & 0.88 & over-constrained \\
10.0 & 324.5 & 0.82 & severely over-constrained \\
\bottomrule
\end{tabular}
\label{tab:kl_ablation}
\end{table}
\subsection{Effect of Contrastive Loss Temperature Parameter}\label{sec:ablation_cont}
\begin{figure}
\centering
\includegraphics[width=0.75\textwidth]{imgs/tau_ablation_score_uncertainty.png}
\caption{
Effect of temperature $\tau$ in the cycle-consistent contrastive loss on Push-T benchmark performance.}
\label{fig:temperature}
\end{figure}
In this section, we investigate the influence of the temperature parameter $\tau$ in the cycle-consistent contrastive loss on the performance of our DP-AG.
As in self-supervised learning settings~\citep{chen2020simple}, the temperature parameter controls the sharpness of similarity scores between a query and its positive pairs in the contrastive learning.
A low $\tau$ enforces highly sharp alignment, while a high $\tau$ leads to smoother matching across samples.
In our DP-AG, unlike its usage for self-supervised learning, our contrastive loss is designed to enforce cycle consistency between perception and action noise predictions during latent evolution.
To evaluate the impact of the temperature parameter on this novel perception-action interplay, we vary $\tau$ across a broad range: $\tau \in \{0.01,\ 0.1,\ 0.2,\ 0.5,\ 0.8,\ 1.0,\ 1.5\}$.
We repeat the experiments with 5 different initialization seeds for training.
\textbf{Results.}
Figure~\ref{fig:temperature} presents the results, with a summary provided in Table~\ref{tab:tau_ablation}.
We observe:
\begin{itemize}[leftmargin=*]
\item Very low temperature ($\tau=0.01,\ 0.1$) could overly sharp alignment that causes unstable training and brittle latent evolution.
\item Moderate temperatures ($\tau=0.2,\ 0.5$) improve stability compared to $\tau=0.01$, but still do not fully leverage latent adaptability.
\item High temperature ($\tau=0.8$) achieves the best performance, which yields the highest success rates and smoothest latent evolution.
\item Very high temperature ($\tau=1.5$) could weaken cycle consistency, which leads to performance degradation.
\end{itemize}
\begin{table}[h]
\centering
\small
\setlength{\tabcolsep}{8pt}
\caption{Effect of temperature $\tau$ on cycle-consistent contrastive loss (Push-T benchmark).}
\begin{tabular}{lccc}
\toprule
$\tau$ & Converged Train Action MSE & Score & Stability \\
\midrule
0.01 & 433.8 & 0.82 & brittle \\
0.1 & 243.4 & 0.85 & unstable \\
0.2 & 186.9 & 0.90 & improved \\
0.5 & 103.1 & 0.92 & stable \\
0.8 & 65.8 & 0.93 & best stability \\
1.0 & 96.3 & 0.92 & slightly diffuse \\
1.5 & 105.5 & 0.89 & diffuse \\
\hline
\end{tabular}
\label{tab:tau_ablation}
\end{table}
\textbf{Analysis.}
In our DP-AG, a higher temperature (e.g., $\tau=0.8$) provides better flexibility for aligning noise predictions during latent evolution, allowing dynamic perception refinement while maintaining cycle consistency.
Extremely low $\tau$ values over-constrain the model, while excessively high values ($\tau > 1.0$) can weaken the noise prediction alignment.
\subsection{Comparison with Input Perturbation Smoothness}
To evaluate whether the stability achieved by DP-AG could be reproduced by alternative smoothness regularization, we compared against baselines that enforce consistency under input perturbations.
For each observation, we generated a perturbed version by adding Gaussian noise, then minimized MSE or cosine similarity between their predicted action noise scores.
Table~\ref{tab:perturbation_smoothness} reports results on the Push-T benchmark.
Action smoothness is quantified as normalized inverse jerk (higher is smoother).
\textbf{Results.}
DP-AG clearly outperforms perturbation-based baselines, achieving higher success rates and smoother action sequences.
This demonstrates that contrastive regularization not only enforces smoothness but also preserves semantic alignment, leading to superior policy performance.
\begin{table}[h!]
\centering
\caption{Comparison of DP-AG with input perturbation smoothness baselines on Push-T.}
\label{tab:perturbation_smoothness}
\vspace{2mm}
\begin{tabular}{lcccc}
\toprule
Method & SR (img) & SR (kp) & Smoothness (img $\uparrow$) & Smoothness (kp $\uparrow$) \\
\midrule
Perturbation MSE & 0.85 & 0.92 & 0.83 & 0.87 \\
Perturbation Cosine & 0.88 & 0.95 & 0.86 & 0.92 \\
{\bf DP-AG (Contrastive Loss)} & \textbf{0.93} & \textbf{0.99} & \textbf{0.91} & \textbf{0.95} \\
\bottomrule
\end{tabular}
\end{table}
\section{More Comparisons of Training Convergence}\label{sec:more_convergence}
\begin{figure}
\centering
\includegraphics[width=0.9\textwidth]{imgs/square-action-mse.png}
\caption{Training action MSE over epochs on Robomimic Square.}
\label{fig:squre-convergence}
\end{figure}
\begin{figure}
\centering
\includegraphics[width=0.9\textwidth]{imgs/lift-action-mse.png}
\caption{Training action MSE over epochs on Robomimic Lift.}
\label{fig:lift-convergece}
\end{figure}
In this section, we compare the training convergence behavior between the baseline DP and our DP-AG in Robomimic Lift and Square tasks.
Training convergence is evaluated by measuring the MSE between predicted and ground-truth actions over training epochs.
\textbf{Setup.}
Both the baseline DP and our DP-AG models are trained under the same settings, which use the same dataset splits, optimizer configurations, and diffusion schedules, etc.
Training is conducted on an Nvidia RTX 4090 GPU with a batch size of 64, and convergence is measured by tracking the action MSE over 300 epochs for each method.
\textbf{Results.}
The convergence curve for the Robomimic Square task is shown in Figure~\ref{fig:squre-convergence}, and the curve for the Robomimic Lift task is shown in Figure~\ref{fig:lift-convergece}.
These results demonstrate that the action-guided perception updates in our DP-AG not only improve the policy performance but also enhance the training dynamics.
By dynamically refining latent features through VJP-guided evolution, DP-AG accelerates the denoising process, which leads to more sample-efficient and stable training compared to the static latent used in the baseline DP.
\section{Real-World Manipulation Experiments on UR5 Robotic Arm}\label{sec:real-world}
We conduct real-world manipulation experiments using the Universal UR5 robotic arm (\url{https://www.universal-robots.com/products/ur5e/}), a 6-degree-of-freedom platform widely adopted in both industrial and research contexts for its versatility and reliability.
In this section, we compare the performance of our proposed DP-AG against the baseline Diffusion Policy (DP) across three visuomotor tasks: Painting, Candy Push, and Peg-in-Hole Insertion.
These tasks are specifically designed to evaluate the perception–action interplay, testing each policy’s ability to generalize from limited sensory inputs (e.g., RGB images) to precise motor control commands.
All evaluations are conducted in a controlled real-world setting, allowing us to evaluate each method’s effectiveness, robustness, and potential for the real-world deployment in practical robotic applications.
\textbf{Dataset Collection.}
We prepared task-specific datasets for Painting, Candy Push, and Peg-in-Hole Insertion to support our real-world experiments using the UR5 robotic arm.
Each dataset was collected via kinesthetic teaching using a teach pendant.
Human operators manually guided the robot through the desired trajectories, while synchronized visual and proprioceptive data were recorded in real time.
RGB images were captured at 30 frames per second (FPS) using two calibrated cameras: one fixed in front of the workspace and another mounted on the robot wrist to provide egocentric views.
Simultaneously, robot joint states and Cartesian end-effector poses were recorded at 100 Hz via the UR5’s internal API.
To ensure alignment between visual and action modalities, timestamps from both the camera streams and robot sensors were synchronized using a shared system clock.
Each trajectory thus consisted of time-aligned sequences of images and robot poses.
For the Painting task, we collected 100 trajectories by having the operator trace either heart-shaped or circular paths on a horizontal plane. Ground-truth 2D target contours were predefined, and the demonstrations were guided to closely follow these shapes.
The Candy Push task involved 80 trials where small candy objects were randomly placed within a $20\text{cm}\times20\text{cm}$ area.
The human operator manually guided the robot to push each candy toward a designated target area marked on the table surface.
Success or failure labels were assigned based on whether all candies reached the target area within a tolerance threshold.
In the Peg-in-Hole Insertion task, 50 demonstrations were gathered by inserting pegs into variable hole positions, with occlusion and visual ambiguity intentionally introduced.
Success was defined as full insertion without slippage or bounce-back.
All demonstrations were segmented and temporally synchronized with the image streams using trajectory timestamps.
Each dataset was then split into 80\% of the trajectories for training and 20\% for evaluation.
\textbf{Training and Implmentation Details.}
The training procedure and implementation closely follow those used in the simulation benchmark experiments.
We build our training pipeline on top of the DP framework, incorporating two core changes in our DP-AG: (1) two separate linear projection heads are added to map the static latent features $z_t$ to the drift term $\mu_\phi(z_t)$ and log-variance $\log \sigma_\phi^2(z_t)$, and (2) latent evolution is guided using a Vector–Jacobian Product (VJP)-driven stochastic differential equation (SDE), implemented without introducing new trainable parameters.
We extract visual features from RGB frames using a pre-trained ResNet-152 encoder.
Proprioceptive inputs, including joint angles and velocities, are processed through a two-layer MLP. For diffusion noise prediction, we use the same conditional U-Net architecture as in the simulation experiments, where white noise serves as input and observation features act as conditioning signals.
During training, we apply random cropping with sizes tuned per task; at evaluation time, we apply center cropping for consistency.
No color jitter or random horizontal flipping is used, as such augmentations may disrupt the temporal consistency critical for dynamics modeling.
Actions are normalized to the $[-1, 1]$ range per dimension to align with the denoising diffusion process, and outputs are clipped to stay within bounds.
Peg-in-Hole tasks use 6D continuous rotation representations to avoid gimbal lock, while Candy Push and Painting use 3D axis-angle vectors for velocity control.
Training follows the iDDPM objective with 100 diffusion steps.
Each policy is trained for 200 epochs with a batch size of 64, using AdamW optimization and a cosine learning rate schedule with a 500-step warmup and a base learning rate of $1 \times 10^{-4}$.
\textbf{Evaluation Metrics.}
We evaluated each policy using 20\% held-out trajectories per task, with each trial repeated 20 times.
Four task-specific metrics were used to evaluate performance: success rate, smoothness (average jerk), IoU (for Painting), and time to complete.
\begin{itemize}[leftmargin=*]
\item {\bf Success Rate (\%)}: Measures how often the robot completes the task correctly among all trials. In Candy Push, a trial is successful if all candies are pushed into the target zone.
For Peg-in-Hole, success means the peg is fully inserted.
We report mean and standard deviation over 20 trials to demonstrate consistency under variations like random object positions or occlusions.
\item {\bf Smoothness (Average Jerk)}: Captures the quality of motion based on jerk (the rate of change of acceleration). It is derived from the end-effector’s position data (recorded at 100 Hz) using finite difference approximations.
For a position sequence $\{x_1, x_2, \dots, x_n\}$ at uniform time intervals $\Delta t$, average jerk is computed as:
Velocity (first derivative):
\begin{equation}
v_i = \frac{x_{i+1} - x_i}{\Delta t}.
\end{equation}
Acceleration (second derivative):
\begin{equation}
a_i = \frac{v_{i+1} - v_i}{\Delta t} = \frac{x_{i+2} - 2x_{i+1} + x_i}{\Delta t^2}.
\end{equation}
Jerk (third derivative):
\begin{equation}
j_i = \frac{a_{i+1} - a_i}{\Delta t} = \frac{x_{i+3} - 3x_{i+2} + 3x_{i+1} - x_i}{\Delta t^3}.
\end{equation}
Average jerk across the trajectory:
\begin{equation}
\text{Average Jerk} = \frac{1}{n} \sum_{i=1}^{n} \left\| j_i \right\|.
\end{equation}
Lower jerk values indicate smoother movements, which are especially important in tasks like Painting and insertion.
\item {\bf Intersection over Union (IoU in \%)}: Used only in Painting, IoU evaluates how well the painted shape matches the ground-truth contour (i.e., heart, circle), which is calculated as:
\begin{equation}
\text{IoU} = \frac{\text{Intersection Area}}{\text{Union Area}}.
\end{equation}
A higher IoU means better adherence to the intended shape.
\item {\bf Time to Complete (s)}: Records how long it takes to finish the task, which evaluates the efficiency alongside accuracy and smoothness. Faster times, particularly in Candy Push, suggest better adaptation to object dynamics.
\end{itemize}
These tasks evaluate perception-action interplay and generalization from partially observable inputs. Painting demands smooth motion, where jerk and IoU are important.
Candy Push tests adaptation to varying layouts, with success rate and time as the key indicators.
Peg-in-Hole requires 3D reasoning from 2D data, where success rate and jerk highlight stability.
The dataset split ensures robust training and evaluation, consistent with the standard practice.
In future work, we plan to explore sim-to-real transfer for our DP-AG, enabling real-world deployment with minimal or no additional data collection or model retraining.
%%%%%%%%%%% To incorporate the rebuttal %%%%%%%%%%%%%
% \section{Visualization of VJP-Guided Latent Evolution in Push-T}
% \label{app:tsne}
% To further illustrate how the proposed VJP mechanism shapes observation latents, we analyze the Push-T benchmark, which presents a challenging perception-action coupling task: the policy must infer block pose, contact geometry, and motion direction directly from high-dimensional visual input.
% \textbf{Setup.}
% In each evaluation episode, the observation horizon is encoded with the vision encoder to produce static latents, as in diffusion policy (DP).
% During action generation, we then evolve these latents using our proposed VJP updates within the diffusion process (DP-AG).
% For analysis, we snapshot the latent state at five denoising steps.
% We then embed the static latents and their VJP-evolved counterparts into a shared 2-D space using t-SNE (Figure~\ref{fig:tsne_vjp}), allowing us to track the latent drift induced by the VJP.
% \begin{figure}[h]
% \centering
% \includegraphics[width=1.0\linewidth]{imgs/tsne.pdf}
% \caption{{\bf t-SNE visualization of latent evolution in Push-T}.
% Each static latent (\textcolor{cyan}{blue}) is accompanied by its VJP trajectory across five steps (\textcolor{orange}{orange} $\rightarrow$ \textcolor{green}{green} $\rightarrow$ \textcolor{red}{red} $\rightarrow$ \textcolor{purple}{purple} $\rightarrow$ \textcolor{brown}{brown}).
% The structured drift highlights how VJP updates reshape observation latents using action feedbacks.}
% \label{fig:tsne_vjp}
% \end{figure}
% \textbf{Results.}
% As shown in Figure~\ref{fig:tsne_vjp}, the VJP-induced latent drift in a real control setting is both interpretable and beneficial, yielding a structured trajectory that aligns perception with action demands.
% To be specific, several consistent patterns are observed across multiple Push-T rollouts:
% \begin{itemize}[leftmargin=*]
% \item \textbf{Structured drift.} Each static latent (blue) is accompanied by a smooth trajectory of points (\textcolor{orange}{orange} $\rightarrow$ \textcolor{green}{green} $\rightarrow$ \textcolor{red}{red} $\rightarrow$ \textcolor{purple}{purple} $\rightarrow$ \textcolor{brown}{brown}), corresponding to the five VJP steps.
% The trajectories are short and consistent, demonstrating that the latent evolution is systematic rather than random.
% \item \textbf{Task-relevant convergence.} Across samples, the VJP updates move latents in similar directions, suggesting that the evolution consistently emphasizes task-relevant cues ({\em e.g.,} object pose and contact dynamics) necessary for accurate action prediction in Push-T.
% \item \textbf{Bridging perception and action.} While the static latents encode raw visual content, the evolved latents progressively reshape this representation towards what the policy needs for predicting the push trajectory.
% This explains the improved action success rate observed when comparing DP-AG against the static baseline (DP).
% \end{itemize}
\section{Visualizations of Action-Guided Latent Evolution}
\label{app:viz_vjp_pusht}
To clarify the semantic meaning of the latent drift in our DP-AG and why it matters for action prediction, we visualize how observation latents evolve under action guidance on the Push-T benchmark.
The key question is:
\emph{“If the end-effector must move toward its desired future state, how should the latent representation of the current scene be adjusted?”}
In DP-AG, the answer is given by the VJP, which nudges the latent toward features that are most predictive of the next action.
This latent drift is not arbitrary.
It is shaped by the policy’s uncertainty: when the action head is confident, updates are small; when the action is ambiguous, the latent is pulled toward features that disambiguate the correct motion.
Thus, drift is expected to improve prediction by refining perception precisely where it is action-relevant.
To make these refinements interpretable, we decode the evolving latents with a lightweight VAE.
This choice directly addresses the concern of whether latent drift has semantic meaning: decoded frames expose what the policy is implicitly “re-seeing” as actions unfold.
If latent evolution were meaningless noise, decoded frames would be incoherent; instead, we observe structured changes that align with the intended push trajectory.
\begin{figure}
\centering
\includegraphics[width=\linewidth]{imgs/vjp_recon_only.pdf}
\caption{\textbf{Decoded VJP-evolved observation latents.}
Columns are independent episodes.
The top row ($k{=}0$) decodes the static latent.
Rows $k{=}1,{\ldots},5$ decode the latent after each VJP step.}
\label{fig:vjp_recon_only}
\end{figure}
\begin{figure}
\centering
\includegraphics[width=\linewidth]{imgs/vjp_recon_gradcam.pdf}
\caption{\textbf{Policy evidence on decoded frames.}
Grad-CAM from the end-effector $(x,y)$ head of the policy is overlaid on the same reconstructions as Figure~\ref{fig:vjp_recon_only}.}
\label{fig:vjp_recon_gradcam}
\end{figure}
\subsection{Decoded Reconstructions and Policy Attention.}
\textbf{Implementation.}
Let $z_0$ be the static latent from the encoder.
At each VJP step $k$, we update $z_k$ with an action-aligned perturbation normalized to fixed step size.
A lightweight VAE, trained only on static latents and then frozen, decodes each $z_k$ to an image $\tilde{x}_k$.
Gradients are never propagated back through this decoder.
To reveal the policy’s evidence, we also compute Grad-CAM on the $(x,y)$ end-effector head and overlay it on the decoded frames.
\textbf{Visualization.}
Figure~\ref{fig:vjp_recon_only} shows that the decoded VJP-evolved latents yield localized and coherent changes: the end-effector (blue gripper) shifts step by step toward its actual future position, while the background remains stable.
Grad-CAM overlays (Figure \ref{fig:vjp_recon_gradcam}) confirm that these same regions receive the strongest policy attention.
Notably, the static latent ($k=0$) contains no future cue, whereas the evolved latents ($k \geq 1$) highlight the precise spatial displacements the policy will execute.
This demonstrates that the action-guided latent drift has both semantic meaning (visible in decoded images) and functional value (aligned with the policy’s attention).
Together, the two figures demonstrate that VJP-induced latent drift produces interpretable image-space refinements and that these refinements are spatially aligned with the policy’s own predictions.
This coupling connects perceptual updates directly to the end-effector objective, demonstrating how action-guided latent evolution sharpens the visual grounding of control.
\begin{figure}
\centering
\includegraphics[width=\linewidth]{imgs/grad_cam.pdf}
\caption{\textbf{Grad-CAM (control view).}
Spatial evidence of the policy without reconstructions.}
\label{fig:grad_cam_only}
\end{figure}
\subsection{Attention-Only Policy View.}
As a control, we project Grad-CAM directly onto the raw input frames without decoding (Figure \ref{fig:grad_cam_only}).
These heatmaps still trace the end-effector’s trajectory, confirming that the decoder does not hallucinate the effect.
Instead, the VAE reconstructions provide a semantic lens: they make explicit how the evolving latent {\em imagines} the scene differently after each update.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{imgs/tsne.pdf}
\caption{\textbf{t-SNE of static vs.\ evolved latents.}
Each static latent (\textcolor{cyan}{blue}) is paired with its VJP trajectory over five steps (\textcolor{orange}{orange} $\rightarrow$ \textcolor{green}{green} $\rightarrow$ \textcolor{red}{red} $\rightarrow$ \textcolor{purple}{purple} $\rightarrow$ \textcolor{brown}{brown}).}
\label{fig:tsne_vjp}
\end{figure}
\subsection{t-SNE of Action-Guided Latent Evolution.}
While decoded frames and Grad-CAM demonstrate where VJP modifies the observation in pixel space, they do not reveal how these updates are organized in latent space.
t-SNE embeddings provide a complementary view by projecting both static and VJP-evolved latents into two dimensions, allowing us to inspect whether action-guided updates follow structured and task-aligned trajectories rather than random drift.
\textbf{Visualization.}
The t-SNE embeddings in Figure~\ref{fig:tsne_vjp} show that VJP does not cause latents to scatter randomly.
Instead, each static latent ($k{=}0$, \textcolor{cyan}{blue}) serves as an anchor from which the evolved latents ($k{=}1,\ldots,5$) trace short, smooth, and consistently oriented trajectories.
Across episodes, these trajectories align to form clusters that correspond to motion modes, such as push direction, indicating that VJP structures the latent space around action-relevant dynamics rather than noise.
Importantly, the cycle-consistent contrastive loss further reinforces this structure: it pulls each evolved latent back toward its static anchor while pushing it away from unrelated samples, preserving semantic grounding and preventing arbitrary drift.
As a result, the trajectories remain compact and discriminative, ensuring that latent evolution not only tracks action dynamics but also maintains clear and task-relevant separation.
In summary, across decoded frames, attention overlays, and t-SNE latent embeddings, a coherent picture emerges:
\begin{enumerate}[leftmargin=*]
\item VJP-driven drift produces semantically interpretable refinements (decoded end-effector positional shifts).
\item These refinements highlight exactly the regions that the policy deems most informative for predicting the next action.
\item In latent space, drift follows smooth and task-aligned trajectories rather than arbitrary noise.
\end{enumerate}
Together, these results demonstrate that action-guided latent evolution is not only mathematically principled but also semantically and functionally meaningful.
\section{Extension to Online Mode}
Although our main experiments are conducted in offline imitation learning, DP-AG is inherently well-suited for online training.
In this section, we present a streaming variant of DP-AG that remains fully consistent with the original formulation.
Its key mechanism, action-guided latent evolution via VJP, naturally supports the incremental refinement of perceptual features as new action-observation pairs arrive.
Below, we describe the specific modifications that enable DP-AG to function effectively in the online regime:
\textbf{Online latent evolution.}
At each diffusion step $k \in \{1,\dots,K\}$ between $t$ and $t{+}1$, the latent observation is updated using the same VJP-guided rule as in the main model:
\begin{equation}\label{eq:online_latent_update}
\tilde{z}_t^{k} \;=\; \mu_\phi(z_t) \;+\; \gamma \, \sigma_\phi(z_t) \odot
\left(\frac{\partial \epsilon_\theta(\hat{a}_t^{k}, z_t, k)}{\partial z_t}\right)^{\!\top}
\epsilon_\theta(\hat{a}_t^{k}, z_t, k),
\end{equation}
where $\gamma$ controls the step size of action-guided latent refinement.
Noise predictions conditioned on static and evolved latents are
\[
\varepsilon_k = \epsilon_{\theta}(\hat{a}_t^k, z_t, k),
\qquad
\tilde{\varepsilon}_k = \epsilon_{\theta}(\hat{a}_t^k, \tilde{z}_t^k, k),
\]
and are aligned through the cycle-consistent InfoNCE loss
\begin{equation}
\mathcal{L}_{\text{cont}}
= -\frac{1}{B}\sum_{i=1}^B
\log \frac{\exp\!\big( \mathrm{sim}(\varepsilon^i_k, \tilde{\varepsilon}^i_k) / \tau \big)}
{\sum_{j\neq i} \exp\!\big( \mathrm{sim}(\varepsilon^i_k, \tilde{\varepsilon}^j_k) / \tau \big)}.
\end{equation}
The overall loss remains unchanged:
\begin{equation}
\mathcal{L}_{\text{DP-AG}} = \mathcal{L}_{\text{DP}} + \lambda_{\text{cont}} \,\mathcal{L}_{\text{cont}} + \lambda_{\text{KL}} \,\mathcal{L}_{\text{KL}}.
\end{equation}
\textbf{Streaming update rule.}
Unlike the offline variant, which optimizes $\{\theta,\phi\}$ over a fixed dataset, the online version continuously updates the model as new observations arrive.
At each environmental time $t$, DP-AG both executes an action and uses the resulting transition as fresh training data:
\begin{enumerate}[leftmargin=*]
\item \textbf{Sample latent:} Draw $z_t \sim q_\phi(z_t|o_t)$ from the variational posterior of the current observation.
\item \textbf{Within-step diffusion (for $k=1,...,K$):}
\begin{enumerate}
\item Compute the noise prediction $\varepsilon_k = \epsilon_\theta(\hat{a}_t^k, z_t, k)$.
\item Update the latent via Eq.~\eqref{eq:online_latent_update} to obtain $\tilde{z}_t^k$.
\item Compute $\tilde{\varepsilon}_k$ and accumulate the contrastive loss $\mathcal{L}_{\text{cont}}$.
\item Accumulate the diffusion loss $\mathcal{L}_{\text{DP}}$.
\end{enumerate}
\item \textbf{Action execution:} After $K$ denoising steps, decode and execute the action $a_t$ in the environment.
\item \textbf{Immediate update:} Use the freshly collected pair $(o_t, a_t)$, optionally together with a short replay buffer containing the last $N$ transitions, to perform a lightweight gradient update of $\theta$ and $\phi$ on $\mathcal{L}_{\text{DP-AG}}$ using online RMSProp~\citep{ma2025efficient}.
\end{enumerate}
This schedule makes DP-AG online: model parameters are updated continually from streaming interaction data, rather than only once on a static offline dataset.
As a result, the policy can adapt on the fly to distributional shifts or novel dynamics while acting.
\textbf{Experimental setup.}
We evaluated both the offline and online variants of DP-AG on Push-T (both image- and keypoint-based settings) and Dynamic Push-T benchmarks.
Performance was measured as the average success rate over 50 online evaluation episodes.
Unless otherwise noted, hyperparameters matched those of the offline configuration.
\begin{table}[h]
\centering
\caption{Online streaming results of DP-AG compared to baselines (success rates, mean$\pm$std).}
\label{tab:online}
\begin{tabular}{lccc}
\toprule
Method & Push-T (img) & Push-T (kp) & Dynamic Push-T (img) \\
\midrule
Diffusion Policy (DP) & 0.87$\pm$0.04 & 0.95$\pm$0.03 & 0.65$\pm$0.85 \\
DP-AG (offline) & 0.93$\pm$0.02 & 0.99$\pm$0.01 & 0.80$\pm$0.53 \\
DP-AG (online) & 0.90$\pm$0.03 & 0.96$\pm$0.02 & 0.76$\pm$0.89 \\
\bottomrule
\end{tabular}
\end{table}
\textbf{Results.}
As shown in Table~\ref{tab:online}, the streaming variant of DP-AG yields slightly lower scores than its offline counterpart, which is expected given the limited replay and noisier updates inherent to online training.
Nevertheless, it consistently surpasses the standard DP baseline by a clear margin, demonstrating that DP-AG retains effective even under the more challenging online setting with continual real-time updates.
In summary, the online variant of DP-AG remains fully consistent with the offline model.
It employs the same VJP definition and latent update rule, with the time discretization constant absorbed into $\gamma$.
Cycle-consistent contrastive alignment is still applied at every diffusion step, ensuring the perception–action loop is preserved within each $[t, t{+}1]$ horizon.
Moreover, both the KL regularizer on $q_\phi$ and the overall training loss $\mathcal{L}_{\text{DP-AG}}$ remain unchanged; only the optimization schedule differs, shifting from offline minibatch training to lightweight streaming updates.
\section{Multimodal Decision-Making in DP-AG}
DPs are inherently capable of modeling multimodal action distributions, making them well-suited for capturing diverse strategies in complex tasks.
By extending DPs with action-guided latent evolution, DP-AG not only preserves this multimodality but amplifies it.
DP-AG enhances the representation of multiple effective strategies through two complementary components:
\begin{itemize}[leftmargin=*]
\item {\bf Action-conditioned latent evolution.}
DP-AG’s VJP-guided latent updates let the model refine features in response to action feedback, so each sampled trajectory adapts its latent state uniquely.
This supports parallel exploration of multiple strategies by leveraging stochasticity from both the diffusion policy and latent SDE.
\item {\bf Cycle-consistent contrastive alignment.}
The contrastive loss keeps features for the same action close and pushes different actions apart, organizing the latent-action space to separate and preserve multiple plausible strategies for each task.
\end{itemize}
{\bf Experimental setup.}
We evaluated multimodal capability on the Franka Kitchen benchmark, which naturally admits multiple strategies for tasks such as opening a drawer or flipping a switch.
40 successful trajectories per task were collected and encoded as key end-effector waypoints.
Distances between trajectories were computed via Dynamic Time Warping (DTW), followed by unsupervised clustering to identify distinct modes.
This procedure follows the existing work on multimodal diffusion policies~\citep{li2024learning}, which demonstrated that clustering trajectories provides a principled way to quantify behavioral diversity.
We then measured (a) the number of discovered modes, (b) inter-cluster distance (strategy distinctiveness), (c) intra-cluster variance (consistency within a strategy), and (d) success rates across all clusters.
\begin{table}[h!]
\centering
\setlength{\tabcolsep}{2.5pt}
\caption{Multiple strategy discovery and diversity on Franka Kitchen. Higher \#Modes ($\uparrow$) and inter-cluster distance ($\uparrow$), and lower intra-cluster variance ($\downarrow$), indicate better multimodal representation.}
\label{tab:multimodal1}
\vspace{2mm}
\begin{tabular}{lccccccc}
\toprule
Method & \# Modes& SR (t1) & SR (t2) & SR (t3) & SR (t4) & Inter-Cluster Dist. & Intra-Cluster Var.\\
\midrule
FlowPolicy & 1.1 & 0.96 & 0.86 & 0.95 & 0.87 & 2.0 & 1.2 \\
DP (Baseline) & 2.8 & 1.00 & 1.00 & 1.00 & 0.99 & 7.1 & 1.2 \\
\midrule
{\bf DP-AG (Ours)} & 3.2 & 1.00 & 1.00 & 1.00 & 1.00 & 9.5 & 1.3 \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[h!]
\centering
\setlength{\tabcolsep}{4pt}
\caption{Cluster-specific analysis of strategies discovered by our DP-AG on Franka Kitchen.}
\label{tab:multimodal2}
\vspace{2mm}
\begin{tabular}{c c c c c c l}
\toprule
Mode ID & Coverage (\%) & SR (t1) & SR (t2) & SR (t3) & SR (t4) & Strategy Description \\
\midrule
1 & 41 & 1.00 & 1.00 & 1.00 & 1.00 & Left-handed drawer pull \\
2 & 32 & 1.00 & 1.00 & 1.00 & 1.00 & Right-handed drawer pull \\
3 & 27 & 1.00 & 1.00 & 1.00 & 1.00 & Two-step approach / mixed arm \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[h!]
\centering
\setlength{\tabcolsep}{6pt}
\caption{Robustness to distribution shifts in Franka Kitchen (task t4).}
\label{tab:multimodal3}
\vspace{2mm}
\begin{tabular}{l c c c l}
\toprule
Method & \# Modes & SR (t4) & Switch Rate (\%) & Comment \\
\midrule
FlowPolicy & 1.1 & 0.64 & 12 & Rarely adapts, prone to failure \\
DP (Baseline) & 2.2 & 0.81 & 35 & Sometimes adapts, less reliable \\
\midrule
{\bf DP-AG (Ours)} & 2.9 & 0.95 & 58 & Switches to alternatives if blocked \\
\bottomrule
\end{tabular}
\end{table}
\textbf{Results.}
Table~\ref{tab:multimodal1} presents the overall diversity analysis, which demonstrates that our DP-AG uncovers and maintains on average more than three distinct strategies per task, significantly more than FlowPolicy (1.1) and slightly more than baseline DP (2.8).
These strategies are not minor variants: the inter-cluster distance is significantly higher, confirming that the modes correspond to distinct behaviors.
At the same time, intra-cluster variance remains low, ensuring compact and consistent execution.
Table~\ref{tab:multimodal2} provides a cluster-specific breakdown, highlighting interpretable strategies such as left-handed drawer pulls, right-handed pulls, and two-step mixed-arm approaches, each achieving near-perfect success.
Moreover, Table~\ref{tab:multimodal3} evaluates robustness under distribution shifts.
When trajectories were blocked, DP-AG preserved an average of 2.9 modes with a 58\% switch rate to alternative strategies, while DP and FlowPolicy degraded significantly.
These findings demonstrate that our DP-AG not only retains the multimodal decision-making ability of DPs but actively encourages the emergence of more rich, diverse, and interpretable strategies.
Its robustness to distribution shifts highlights its potential for real-world robotic applications where flexibility and adaptability are important.
\section{More Experiments on Trajectory Planning}\label{sec:more_trajectory}
To situate our DP-AG within the broader family of trajectory generation methods, we additionally compare it against Diffuser~\citep{janner2022diffuser} and Hierarchical Diffuser~\citep{chen2024simple}.
Following the experimental protocol of~\citep{chen2024simple}, we adopt two benchmark families: (i) Maze2D and Multi2D tasks from D4RL, which evaluate long-horizon navigation in continuous control environments, and (ii) the multi-stage FrankaKitchen benchmark, which evaluates robotic manipulation requiring both skill composition and generalization to unseen states.
In Maze2D and Multi2D, an agent must reach goals in varied layouts (U-Maze, Medium, Large), with performance measured by average return.
In FrankaKitchen, policies are rolled out from diverse initial states and evaluated by the number of sub-tasks completed within long-horizon episodes.
{\bf Experimental setup.}
We match the evaluation protocol of the baselines, adopting the same trajectory segmentation and planning horizons ({\em e.g.,} $K=15$ for Maze2D and Multi2D, $K=4$ for Gym-MuJoCo, and horizon lengths $H=120$ for U-Maze and $H=255$ for Maze2D Medium).
All models are evaluated over 100 random seeds, reporting the average return or goal completion rate under identical training, validation, and testing splits.
\begin{table}[h!]
\centering
\caption{Trajectory planning on Maze2D and Multi2D (long-horizon). Performance is measured by average return (higher is better).}
\label{tab:maze2d}
\begin{tabular}{lccc}
\toprule
Environment & Diffuser & Hierarchical Diffuser & \textbf{DP-AG (Ours)} \\
\midrule
Maze2D U-Maze & 113.9 $\pm$ 3.1 & 128.4 $\pm$ 3.6 & \textbf{142.7 $\pm$ 2.9} \\
Maze2D Medium & 121.5 $\pm$ 2.7 & 135.6 $\pm$ 3.0 & \textbf{150.2 $\pm$ 2.6} \\
Maze2D Large & 123.0 $\pm$ 6.4 & 155.8 $\pm$ 2.5 & \textbf{172.1 $\pm$ 2.2} \\
Multi2D U-Maze & 128.9 $\pm$ 1.8 & 144.1 $\pm$ 1.2 & \textbf{168.2 $\pm$ 1.2} \\
Multi2D Medium & 127.2 $\pm$ 3.4 & 140.2 $\pm$ 1.6 & \textbf{153.7 $\pm$ 1.3} \\
Multi2D Large & 132.1 $\pm$ 5.8 & 165.5 $\pm$ 0.6 & \textbf{182.2 $\pm$ 0.5} \\
\midrule
Average & 124.4 & 145.0 & \textbf{161.5} \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[h!]
\centering
\caption{Multi-stage robotic manipulation on FrankaKitchen. Performance is measured by average number of completed sub-tasks.}
\label{tab:kitchen}
\begin{tabular}{lccc}
\toprule
Task & Diffuser & Hierarchical Diffuser & \textbf{DP-AG (Ours)} \\
\midrule
Partial Kitchen & 56.2 $\pm$ 5.4 & 73.3 $\pm$ 1.4 & \textbf{82.2 $\pm$ 1.3} \\
Mixed Kitchen & 50.0 $\pm$ 8.8 & 71.7 $\pm$ 2.7 & \textbf{78.5 $\pm$ 2.1} \\
\midrule
Average & 53.1 & 72.5 & \textbf{80.4} \\
\bottomrule
\end{tabular}
\end{table}
\textbf{Results.}
Across all benchmarks, DP-AG achieves the highest returns, with average improvements of +16.5 over the Hierarchical Diffuser and +37.1 over the Diffuser in Maze2D and Multi2D (Table~\ref{tab:maze2d}).
These performance gains indicate that our DP-AG generates longer and higher-quality trajectories, especially in large and complex environments.
In the multi-stage FrankaKitchen benchmark (Table~\ref{tab:kitchen}), DP-AG completes on average 80.4 sub-tasks, significantly outperforming both Hierarchical Diffuser (72.5) and Diffuser (53.1), which highlights DP-AG’s effectiveness in composing diverse skills and generalizing them to extended horizons.
Together, these results indicate that the action-guided perception-action loop in DP-AG provides clear advantages for both navigation and manipulation tasks, beyond what trajectory diffusion alone can achieve.
\section{Extending DP-AG with World Models}\label{app:world_model}
World models such as DreamerV3~\citep{hafner2025dreamerv3} and Unified World Models (UWM)~\citep{zhu2025uwm} learn to ``imagine'' latent trajectories by predicting future states and rewards.
While these predictions enable long-horizon planning, the latent trajectories are typically static within each rollout: once generated, they do not adapt to evolving action feedback.
In contrast, DP-AG continually updates latent features within each action diffusion step, guided by action-conditioned VJPs.
By combining with world models, our DP-AG gains long-horizon foresight while retaining its intra-step refinement.
\textbf{Hybrid Architecture.}
We integrate DP-AG into UWM by adding an intra-step feedback mechanism on top of UWM’s diffusion transformer backbone:
\begin{itemize}[leftmargin=*]
\item \textbf{UWM Backbone.}
UWM jointly learns to predict actions and future visual observations by treating them as parallel diffusion processes.
This unified design allows the same model to serve flexibly as a policy, a forward dynamics predictor, an inverse dynamics model, or a video generator, simply by adjusting which modalities are denoised.
The architecture uses ResNet-based encoders for observations, spatiotemporal patching for latent images, shallow MLPs for actions, and transformer layers conditioned through adaptive normalization with additional register tokens to promote information exchange across modalities.
\item \textbf{DP-AG Intra-step Feedback.} During action generation, DP-AG introduces a feedback signal into UWM’s latent representations.
At each denoising step, the current action prediction influences how features are refined, ensuring that perception and action remain aligned as the sequence unfolds.
This intra-step refinement captures the core principle of DP-AG, where even a fixed observation is reinterpreted in light of evolving actions.
\item \textbf{Bounded Consistency.} To prevent excessive feature drift, we apply a lightweight penalty that encourages smooth updates across steps, complemented by DP-AG’s cycle-consistent contrastive anchor.
This balance preserves stability while still enabling action-driven adaptability.
\end{itemize}
\textbf{Implementation Details.}
\begin{itemize}[leftmargin=*]
\item {\bf Observation and Conditioning.}
Visual inputs are processed by frozen VAEs into compact latent grids, which are patchified and combined with action tokens.
Diffusion time steps for both action and observation branches are encoded as embeddings and injected into each transformer block.
\item {\bf Training Objective.}
UWM is trained with a coupled denoising loss over action and future-image predictions.
For action-free videos, the action branch is masked out, allowing the same objective to learn from both robot trajectories and pure video data.
\item {\bf DP-AG Integration.}
The action-guided feedback mechanism is applied during both imagination and execution.
At training time, refinement losses are added alongside UWM’s standard objectives; at inference time, feedback ensures that action sampling and imagined rollouts remain tightly coupled.
\item {\bf Compute.} Following UWM’s reported setup, training runs efficiently on Nvidia A100 GPUs.
\end{itemize}
\begin{table}[h!]
\centering
\setlength{\tabcolsep}{3.5pt}
\caption{Success rate for combining DP-AG with UWM on LIBERO tasks.}
\label{tab:libero}
\begin{tabular}{lccccc c}
\toprule
Method & Book-Caddy & Soup-Cheese & Bowl-Drawer & Moka-Moka & Mug-Mug & Average \\
\midrule
DP & 0.78 & 0.88 & 0.77 & 0.65 & 0.53 & 0.71 \\
DP-AG & 0.86 & 0.92 & 0.85 & 0.72 & 0.60 & 0.79 \\
UWM & 0.91 & 0.93 & 0.80 & 0.68 & 0.65 & 0.79 \\
UWM\,+\,DP-AG & \textbf{0.94} & \textbf{0.95} & \textbf{0.87} & \textbf{0.75} & \textbf{0.70} & \textbf{0.84} \\
\bottomrule
\end{tabular}
\end{table}
\textbf{Experimental Results on LIBERO.}
Following the UWM protocol, we fine-tune each model on the LIBERO benchmark tasks and report average success rates.
As shown in Table~\ref{tab:libero}, the hybrid model UWM\,+\,DP-AG outperforms both UWM and DP-AG alone, especially in manipulation tasks requiring real-time adaptation ({\em e.g.,} Mug-Mug).
This confirms that action-guided latent updates make world models more responsive, while world models provide long-horizon foresight to complement DP-AG's intra-step refinement.
\textbf{Discussion.}
This extension highlights a conceptual bridge: world models offer look-ahead imagination of possible futures, while DP-AG provides real-time corrective adaptation from action feedback within each imagined step.
Together, they yield policies that are both farsighted and responsive, achieving long-horizon planning while adapting online to evolving cues.
We will conduct an in-depth study on this subject in the future work.
\end{document}
\begin{thebibliography}{42}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Arora and Doshi(2021)]{arora2021survey}
S.~Arora and P.~Doshi.
\newblock A survey of inverse reinforcement learning: Challenges, methods and progress.
\newblock \emph{Artificial Intelligence}, 297:\penalty0 103500, 2021.
\bibitem[Bain and Sammut(1995)]{bain1995framework}
M.~Bain and C.~Sammut.
\newblock A framework for behavioural cloning.
\newblock In \emph{Machine Intelligence 15}, pages 103--129, 1995.
\bibitem[Black et~al.(2024)Black, Brown, Driess, Esmail, Equi, Finn, Fusai, Groom, Hausman, Ichter, Jakubczak, Jones, Ke, Levine, Li-Bell, Mothukuri, Nair, Pertsch, Shi, Tanner, Vuong, Walling, Wang, and Zhilinsky]{black2024pi_0}
K.~Black, N.~Brown, D.~Driess, A.~Esmail, M.~Equi, C.~Finn, N.~Fusai, L.~Groom, K.~Hausman, B.~Ichter, S.~Jakubczak, T.~Jones, L.~Ke, S.~Levine, A.~Li-Bell, M.~Mothukuri, S.~Nair, K.~Pertsch, L.~X. Shi, J.~Tanner, Q.~Vuong, A.~Walling, H.~Wang, and U.~Zhilinsky.
\newblock $\pi_0$: A vision-language-action flow model for general robot control.
\newblock \emph{arXiv preprint arXiv:2410.24164}, 2024.
\bibitem[Brooks(1991)]{brooks1991intelligence}
R.~A. Brooks.
\newblock Intelligence without representation.
\newblock \emph{Artificial intelligence}, 47\penalty0 (1-3):\penalty0 139--159, 1991.
\bibitem[Chen et~al.(2024)Chen, Deng, Kawaguchi, Gulcehre, and Ahn]{chen2024simple}
C.~Chen, F.~Deng, K.~Kawaguchi, C.~Gulcehre, and S.~Ahn.
\newblock Simple hierarchical planning with diffusion.
\newblock In \emph{International Conference on Learning Representations (ICLR)}, 2024.
\bibitem[Chen et~al.(2018)Chen, Rubanova, Bettencourt, and Duvenaud]{chen2018neural}
R.~T. Chen, Y.~Rubanova, J.~Bettencourt, and D.~K. Duvenaud.
\newblock Neural ordinary differential equations.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 31, 2018.
\bibitem[Chen et~al.(2020)Chen, Kornblith, Norouzi, and Hinton]{chen2020simple}
T.~Chen, S.~Kornblith, M.~Norouzi, and G.~Hinton.
\newblock A simple framework for contrastive learning of visual representations.
\newblock In \emph{International Conference on Machine Learning (ICML)}, pages 1597--1607. PMLR, 2020.
\bibitem[Chi et~al.(2023)Chi, Xu, Feng, Cousineau, Du, Burchfiel, Tedrake, and Song]{chi2023diffusion}
C.~Chi, Z.~Xu, S.~Feng, E.~Cousineau, Y.~Du, B.~Burchfiel, R.~Tedrake, and S.~Song.
\newblock Diffusion policy: Visuomotor policy learning via action diffusion.
\newblock \emph{The International Journal of Robotics Research}, page 02783649241273668, 2023.
\bibitem[Codevilla et~al.(2018)Codevilla, M{\"u}ller, L{\'o}pez, Koltun, and Dosovitskiy]{codevilla2018end}
F.~Codevilla, M.~M{\"u}ller, A.~L{\'o}pez, V.~Koltun, and A.~Dosovitskiy.
\newblock End-to-end driving via conditional imitation learning.
\newblock In \emph{2018 IEEE International Conference on Robotics and Automation (ICRA)}, pages 4693--4700. IEEE, 2018.
\bibitem[Florence et~al.(2022)Florence, Lynch, Zeng, Ramirez, Wahid, Downs, Wong, Lee, Mordatch, and Tompson]{florence2022implicit}
P.~Florence, C.~Lynch, A.~Zeng, O.~A. Ramirez, A.~Wahid, L.~Downs, A.~Wong, J.~Lee, I.~Mordatch, and J.~Tompson.
\newblock Implicit behavioral cloning.
\newblock In \emph{Conference on Robot Learning (CoRL)}, pages 158--168. PMLR, 2022.
\bibitem[Gupta et~al.(2020)Gupta, Kumar, Lynch, Levine, and Hausman]{gupta2020}
A.~Gupta, V.~Kumar, C.~Lynch, S.~Levine, and K.~Hausman.
\newblock Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning.
\newblock In L.~P. Kaelbling, D.~Kragic, and K.~Sugiura, editors, \emph{Conference on Robot Learning (CoRL)}, volume 100 of \emph{Proceedings of Machine Learning Research}, pages 1025--1037, 30 Oct--01 Nov 2020.
\bibitem[Hafner et~al.(2019)Hafner, Lillicrap, Fischer, Villegas, Ha, Lee, and Davidson]{hafner2019learning}
D.~Hafner, T.~Lillicrap, I.~Fischer, R.~Villegas, D.~Ha, H.~Lee, and J.~Davidson.
\newblock Learning latent dynamics for planning from pixels.
\newblock In \emph{International Conference on Machine Learning (ICML)}, pages 2555--2565. PMLR, 2019.
\bibitem[Hafner et~al.(2025)Hafner, Pasukonis, Ba, and Lillicrap]{hafner2025dreamerv3}
D.~Hafner, J.~Pasukonis, J.~Ba, and T.~Lillicrap.
\newblock Mastering diverse control tasks through world models.
\newblock \emph{Nature}, pages 1--7, 2025.
\bibitem[He et~al.(2016)He, Zhang, Ren, and Sun]{he2016deep}
K.~He, X.~Zhang, S.~Ren, and J.~Sun.
\newblock Deep residual learning for image recognition.
\newblock In \emph{Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, pages 770--778, 2016.
\bibitem[Higgins et~al.(2017)Higgins, Matthey, Pal, Burgess, Glorot, Botvinick, Mohamed, and Lerchner]{higgins2017betavae}
I.~Higgins, L.~Matthey, A.~Pal, C.~Burgess, X.~Glorot, M.~Botvinick, S.~Mohamed, and A.~Lerchner.
\newblock beta-{VAE}: Learning basic visual concepts with a constrained variational framework.
\newblock In \emph{International Conference on Learning Representations (ICLR)}, 2017.
\bibitem[Ho et~al.(2020)Ho, Jain, and Abbeel]{ho2020denoising}
J.~Ho, A.~Jain, and P.~Abbeel.
\newblock Denoising diffusion probabilistic models.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 33:\penalty0 6840--6851, 2020.
\bibitem[Hu et~al.(2024)Hu, Liu, Liu, and Liu]{hu2024adaflow}
X.~Hu, Q.~Liu, X.~Liu, and B.~Liu.
\newblock Adaflow: Imitation learning with variance-adaptive flow-based policies.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 37:\penalty0 138836--138858, 2024.
\bibitem[Jackson et~al.(2024)Jackson, Matthews, Lu, Ellis, Whiteson, and Foerster]{jackson2024policy}
M.~T. Jackson, M.~T. Matthews, C.~Lu, B.~Ellis, S.~Whiteson, and J.~Foerster.
\newblock Policy-guided diffusion.
\newblock \emph{arXiv preprint arXiv:2404.06356}, 2024.
\bibitem[Janner et~al.(2022)Janner, Du, Tenenbaum, and Levine]{janner2022diffuser}
M.~Janner, Y.~Du, J.~B. Tenenbaum, and S.~Levine.
\newblock Planning with diffusion for flexible behavior synthesis.
\newblock In \emph{International Conference on Machine Learning (ICML)}, 2022.
\bibitem[Jia et~al.(2024)Jia, Blessing, Jiang, Reuss, Donat, Lioutikov, and Neumann]{jia2024towards}
X.~Jia, D.~Blessing, X.~Jiang, M.~Reuss, A.~Donat, R.~Lioutikov, and G.~Neumann.
\newblock Towards diverse behaviors: A benchmark for imitation learning with human demonstrations.
\newblock In \emph{International Conference on Learning Representations (ICLR)}, 2024.
\bibitem[Kim et~al.(2024)Kim, Pertsch, Karamcheti, Xiao, Balakrishna, Nair, Rafailov, Foster, Sanketi, Vuong, Kollar, Burchfiel, Tedrake, Sadigh, Levine, Liang, and Finn]{kim2024openvla}
M.~J. Kim, K.~Pertsch, S.~Karamcheti, T.~Xiao, A.~Balakrishna, S.~Nair, R.~Rafailov, E.~P. Foster, P.~R. Sanketi, Q.~Vuong, T.~Kollar, B.~Burchfiel, R.~Tedrake, D.~Sadigh, S.~Levine, P.~Liang, and C.~Finn.
\newblock Open{VLA}: An open-source vision-language-action model.
\newblock In \emph{8th Annual Conference on Robot Learning (CoRL)}, 2024.
\bibitem[Kingma et~al.(2014)Kingma, Welling, et~al.]{kingma2014auto}
D.~P. Kingma, M.~Welling, et~al.
\newblock Auto-encoding variational bayes.
\newblock \emph{The Second International Conference on Learning Representations (ICLR)}, 2014.
\bibitem[Kingma et~al.(2015)Kingma, Salimans, and Welling]{kingma2015variational}
D.~P. Kingma, T.~Salimans, and M.~Welling.
\newblock Variational dropout and the local reparameterization trick.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 28, 2015.
\bibitem[Li et~al.(2024)Li, Krohn, Chen, Ajay, Agrawal, and Chalvatzaki]{li2024learning}
S.~Li, R.~Krohn, T.~Chen, A.~Ajay, P.~Agrawal, and G.~Chalvatzaki.
\newblock Learning multimodal behaviors from scratch with diffusion policy gradient.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 37:\penalty0 38456--38479, 2024.
\bibitem[Li et~al.(2019)Li, Wong, Chen, and Duvenaud]{li2019scalable}
X.~Li, T.-K.~L. Wong, R.~T.~Q. Chen, and D.~K. Duvenaud.
\newblock Scalable gradients and variational inference for stochastic differential equations.
\newblock In \emph{Second Symposium on Advances in Approximate Bayesian Inference}, 2019.
\bibitem[Lipman et~al.(2023)Lipman, Chen, Ben-Hamu, Nickel, and Le]{lipman2023flow}
Y.~Lipman, R.~T.~Q. Chen, H.~Ben-Hamu, M.~Nickel, and M.~Le.
\newblock Flow matching for generative modeling.
\newblock In \emph{The Eleventh International Conference on Learning Representations (ICLR)}, 2023.
\bibitem[Ma et~al.(2025)Ma, Chen, Wang, Li, and Dai]{ma2025efficient}
H.~Ma, T.~Chen, K.~Wang, N.~Li, and B.~Dai.
\newblock Efficient online reinforcement learning for diffusion policy.
\newblock In \emph{International Conference on Machine Learning (ICML)}, 2025.
\bibitem[Mandlekar et~al.(2022)Mandlekar, Xu, Wong, Nasiriany, Wang, Kulkarni, Fei-Fei, Savarese, Zhu, and Mart\'in-Mart\'in]{mandlekar22a}
A.~Mandlekar, D.~Xu, J.~Wong, S.~Nasiriany, C.~Wang, R.~Kulkarni, L.~Fei-Fei, S.~Savarese, Y.~Zhu, and R.~Mart\'in-Mart\'in.
\newblock What matters in learning from offline human demonstrations for robot manipulation.
\newblock In \emph{Conference on Robot Learning (CoRL)}. PMLR, 2022.
\bibitem[Nichol and Dhariwal(2021)]{nichol2021improved}
A.~Q. Nichol and P.~Dhariwal.
\newblock Improved denoising diffusion probabilistic models.
\newblock In \emph{International Conference on Machine Learning (ICML)}, pages 8162--8171. PMLR, 2021.
\bibitem[O'regan and No{\"e}(2001)]{o2001sensorimotor}
J.~K. O'regan and A.~No{\"e}.
\newblock A sensorimotor account of vision and visual consciousness.
\newblock \emph{Behavioral and Brain Sciences}, 24\penalty0 (5):\penalty0 939--973, 2001.
\bibitem[Osa et~al.(2018)Osa, Pajarinen, Neumann, Bagnell, Abbeel, Peters, et~al.]{osa2018algorithmic}
T.~Osa, J.~Pajarinen, G.~Neumann, J.~A. Bagnell, P.~Abbeel, J.~Peters, et~al.
\newblock An algorithmic perspective on imitation learning.
\newblock \emph{Foundations and Trends{\textregistered} in Robotics}, 7\penalty0 (1-2):\penalty0 1--179, 2018.
\bibitem[Parulekar et~al.(2023)Parulekar, Collins, Shanmugam, Mokhtari, and Shakkottai]{parulekar2023infonce}
A.~Parulekar, L.~Collins, K.~Shanmugam, A.~Mokhtari, and S.~Shakkottai.
\newblock Infonce loss provably learns cluster-preserving representations.
\newblock In \emph{The Thirty Sixth Annual Conference on Learning Theory}, pages 1914--1961. PMLR, 2023.
\bibitem[Shafiullah et~al.(2022)Shafiullah, Cui, Altanzaya, and Pinto]{shafiullah2022behavior}
N.~M. Shafiullah, Z.~Cui, A.~A. Altanzaya, and L.~Pinto.
\newblock Behavior transformers: Cloning $ k $ modes with one stone.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 35:\penalty0 22955--22968, 2022.
\bibitem[Song et~al.(2021)Song, Sohl-Dickstein, Kingma, Kumar, Ermon, and Poole]{song2021scorebased}
Y.~Song, J.~Sohl-Dickstein, D.~P. Kingma, A.~Kumar, S.~Ermon, and B.~Poole.
\newblock Score-based generative modeling through stochastic differential equations.
\newblock In \emph{International Conference on Learning Representations (ICLR)}, 2021.
\bibitem[Sutton et~al.(1999)Sutton, McAllester, Singh, and Mansour]{sutton1999policy}
R.~S. Sutton, D.~McAllester, S.~Singh, and Y.~Mansour.
\newblock Policy gradient methods for reinforcement learning with function approximation.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 12, 1999.
\bibitem[Tan and Xie(2020)]{tan2020extension}
F.~Tan and A.~Xie.
\newblock An extension of the am--gm--hm inequality.
\newblock \emph{Bulletin of the Iranian Mathematical Society}, 46:\penalty0 245--251, 2020.
\bibitem[Wang et~al.(2024)Wang, Bae, Chen, Zhang, Sigal, and de~Silva]{wang2024what}
J.~Wang, W.~Bae, J.~Chen, K.~Zhang, L.~Sigal, and C.~W. de~Silva.
\newblock What has been overlooked in contrastive source-free domain adaptation: Leveraging source-informed latent augmentation within neighborhood context.
\newblock \emph{Transactions on Machine Learning Research}, 2024.
\newblock ISSN 2835-8856.
\newblock Featured Certification.
\bibitem[Wang et~al.(2023)Wang, Hunt, and Zhou]{wang2023diffusion}
Z.~Wang, J.~J. Hunt, and M.~Zhou.
\newblock Diffusion policies as an expressive policy class for offline reinforcement learning.
\newblock In \emph{The Eleventh International Conference on Learning Representations (ICLR)}, 2023.
\bibitem[Watter et~al.(2015)Watter, Springenberg, Boedecker, and Riedmiller]{watter2015embed}
M.~Watter, J.~Springenberg, J.~Boedecker, and M.~Riedmiller.
\newblock Embed to control: A locally linear latent dynamics model for control from raw images.
\newblock \emph{Advances in Neural Information Processing Systems (NeurIPS)}, 28, 2015.
\bibitem[Wu et~al.(2023)Wu, Escontrela, Hafner, Abbeel, and Goldberg]{wu2023daydreamer}
P.~Wu, A.~Escontrela, D.~Hafner, P.~Abbeel, and K.~Goldberg.
\newblock Daydreamer: World models for physical robot learning.
\newblock In \emph{Conference on Robot Learning (CoRL)}, pages 2226--2240. PMLR, 2023.
\bibitem[Zhang et~al.(2025)Zhang, Liu, Fan, Liu, Zeng, and Liu]{zhang2025flowpolicy}
Q.~Zhang, Z.~Liu, H.~Fan, G.~Liu, B.~Zeng, and S.~Liu.
\newblock Flowpolicy: Enabling fast and robust 3d flow-based policy via consistency flow matching for robot manipulation.
\newblock In \emph{Proceedings of the AAAI Conference on Artificial Intelligence}, volume~39, pages 14754--14762, 2025.
\bibitem[Zhu et~al.(2025)Zhu, Yu, Feng, Burchfiel, Shah, and Gupta]{zhu2025uwm}
C.~Zhu, R.~Yu, S.~Feng, B.~Burchfiel, P.~Shah, and A.~Gupta.
\newblock Unified world models: Coupling video and action diffusion for pretraining on large robotic datasets.
\newblock In \emph{Proceedings of Robotics: Science and Systems (RSS)}, 2025.
\end{thebibliography}