Learning Task-Invariant Properties via Dreamer: Enabling Efficient Policy Transfer for Quadruped Robots

1. Problem Statement

The core challenge addressed is the Sim-to-Real transfer gap in quadruped robot locomotion. While training in simulation is cost-effective, policies often fail in the real world due to:

• Dynamics Discrepancies: Differences in friction, mass, and sensor noise between simulation and reality.
• Over-reliance on Specific Dynamics: Existing methods (like standard Dreamer or World Model Perception) learn representations heavily tied to specific simulation parameters, making them fragile when faced with real-world variations.
• Data Inefficiency & Catastrophic Forgetting: Fine-tuning on limited real-world data often leads to "representation collapse" (loss of general features) or "catastrophic forgetting" (loss of previously learned skills).
• Manual Feature Design: Identifying robust, transferable features manually is costly, biased, and requires deep domain expertise.

2. Methodology: The DreamTIP Framework

The authors propose DreamTIP (Learning Task-Invariant Properties via Dreamer), a two-stage framework that integrates Large Language Models (LLMs) with a World Model architecture.

A. Task-Invariant Properties (TIP) Extraction via LLM

Instead of manual design, the framework uses an LLM to identify Task-Invariant Properties—high-level physical or behavioral constraints that are crucial for task success but robust to dynamic changes (e.g., contact stability, terrain clearance).

• Process: The LLM takes high-level task descriptions ($I_{text}$) and privileged state observations ($I_{priv}$) as input.
• Output: It generates a transformation function (TIP_Extractor) that maps raw privileged states into a feature vector ($f_t$) representing these invariant properties.
• Examples:
  • Base Angle: Pitch and roll derived from projected gravity.
  • Contact Stability: Statistical measures (mean/variance) of foot contact forces.
  • Terrain Clearance: The difference between body height and front depth image data.

B. Enhanced World Model (DreamTIP)

The framework builds upon the Dreamer architecture (Recurrent State-Space Model - RSSM).

• Auxiliary Prediction Target: A new predictor (MLP) is added to the Dreamer model. It takes the latent state ($h_t, z_t$) and attempts to predict the Task-Invariant Properties ($\hat{f}_t$).
• Loss Function: The training objective includes a standard reconstruction loss plus a Maximum Likelihood Estimation (MLE) term for the TIPs:
  $$\mathcal{L}_{train} = \mathcal{L}_{D} - \mathbb{E}[\sum \ln p_\theta(f_t | h_t, z_t)]$$
  This forces the latent representation to encode dynamics-robust features, reducing reliance on specific simulation parameters.

C. Efficient Real-World Adaptation

To bridge the final gap with minimal real-world data, an efficient adaptation strategy is employed:

1. Mixed Replay Buffer: Combines simulated data (Sim buffer) with a small amount of real-world rollouts (Mix buffer) to prevent distribution shift.
2. Frozen Reference Model: A pre-trained DreamTIP model ($M_{sg}$) is frozen to serve as a "teacher."
3. Regularization Constraints: During adaptation, the model minimizes the Negative Cosine Similarity between the stochastic state representations of the adaptable model and the frozen reference model.
   • Goal: Enforce semantic consistency and prevent representation collapse without being sensitive to feature scale.
4. Module Freezing: The Recurrent Model (RNN) is frozen during fine-tuning to accelerate alignment with real dynamics while preserving temporal memory.

3. Key Contributions

1. LLM-Driven Feature Discovery: A novel approach using LLMs to automatically extract Task-Invariant Properties (e.g., contact stability, clearance) that serve as robust auxiliary targets for world model learning, eliminating manual feature engineering.
2. DreamTIP Architecture: An extension of the Dreamer framework that explicitly learns these invariant properties, creating representations that generalize across tasks and dynamic variations.
3. Stable Adaptation Strategy: A novel fine-tuning protocol using a mixed replay buffer, frozen reference models, and cosine similarity regularization to achieve rapid adaptation with minimal real data while avoiding catastrophic forgetting.
4. Comprehensive Validation: Extensive testing on complex terrains (Stairs, Climbing, Tilt, Crawling) demonstrating superior performance over state-of-the-art baselines.

4. Experimental Results

Experiments were conducted on Unitree Go2 robots in both Isaac Gym simulation and real-world environments.

• Simulation Performance:

  • Achieved an average performance improvement of 28.1% across eight distinct transfer tasks compared to baselines (WMP, DreamTIP-DWL).
  • In the challenging "Crawl" task (23 cm gap), DreamTIP achieved a reward of 25.35, significantly outperforming the baseline (WMP) which scored 5.66 (an 83% performance drop for the baseline vs. 30% for DreamTIP as difficulty increased).

• Real-World Performance:

  • Climb Task (52 cm): The baseline (WMP) achieved only a 10% success rate. DreamTIP achieved a 100% success rate.
  • Tilt Task (33 cm): DreamTIP achieved 80% success vs. 40% for the baseline.
  • Robustness: Even with a 2kg counterweight added to shift the center of mass, DreamTIP maintained high success rates, whereas baselines failed.

• Ablation Studies:

  • LLM vs. Manual/Privileged: LLM-designed TIPs outperformed direct privileged information prediction (DreamTIP-DWL).
  • Adaptation: Removing the adaptation mechanism or the regularization constraints led to significant performance drops, confirming the necessity of the mixed buffer and stability constraints.
  • Data Efficiency: The method achieved optimal results with as few as 5 real-world trajectories for fine-tuning.

5. Significance

This work represents a significant step forward in Embodied AI and Robotics:

• Bridging the Sim-to-Real Gap: It offers a scalable solution to the "reality gap" without requiring expensive high-fidelity simulators or massive real-world datasets.
• Generalization: By focusing on invariant physical properties rather than specific dynamics, the policy learns to "understand" the task physics rather than just memorizing simulation parameters.
• Automation: Leveraging LLMs to design reward structures and feature extractors democratizes the development of robust robot controllers, reducing the need for expert manual tuning.
• Practical Deployment: The ability to deploy robust policies on real quadrupeds with minimal fine-tuning makes this approach highly viable for real-world applications in search and rescue, exploration, and logistics.