Technical Summary: Unifying Value Alignment and Assignment in Cross-Domain Offline Reinforcement Learning with Heterogeneous Datasets

1. Problem Statement

Cross-domain offline Reinforcement Learning (RL) aims to learn a policy in a target domain using a limited target dataset and an abundant source dataset that exhibits a dynamics shift (i.e., different transition dynamics $P_{src}$ vs. $P_{tar}$). While recent methods have attempted to address the Out-of-Distribution (OOD) dynamics issue through data filtering based on dynamics alignment (e.g., IGDF, OTDF) or value alignment (e.g., DVDF), these approaches typically assume the source dataset is collected from a single environment with a single behavior policy.

This paper addresses a more general and realistic setting: Heterogeneous Cross-Domain Offline RL. In this setting, the source dataset is aggregated from multiple source domains with distinct transition dynamics and collected by diverse behavior policies. The authors identify a critical, previously overlooked issue in this setting called Value Misassignment.

Value Misassignment occurs when a single value function (or advantage function) is trained on the entire heterogeneous source dataset without distinguishing between the underlying dynamics modalities. This leads to:

1. Incorrect Advantage Estimation: The advantage values assigned to samples do not reflect their true optimality relative to their specific dynamics sub-domain.
2. Suboptimal Data Filtering: Data filtering mechanisms relying on these misassigned values are misled, selecting suboptimal samples (e.g., medium-quality data from a specific dynamic) while discarding optimal ones.
3. Loosened Suboptimality Gap: Theoretical analysis shows that value misassignment prevents the tight bounding of the suboptimality gap in the target domain, degrading agent performance.

2. Methodology: V2A

To resolve value misassignment, the authors propose V2A (Value Alignment and Assignment), a framework that unifies dynamics alignment, value alignment, and value assignment. The method proceeds in three main stages:

2.1. Temporally-Consistent Modality Representation Learning

The first step is to disentangle the different dynamics modalities within the source dataset.

• Approach: The authors employ an Expectation-Maximization (EM) framework to learn a latent representation $z$ for each trajectory.
• Innovation: Unlike standard variational autoencoders that encode transitions individually (leading to temporal inconsistency), V2A uses a Temporally-Consistent ELBO. It encodes the entire trajectory $\tau$ into a single latent variable $z$ using a Recurrent Neural Network (RNN/LSTM), while decoding transitions at the step level.
• Goal: This ensures that all transitions within a trajectory share the same dynamics representation $z$, effectively clustering trajectories by their underlying dynamics.

2.2. Modality-Aware Advantage Learning

Once the dynamics modalities ($z$) are inferred, the source dataset is relabeled as $D'_{src} = \{(s, a, r, s', z)\}$.

• Approach: The authors propose Modality-Aware Advantage Learning. Instead of training a standard advantage function $A(s, a)$, they train a modality-conditioned advantage function $A(s, a, z)$.
• Mechanism: The latent representation $z$ is fed as an input to the $Q$-function and $V$-function during pre-training (using a Sparse-QL scheme).
• Result: This allows the advantage function to distinguish between different dynamics modalities, assigning correct advantage values to state-action pairs based on their specific dynamics context, thereby correcting the value misassignment.

2.3. Unified Data Filtering and Policy Learning

The final stage integrates these components into a data filtering paradigm for target policy learning.

• Score Function: A new score function $f(s, a, s', z)$ is defined to filter source data:
  $$f(s, a, s', z) = \lambda \cdot h(s, a, s') + (1 - \lambda) \cdot \text{Norm}(A(s, a, z))$$
  where $h(\cdot)$ measures dynamics alignment (via contrastive learning or optimal transport) and $A(s, a, z)$ measures the modality-aware value alignment.
• Filtering: Source samples are selected based on the top-$\xi$ quantile of this score.
• Policy Optimization: The target policy is trained on the filtered source data combined with the limited target data using Implicit Q-Learning (IQL).

3. Key Contributions

1. Problem Formulation: The paper formalizes the Heterogeneous Cross-Domain Offline RL setting, where source data comes from multiple domains and behavior policies, and identifies Value Misassignment as a fundamental limitation in existing methods.
2. Theoretical Insight: The authors provide a theoretical bound on the suboptimality gap, demonstrating that value misassignment loosens this bound. They prove that disentangling dynamics modalities is necessary to achieve a tighter bound.
3. Algorithm (V2A): They propose V2A, which integrates:
   • Temporally-consistent modality representation learning (via EM).
   • Modality-aware advantage learning.
   • A unified data filtering paradigm.
4. Empirical Validation: Extensive experiments on MuJoCo tasks (HalfCheetah, Hopper, Walker2d, Ant) with heterogeneous source datasets show that V2A significantly outperforms strong baselines (IGDF, OTDF, DVDF, BOSA, DARA).

4. Experimental Results

• Performance: V2A achieves significant performance improvements over baselines. For instance, when built upon IGDF, V2A improved the total normalized score by 21.4% compared to the baseline IGDF, and by 21.4% over DVDF (which only improved by 6.8% over IGDF). Similar gains were observed with OTDF.
• Ablation Studies:
  • Temporally-Consistent ELBO: Removing the trajectory-level consistency (using sample-level ELBO instead) resulted in lower performance, confirming the necessity of temporal consistency for accurate modality extraction.
  • Modality Awareness: Visualizations of the learned advantage distributions show that V2A produces sharper, more distinct advantage distributions compared to the flatter, less discriminative distributions of DVDF, indicating better estimation of sample optimality.
  • Parameter Sensitivity: The method is robust to the trade-off coefficient $\lambda$ (optimal around 0.6) and data selection ratio $\xi$ (optimal around 0.5).

5. Significance and Claims

The paper claims that V2A is the first framework to effectively address the value misassignment problem inherent in heterogeneous cross-domain offline RL. By unifying dynamics alignment, value alignment, and value assignment, V2A enables efficient policy transfer even when source data is a complex mixture of diverse dynamics and behaviors.

The authors emphasize that while previous works focused on dynamics or value alignment in isolation (or assumed homogeneous source data), their work highlights that ignoring the multi-modal nature of source dynamics leads to incorrect value estimation, which in turn undermines data filtering. The proposed method offers a simple yet effective solution that significantly bridges the gap between theoretical settings and realistic, heterogeneous data scenarios in robotics and embodied AI.

Limitations Acknowledged:

• The temporally-consistent modality representation learning requires additional training time (though this can be precomputed).
• The theoretical analysis relies on mild assumptions (e.g., second-order sufficient conditions for convexity in the neighborhood of the optimal policy).