DADP: Domain Adaptive Diffusion Policy

1. Problem Statement

The paper addresses the challenge of Domain Adaptive Policy Learning in reinforcement learning (RL). While learning-based policies have achieved success in specific environments, they often fail to generalize to unseen domains (Out-of-Distribution or OOD) where transition dynamics change (e.g., friction, gravity, or morphological changes).

Existing approaches typically rely on domain representation learning to capture domain-specific information. However, the authors identify a critical flaw in current methods:

• Entanglement of Static and Dynamic Information: Most methods use recent history (context) to predict the next state. This causes the learned representation to entangle static domain factors (e.g., gravity, limb length) with time-varying dynamical properties (e.g., instantaneous velocity, higher-order derivatives).
• Consequence: This mixture confuses the conditioned policy, leading to representation drift within the same domain and limiting zero-shot adaptation capabilities.
• Ineffective Utilization: Even when representations are learned, standard methods often use them merely as extra input concatenation, failing to fully leverage the structure of the generative process in diffusion models.

2. Methodology: DADP

The authors propose DADP (Domain-Adaptive Diffusion Policy), which achieves robust adaptation through two core innovations: Unsupervised Disentanglement and Domain-Aware Diffusion Injection.

A. Unsupervised Disentanglement: Lagged Context Dynamical Prediction

To extract pure static domain representations ($\xi$) while filtering out time-varying information ($\omega_t$), the authors introduce Lagged Context Dynamical Prediction.

• Mechanism: Instead of using the immediate past history ($\tau_t$) to predict the next state ($s_{t+1}$), the model uses a context with a temporal offset $\Delta t$ ($\tau_{t-\Delta t}$).
• Logic:
  • When $\Delta t$ is small (e.g., 1), the context contains instantaneous cues (velocity) that help predict the immediate next state, causing the representation to encode these transient properties.
  • As $\Delta t \to \infty$ (e.g., using context from a different episode in the same domain), the instantaneous velocity cues become irrelevant for predicting the current state.
  • The model is forced to rely solely on static domain factors (like gravity or mass) to minimize prediction loss.
• Result: This unsupervised strategy effectively disentangles the static domain representation from time-varying dynamics without requiring ground-truth labels.

B. Domain-Aware Diffusion Injection

Once high-quality static representations ($z$) are learned, DADP integrates them directly into the generative process of the diffusion policy, rather than just conditioning the network.

• Biased Prior Distribution: Instead of starting the denoising process from a standard Gaussian noise distribution, DADP initializes the process from a Gaussian Mixture centered around the learned domain representation $z$.
  • Formula: $a_t = \sqrt{\bar{\alpha}_t}(a_0 - z) + z + \sqrt{1-\bar{\alpha}_t}\epsilon$.
  • This biases the prior distribution toward domain-specific action modalities.
• Reformulated Diffusion Target: The model is trained to predict a composite term $\hat{\epsilon}_\theta$ that includes both the noise and the representation shift.
  • This provides extra supervision at every denoising step, guiding the trajectory toward the correct domain manifold more efficiently.

3. Key Contributions

1. Lagged Context Dynamical Prediction: A novel, simple, and effective unsupervised method to learn static domain representations by breaking temporal correlations between context and current state.
2. Representation-Prediction Injection: A new paradigm for utilizing domain representations in diffusion policies by biasing the prior distribution and reformulating the diffusion target, rather than simple concatenation.
3. Superior Zero-Shot Performance: Demonstrated state-of-the-art (SOTA) performance in zero-shot domain adaptation across diverse locomotion and manipulation tasks.
4. Open-Source Release: The authors released a complete codebase, datasets, and data generation pipeline to facilitate community research.

4. Experimental Results

The authors evaluated DADP on challenging benchmarks including MuJoCo (HalfCheetah, Walker2d, Ant, Hopper) and Adroit (manipulation tasks).

• Performance: DADP consistently outperformed SOTA baselines (CORRO, Prompt-DT, Meta-DT) in both In-Distribution (IID) and Out-of-Distribution (OOD) settings.
  • Example: On Walker2d OOD, DADP achieved a normalized return of 2834, significantly higher than Meta-DT (889) and CORRO (66).
  • Example: On HalfCheetah OOD, DADP achieved 3001, compared to Meta-DT's 3174 (note: DADP showed higher stability with lower std dev).
• Ablation Studies:
  • Effect of $\Delta t$: Increasing the lag $\Delta t$ significantly improved representation quality (measured by linear probe accuracy and reconstruction loss), confirming the disentanglement hypothesis.
  • Utilization Strategy: The "w/ Predict" strategy (biasing prior + reformulating target) yielded the highest performance, outperforming simple conditioning or prior biasing alone.
• Mastery Metric: DADP achieved the highest "mastery" (percentage of domains where the policy reaches 60% of expert performance), indicating its ability to robustly control diverse embodiments.

5. Significance

• Theoretical Insight: The paper provides a crucial insight that time-local cues in context are detrimental to learning static domain representations. By introducing a temporal lag, the authors show how to unsupervisedly filter these out.
• Generative Modeling: It advances the application of diffusion models in RL by demonstrating that injecting domain knowledge into the prior distribution and denoising target is more effective than standard conditioning.
• Practical Impact: The method enables a single policy to adapt to unseen physical environments (e.g., different robot morphologies or friction coefficients) without fine-tuning, addressing a major bottleneck in deploying RL agents in the real world.
• Scalability: The approach is data-efficient (unsupervised) and does not require complex contrastive learning objectives or extra data generation models.

In conclusion, DADP represents a significant step forward in making reinforcement learning policies robust to dynamic changes by fundamentally rethinking how domain information is extracted and utilized within generative models.