Translate Policy to Language: Flow Matching Generated Rewards for LLM Explanations

1. Problem Statement

As intelligent agents (RL agents, LLMs, robots) become more integrated into human environments, the ability to explain their decision-making processes in natural language is critical for trust, safety, and alignment. However, existing methods face two primary challenges:

• Task Specificity: Most explanation methods are tailored to specific domains (e.g., self-driving or stock prediction), lacking generalizability.
• Reward Modeling Limitations: Training explanation-generating LLMs via Reinforcement Learning from AI Feedback (RLAIF) typically relies on "proxy" LLMs to generate rewards. These proxy rewards are often noisy and fail to capture the pluralistic and probabilistic nature of human judgment. Furthermore, standard RLAIF approaches struggle to explicitly manage the deviation between proxy rewards and true human reward distributions, leading to suboptimal explanations that may lack logical soundness or fail to enable third parties to infer the agent's actual decision.

2. Methodology

The authors propose a general-purpose framework that trains an EXPLANATION LLM to recover an underlying agent policy and generate natural language explanations using Reinforcement Learning from AI Feedback (RLAIF), where the reward model is a Generative Continuous Normalizing Flow (CNF), specifically a Rectified Flow (RF).

Core Components:

1. The Goal: Train an LLM ($\pi_e$) to generate an explanation ($y$) given only a decision context ($c$), such that a third party can infer the agent's true decision ($a$) from $y$. The actual decision $a$ is hidden from the explanation LLM during generation to prevent superficial restatement.
2. Proxy LLMs: Instead of costly human feedback, the system uses multiple independent PROXY LLMs to generate initial reward signals. These LLMs are queried with the context and the generated explanation to predict the likelihood of the true decision.
3. Rectified Flow Reward Model ($\phi$):
   • Architecture: The RF model is embedded within an LLM backbone. It takes the context and explanation as input and uses a cross-attention mechanism to selectively attend to linguistic cues. It outputs a vector field that transforms a base noise distribution (Gaussian) into a complex reward distribution over possible decisions.
   • Training: The RF model is trained to minimize the Wasserstein distance between the distribution of logits generated by the PROXY LLMs and the flow's output. It learns to "denoise" the proxy rewards, effectively recovering the underlying human reward distribution.
   • Theoretical Guarantee: The paper proves that if the noise in the proxy LLM rewards follows the same functional form as the flow's base distribution (e.g., Gaussian), the RF model provably bounds the deviation from the true human reward distribution. This allows the model to correct for proxy errors.
4. Training Loop:
   • Iterative Process: The EXPLANATION LLM and the RF reward model are trained alternately.
   • Reward Signal: The RF generates a distribution over decisions. The reward is derived from the change in the probability (logit) of the true decision as sentences are added to the explanation (sentence-level rewards), providing dense feedback.
   • Optimization: The EXPLANATION LLM is fine-tuned using PPO (Proximal Policy Optimization) with LoRA to maximize the rewards generated by the RF model.

3. Key Contributions

• Flow-Matching for RLAIF: This is the first work to utilize generative flow matching (specifically Rectified Flow) as a reward model for training LLMs to generate explanations. It addresses the limitation of deterministic reward models by modeling the distribution of human preferences.
• Theoretical Error Bounds: The authors provide a theoretical proof showing that the Rectified Flow model can bound the error relative to the true human reward distribution, provided the noise model matches the flow's base distribution. This offers a rigorous justification for using generative models to correct proxy LLM errors.
• Generalizable Framework: The method is task-agnostic. It does not require the explanation LLM to know the specific agent architecture or task details, relying solely on contextual information to recover the policy.
• Denoising Proxy Rewards: The framework effectively distinguishes between "true negatives" (bad explanations) and "false negatives" (good explanations misclassified by proxy LLMs), assigning appropriate rewards to both.

4. Experimental Results

The method was evaluated on four diverse benchmarks: SMAC (Multi-agent RL), AI2-THOR (Embodied AI), MMLU (Professional Knowledge QA), and MathQA (Math Word Problems).

• Performance vs. Baselines: The proposed method outperformed state-of-the-art baselines, including SFT, RLHF (PPO, DPO, KTO), RLAIF (Skywork), and reasoning frameworks (CoT+SFT, ReFT).
  • Accuracy Gains: It achieved 6.9% to 12.4% higher accuracy in inferring true decisions compared to using raw Proxy LLM rewards.
  • Overall Improvement: It outperformed all baselines by 1.9% to 18.2% across all tasks and backbone models (Llama-3.1, Qwen, Gemma).
• Human Evaluation: Human evaluators rated the generated explanations as having:
  • Highest Accuracy: 25.7% higher than the strongest baseline (DPO) in predicting the correct decision.
  • Logical Soundness & Actionability: Significantly higher scores for logical coherence and practical utility.
  • Lower Cognitive Load: Easier for humans to understand and process.
• Ablation Studies:
  • Removing the Rectified Flow (using raw Proxy LLM rewards) caused a significant performance drop.
  • Removing the cross-attention mechanism in the RF model led to stagnation on negative samples, proving the importance of the architecture in generalizing to unseen data.
  • Sentence-level rewards provided better guidance than sparse (single-token) rewards.

5. Significance

This paper represents a significant advancement in Explainable AI (XAI) and RLHF/RLAIF:

• Bridging the Gap: It bridges the gap between the probabilistic nature of human judgment and the deterministic nature of standard reward models. By treating rewards as distributions, it captures the pluralism inherent in human feedback.
• Robustness: The theoretical and empirical demonstration that generative flows can "denoise" proxy rewards suggests a new paradigm for training LLMs where high-quality human data is scarce.
• Scalability: The method offers a scalable, generalizable solution for explaining complex agent behaviors (from robotics to LLM reasoning) without requiring task-specific engineering or massive human annotation efforts.
• Trust and Safety: By generating explanations that are logically sound and allow third parties to accurately infer agent decisions, the work directly contributes to building more transparent and trustworthy AI systems.