Preference Goal Tuning: Post-Training as Latent Control for Frozen Policies

1. Problem Statement

Goal-conditioned policies (e.g., foundation models for robotics or embodied AI) have demonstrated strong capabilities in interpreting instructions and executing diverse behaviors. However, their downstream performance is highly sensitive to the specific prompts or instructions used.

• The Dilemma:
  • Prompt Engineering: Relies on discrete, derivative-free text search. It respects the frozen policy but often fails to elicit optimal behavior due to the limitations of natural language as a control interface.
  • Standard Fine-Tuning: Updates policy parameters to align with tasks. While effective, it risks catastrophic forgetting, overfitting to narrow task distributions, and destroying the broad generalization capabilities acquired during pre-training.
• The Gap: There is a need for an adaptation mechanism that offers the optimization power of fine-tuning while maintaining the structural stability and generalization of a frozen foundation model.

2. Methodology: Preference Goal Tuning (PGT)

The authors propose Preference Goal Tuning (PGT), a post-training framework that reformulates adaptation as a latent control problem. Instead of updating the policy weights, PGT optimizes a continuous latent goal embedding ($g$) to modulate the behavior of a frozen policy.

Core Formulation

• Frozen Policy: The policy parameters ($\theta$) remain fixed: $\pi(a|s, g; \theta)$.
• Latent Control: The goal $g \in \mathbb{R}^d$ is treated as a learnable parameter. The policy defines a "behavioral manifold" indexed by $g$.
• Objective: Find the optimal $g^*$ that induces a trajectory distribution aligned with task preferences, minimizing divergence from a target distribution defined by human or reward-based preferences.

Algorithmic Process

PGT operates in two main phases, often iteratively:

1. Preference Sample Generation:
   • Start with an initial prompt encoded into a latent goal $g$.
   • Collect trajectories ($\tau$) by interacting with the environment using the frozen policy and $g$.
   • Label trajectories as positive (preferred) or negative (non-preferred). Labels can come from human annotators or reward signals (e.g., cumulative reward in Minecraft).
2. Behavior Preference Propagation:
   • Update the latent goal $g$ using a preference learning objective (e.g., DPO - Direct Preference Optimization).
   • The objective maximizes the likelihood of preferred trajectories relative to a reference goal ($g_{ref}$) while suppressing undesirable ones.
   • Loss Function: The optimization minimizes the negative log-likelihood of the preference pairs:
     $$\mathcal{L}(g, g_{ref}) = \mathbb{E}_{(\tau_w, \tau_l) \sim D} \left[ -\log \sigma \left( \beta \Delta \right) \right]$$
     Where $\Delta$ is the difference in log-likelihood ratios between the current goal and the reference goal for the winning ($\tau_w$) and losing ($\tau_l$) trajectories.

Key Theoretical Insight

By restricting optimization to a low-dimensional latent vector rather than millions of parameters, PGT imposes a strong inductive bias that favors generalization over memorization. It effectively searches for the optimal conditioning input without altering the underlying physical dynamics encoded in the policy.

3. Key Contributions

1. Latent Control Paradigm: Introduces a novel post-training approach where adaptation is achieved solely by optimizing the latent goal embedding, keeping the foundation policy frozen. This decouples task alignment from physical dynamics.
2. Preference-Based Optimization: Adapts preference learning (specifically DPO and variants) to the latent goal space, allowing for efficient alignment with task preferences using minimal data.
3. Robustness to OOD: Demonstrates that isolating task adaptation to the latent goal prevents catastrophic forgetting and maintains out-of-distribution (OOD) generalization, which standard fine-tuning often degrades.
4. Continual Learning Capability: Shows that PGT enables scalable continual learning by storing a compact latent vector per task, avoiding task interference without complex replay or regularization mechanisms.

4. Experimental Results

The framework was evaluated on the Minecraft SkillForge benchmark (17 tasks) using two foundation policies (GROOT and STEVE-1) and on the LIBERO-goal robotic manipulation benchmark.

• Performance Gains:
  • PGT achieved average relative improvements of 72.0% (GROOT) and 81.6% (STEVE-1) over pre-trained policies on in-distribution tasks.
  • It consistently outperformed expert-crafted prompts and standard prompt engineering baselines.
• Out-of-Distribution (OOD) Generalization:
  • In OOD settings (changed seeds, biomes, tool configurations), PGT improved performance by 73.8% and 36.9% respectively.
  • Crucially, PGT surpassed full fine-tuning by 13.4% in OOD settings. Full fine-tuning improved in-distribution performance but often degraded generalization due to overfitting.
• Parameter Efficiency:
  • PGT optimizes only 512 parameters (the latent goal vector), whereas full fine-tuning involves ~100M parameters.
  • It outperformed other Parameter-Efficient Fine-Tuning (PEFT) methods like LoRA, BitFit, and VeRA in OOD scenarios.
• Continual Learning:
  • In sequential task acquisition, PGT avoided catastrophic forgetting entirely, whereas baselines like Naive Continual Learning (NCL) and Elastic Weight Consolidation (EWC) suffered significant performance drops on earlier tasks.
• Cross-Domain Validity:
  • Experiments on the OpenVLA policy for robotic manipulation (LIBERO-goal) confirmed that PGT is a general post-training mechanism applicable beyond Minecraft, improving success rates in continuous control domains.

5. Significance and Impact

• Preservation of Foundation Capabilities: PGT offers a "non-destructive" way to adapt large models. It allows the model to learn new tasks without erasing the broad, compositional knowledge gained during pre-training.
• Sample Efficiency: By optimizing a low-dimensional vector, PGT requires significantly less data and compute than full fine-tuning, making it highly practical for real-world deployment.
• Scalable Adaptation: The ability to store a single compact latent vector per task suggests a viable path for scalable, multi-task agents that do not suffer from interference or the need for massive replay buffers.
• Safety and Control: The framework provides a continuous, differentiable interface for behavior modulation, bridging the gap between discrete text prompts and the continuous control required for complex physical tasks.

In conclusion, Preference Goal Tuning represents a significant shift in how foundation models are adapted, moving from weight modification to latent space navigation, thereby achieving superior robustness, generalization, and efficiency in goal-conditioned decision-making.