Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PROF: An LLM-based Reward Code Preference Optimization Framework for Offline Imitation Learning

Shengjie Sun, Jiafei Lyu, Runze Liu, Mengbei Yan, Bo Liu, Deheng Ye, Xiu Li

TL;DR

PROF introduces a fully automatic framework for reward design in offline imitation learning by generating executable reward function codes with LLMs and refining them via Reward Preference Ranking and TextGrad. It eliminates the need for environment interactions and RL training during reward design by using a data-driven dominance score computed from expert and noisy trajectories. Across MuJoCo, AntMaze, and Adroit tasks on the D4RL benchmark, PROF matches or surpasses recent baselines and improves several offline RL algorithms when combined with PROF-designed rewards. The approach yields human-readable reward functions and demonstrates robustness to different LLM APIs and task complexities, highlighting practical impact for safe and scalable offline IL. Overall, PROF shows that automated, interpretable reward design guided by expert preferences can substantially enhance offline policy learning.

Abstract

Offline imitation learning (offline IL) enables training effective policies without requiring explicit reward annotations. Recent approaches attempt to estimate rewards for unlabeled datasets using a small set of expert demonstrations. However, these methods often assume that the similarity between a trajectory and an expert demonstration is positively correlated with the reward, which oversimplifies the underlying reward structure. We propose PROF, a novel framework that leverages large language models (LLMs) to generate and improve executable reward function codes from natural language descriptions and a single expert trajectory. We propose Reward Preference Ranking (RPR), a novel reward function quality assessment and ranking strategy without requiring environment interactions or RL training. RPR calculates the dominance scores of the reward functions, where higher scores indicate better alignment with expert preferences. By alternating between RPR and text-based gradient optimization, PROF fully automates the selection and refinement of optimal reward functions for downstream policy learning. Empirical results on D4RL demonstrate that PROF surpasses or matches recent strong baselines across numerous datasets and domains, highlighting the effectiveness of our approach.

PROF: An LLM-based Reward Code Preference Optimization Framework for Offline Imitation Learning

TL;DR

PROF introduces a fully automatic framework for reward design in offline imitation learning by generating executable reward function codes with LLMs and refining them via Reward Preference Ranking and TextGrad. It eliminates the need for environment interactions and RL training during reward design by using a data-driven dominance score computed from expert and noisy trajectories. Across MuJoCo, AntMaze, and Adroit tasks on the D4RL benchmark, PROF matches or surpasses recent baselines and improves several offline RL algorithms when combined with PROF-designed rewards. The approach yields human-readable reward functions and demonstrates robustness to different LLM APIs and task complexities, highlighting practical impact for safe and scalable offline IL. Overall, PROF shows that automated, interpretable reward design guided by expert preferences can substantially enhance offline policy learning.

Abstract

Offline imitation learning (offline IL) enables training effective policies without requiring explicit reward annotations. Recent approaches attempt to estimate rewards for unlabeled datasets using a small set of expert demonstrations. However, these methods often assume that the similarity between a trajectory and an expert demonstration is positively correlated with the reward, which oversimplifies the underlying reward structure. We propose PROF, a novel framework that leverages large language models (LLMs) to generate and improve executable reward function codes from natural language descriptions and a single expert trajectory. We propose Reward Preference Ranking (RPR), a novel reward function quality assessment and ranking strategy without requiring environment interactions or RL training. RPR calculates the dominance scores of the reward functions, where higher scores indicate better alignment with expert preferences. By alternating between RPR and text-based gradient optimization, PROF fully automates the selection and refinement of optimal reward functions for downstream policy learning. Empirical results on D4RL demonstrate that PROF surpasses or matches recent strong baselines across numerous datasets and domains, highlighting the effectiveness of our approach.

Paper Structure

This paper contains 39 sections, 14 equations, 2 figures, 12 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Offline Imitation Learning.
  4. Reward Design via Large Models.
  5. Preliminaries
  6. Reinforcement Learning.
  7. Text Gradient.
  8. Method
  9. Reward Generation
  10. Preference Ranking
  11. Expert Demonstration Return Threshold.
  12. Noisy Trajectory Construction.
  13. Dominance Score.
  14. Reward Preference Ranking.
  15. Iterative Optimization
  16. ...and 24 more sections

Figures (2)

  • Figure 1: The framework of PROF. PROF initiates by generating $n$ candidate reward functions, which are stored in a buffer $\mathcal{B}$. The algorithm proceeds through $T$ rounds of iterative optimization. In each round, the reward functions with the highest and lowest dominance scores are selected from the buffer to construct the loss feedback. Leveraging TextGrad, gradients are computed automatically and backpropagation is applied to optimize each candidate independently, yielding $n$ new reward functions. These newly optimized candidates are added to the buffer, ensuring diversity and continual improvement. After $T$ iterations, PROF outputs the reward function with the highest dominance score.
  • Figure 2: The performance of PROF across different numbers of iterative optimization rounds. From left to right are MuJoCo and AntMaze domains. We report the total D4RL normalized score calculated across $5$ seeds for each domain.