MOBODY: Model Based Off-Dynamics Offline Reinforcement Learning

1. Problem Definition

The paper addresses Off-Dynamics Offline Reinforcement Learning (RL). In this setting, an agent must learn a policy to maximize cumulative reward in a target domain ($M_{trg}$) using only pre-collected offline datasets from two sources:

1. Source Domain ($M_{src}$): A large dataset with dynamics that differ significantly from the target (e.g., a simulator with different gravity or friction).
2. Target Domain ($M_{trg}$): A very limited dataset (often orders of magnitude smaller than the source, e.g., a ratio of 200:1).

The Core Challenge:
Existing methods typically handle this by either:

• Reward Regularization: Penalizing source data based on the estimated dynamics gap.
• Data Filtering: Discarding source transitions that exhibit high dynamics shifts.

Limitation of Existing Methods: These approaches force the policy to optimize only on "low-shift" regions of the state space. Consequently, they fail to explore high-reward states in the target domain that lie outside these low-shift regions, leading to suboptimal performance when the dynamics shift is significant or when optimal trajectories require navigating high-shift regions.

2. Methodology: MOBODY

The authors propose MOBODY, a Model-Based approach that learns the target dynamics explicitly and uses them for policy exploration, rather than relying solely on static low-shift data.

A. Dynamics Learning with Shared Structure

MOBODY leverages the observation that while dynamics differ, the underlying structural knowledge (e.g., robot kinematics) is shared. To learn the target dynamics using limited target data and abundant source data, MOBODY employs a specific architecture:

• Shared State Encoder ($\phi_E$): Encodes states from both domains into a unified latent state representation ($z_s$).
• Separate Action Encoders ($\psi_{src}, \psi_{trg}$): Since achieving the same next state from the same current state requires different actions in different domains, MOBODY uses distinct encoders for source and target actions to map them into the shared latent space ($z_{sa}$).
• Unified Transition Function ($\phi_T$): A single transition function maps the unified latent state-action representation to the next state.

Loss Functions for Dynamics Learning:

1. Transition Loss: Minimizes the MSE between predicted and actual next states for both domains.
2. Encoder Loss: Ensures the latent state-action representation ($z_{sa}$) is close to the next state representation ($z_{s'}$), forcing the action encoders to capture domain-specific dynamics.
3. Cycle Transition Loss (VAE-style): Uses a Variational Autoencoder structure to learn robust state representations and prevent mode collapse, ensuring the transition function can reconstruct the current state when no action is taken.
4. Reward Learning: Learns a unified reward function $r(s, a, s')$ using both source and target data, penalized by model uncertainty (ensemble standard deviation) to ensure conservatism.

B. Policy Optimization with Target Q-Weighted Behavior Cloning

Once the target dynamics model is learned, MOBODY performs model-based offline RL:

1. Data Augmentation: It generates synthetic "rollout" data ($D_{fake}$) by simulating transitions using the learned target dynamics.
2. Enhanced Dataset: The policy is trained on a combination of:
   • Regularized source data (using DARA reward augmentation).
   • Limited real target data.
   • Synthetic rollout data.
3. Target Q-Weighted Behavior Cloning (BC):
   • Standard Behavior Cloning (BC) tends to imitate actions from the source dataset, which may be suboptimal in the target domain due to dynamics shifts.
   • MOBODY introduces a Target Q-weighted BC loss. Instead of uniformly imitating offline actions, it re-weights the BC loss based on the Target Q-value ($Q_{trg}$).
   • Mechanism: Actions with high target Q-values are up-weighted, guiding the policy toward actions that are expected to perform well under the target dynamics, rather than just mimicking the source behavior or maximizing source advantages.

3. Key Contributions

1. Novel Paradigm: Proposes the first model-based framework for off-dynamics offline RL that actively explores the target domain via learned dynamics, overcoming the limitations of conservative reward regularization and data filtering.
2. Architecture Design: Introduces a novel dynamics learning framework that shares state representations and transition functions across domains while using separate action encoders to account for domain-specific dynamics.
3. Regularization Strategy: Develops a simple yet effective Target Q-weighted Behavior Cloning loss that biases the policy toward high-value actions in the target domain, avoiding out-of-distribution (OOD) errors common in offline RL.
4. Comprehensive Evaluation: Extensive experiments on MuJoCo (Ant, HalfCheetah, Walker2d, Hopper) and Adroit (Pen, Door) benchmarks.

4. Experimental Results

The authors evaluated MOBODY against state-of-the-art baselines (DARA, BOSA, SRPO, REAG, MOPO, IQL, TD3-BC) under various shift types (Gravity, Friction, Kinematics, Morphology) and levels (0.1x to 5.0x).

• Overall Performance: MOBODY achieved an average 44% improvement over the best baseline (REAG) in gravity/friction shift settings and 25% improvement in kinematic/morphology shift settings.
• Large Shift Scenarios: MOBODY demonstrated superior robustness in "large shift" scenarios (e.g., gravity 0.1x or 5.0x) where existing methods (like DARA) often failed completely or performed poorly due to their inability to explore high-reward regions outside low-shift zones.
• Ablation Studies:
  • Removing the learned dynamics (replacing with target-only training) caused significant performance drops, validating the necessity of the shared-structure learning.
  • Removing the Target Q-weighted BC loss resulted in lower performance, confirming its role as an essential regularizer.
  • Model-based rollouts were identified as the primary driver of performance gains, with the Q-weighted BC acting as a crucial auxiliary component.

5. Significance

• Bridging the Gap: MOBODY effectively bridges the gap between source and target domains without requiring online interaction, making it highly applicable to real-world scenarios where safety or cost prevents extensive testing in the target environment.
• Exploration vs. Exploitation: Unlike previous methods that restrict exploration to "safe" (low-shift) regions, MOBODY enables safe exploration of the entire target state space by relying on a learned model, allowing the discovery of high-reward trajectories that were previously inaccessible.
• Generalization: The method shows strong generalization across different types of dynamics shifts (physical parameters vs. structural changes), suggesting a robust approach to domain adaptation in RL.

In conclusion, MOBODY represents a significant advancement in offline RL by shifting the focus from filtering mismatched data to learning and utilizing the target dynamics to guide policy exploration.