Technical Summary: Bridging Domain Gaps with Target-Aligned Generation for Offline Reinforcement Learning

Problem Statement

Cross-domain offline reinforcement learning (RL) aims to transfer a policy from a source domain to a target domain using only pre-collected datasets, without any online interaction. This setting is critical for applications like heterogeneous robot control and simulation-to-real transfer, where online data collection is cost-prohibitive or unsafe.

A primary challenge in this regime is the distributional mismatch between source and target domains, particularly when the target dataset is extremely limited ($|D_{tar}| \ll |D_{src}|$). Existing methods often attempt to bridge this gap by filtering source data to find transitions "near" the target distribution. However, the authors demonstrate that when the dynamics gap is large, directly incorporating even "nearest-neighbor" source transitions can induce substantial distributional mismatch, degrading performance compared to training on the target data alone. Conversely, relying solely on limited target data often fails to provide sufficient state coverage for effective policy learning.

Methodology: Target-Aligned Coverage Expansion (TCE)

To address the trade-off between leveraging source data and avoiding distributional mismatch, the authors propose Target-Aligned Coverage Expansion (TCE). TCE is a framework that theoretically guides how source data should be utilized: either by direct incorporation of target-near transitions or by using source data to expand state coverage via target-aligned generation.

Theoretical Foundation

The method is grounded in a theoretical analysis of the performance gap between a target Markov Decision Process (MDP) and a mixture MDP (combining source and target dynamics). The authors derive a bound (Theorem 4.1) showing that the performance gap is governed by two components:

1. Dynamics Gap: The total variation distance between the source transition dynamics ($P_{src}$) and the target dynamics ($P_{tar}$).
2. Generation Error: The distance between a generated approximate target transition ($\hat{P}_{tar}$) and the true target dynamics.

TCE controls these two error sources by adjusting two coefficients:

• $\lambda_{mix}$: Controls the fraction of source data directly incorporated (minimizing generation error but potentially increasing dynamics gap).
• $\lambda_{cov}$: Controls the extent of state coverage expansion via generation (minimizing dynamics gap but potentially increasing generation error).

Algorithmic Framework

TCE employs a dual score-based generative model built on Stochastic Differential Equations (SDEs) to synthesize transitions that are consistent with target dynamics but cover an expanded state space.

1. Dual Score Networks:

   • Mixture-State Score Network ($q^{mix}_\theta$): Trained on a controllable mixture of target states and source states that are "near" the target (selected via a nearest-neighbor criterion). This network models the expanded state distribution.
   • Target-Transition Score Network ($q^{tran}_\theta$): Trained solely on target transitions. It generates the next state conditioned on the current state, ensuring the transition dynamics align with the target domain.

2. Two-Stage Sampling Procedure:

   • Stage 1 (State Sampling): Samples a state $\hat{s}_t$ from the mixture-state distribution using the reverse SDE guided by $q^{mix}_\theta$.
   • Stage 2 (Transition Sampling): Generates the next state $\hat{s}_{t+1}$ conditioned on $\hat{s}_t$ using the reverse SDE guided by $q^{tran}_\theta$.
   • Auxiliary Models: An inverse dynamics model (trained on target data) and a reward model (trained on both domains) are used to recover the action $\hat{a}_t$ and reward $\hat{r}_t$, completing the transition tuple.

3. Policy Learning:
   The final training dataset ($D_{train}$) is constructed by combining the target dataset ($D_{tar}$), the generated target-aligned transitions ($D^{gen}_{\lambda_{cov}}$), and optionally, a filtered subset of source transitions ($D^{src}_{\lambda_{mix}}$). An offline RL algorithm (Implicit Q-Learning, IQL) is trained on this augmented dataset with a KL-regularization term to penalize deviation from the target behavior policy.

Key Contributions

• Target-Aligned Coverage Expansion (TCE): A novel framework that moves beyond simple source data filtering. It explicitly decides whether to use source data directly or to leverage it for generating target-consistent transitions based on the magnitude of the domain gap.
• Theoretical Analysis: A formal decomposition of the cross-domain performance gap into dynamics mismatch and generation error, providing a theoretical basis for balancing data mixing and generation.
• Dual Score-Based Architecture: A two-stage generative approach that decouples state coverage expansion from transition modeling. This design prevents overfitting to limited target data, which is a common failure mode when directly learning state-action-next-state mappings in low-data regimes.
• Empirical Validation: Extensive experiments demonstrating that TCE consistently outperforms state-of-the-art baselines across diverse cross-domain settings.

Experimental Results

The authors evaluated TCE on MuJoCo continuous control tasks (HalfCheetah, Hopper, Walker2d, Ant) and Adroit dexterous manipulation tasks, covering morphology, kinematic, and gravity shifts.

• Performance: TCE achieved the highest average performance in 33 out of 36 morphology-shift tasks in MuJoCo and outperformed all baselines on Adroit tasks.
• Robustness to Large Gaps: In scenarios with large dynamics gaps (e.g., significant morphology changes or extreme gravity shifts), the variant TCE(OG) (using only generated transitions, $\lambda_{mix}=0$) significantly outperformed methods that directly mixed source data. This confirms that naive source incorporation can be detrimental under large gaps.
• Small Gaps: In milder settings (e.g., gravity shifts), the variant TCE(SM) (incorporating source data, $\lambda_{mix} > 0$) showed superior performance, validating the framework's adaptability.
• Ablation Studies:
  • Coverage Coefficient ($\lambda_{cov}$): Performance follows a trade-off curve; expanding coverage too aggressively increases generation error and degrades performance. A moderate value (e.g., 0.2) was found optimal.
  • Two-Stage Design: The decoupled two-stage design was shown to be critical. A "one-stage" baseline that directly estimated next states from state-action pairs suffered from significantly higher errors and lower performance due to overfitting on limited target data.

Significance and Claims

The paper claims that TCE provides a simple and practical solution for cross-domain offline RL, particularly in safety-critical or cost-sensitive domains where online interaction is impossible. By theoretically grounding the decision to mix or generate data, TCE mitigates the distributional mismatch that plagues existing methods.

The authors emphasize that their approach is not a universal panacea; it assumes shared state and action spaces and relies on the target dataset, however small, to guide the generation. They acknowledge limitations in scenarios with extreme mismatch where source and target share little support. Nevertheless, the results suggest that target-aligned generation is a more robust strategy than direct data mixing when the domain gap is substantial, enabling effective policy learning with limited target data.