Progress Reward Model for Reinforcement Learning via Large Language Models

AI summary

Overview

Research area: Reinforcement learning (RL) augmented by large language models (LLMs), specifically reward shaping and hierarchical task decomposition for robotics control.

Technical level: Advanced. The paper assumes familiarity with Markov Decision Processes, Q-functions, potential-based reward shaping, and policy-gradient/actor-critic algorithms (SAC, PPO).

Scope: The paper proposes PRM4RL, a framework that uses an LLM once to decompose a long-horizon robotics task into subtasks and to write a code-based progress function, then converts that progress function into a dense, theoretically grounded reward for training a low-level RL policy.

What This Paper Is About

Reinforcement learning struggles on long, multi-stage robot tasks where the only reward arrives at the very end, because the agent gets almost no feedback to learn from. LLMs can supply useful prior knowledge, but prior work either uses the LLM only as a high-level planner (which does not train low-level control well) or only as a reward writer (which requires extensive human feedback or expensive search over reward code). This paper combines both: the LLM breaks the task into subtasks and writes a per-subtask progress score, and that score is turned into a dense reward with formal guarantees that it does not change the optimal policy.

Key Contributions

1. A combined planning-and-reward framework (PRM4RL). Unlike prior LLM-augmented RL methods that address either high-level planning or low-level reward generation, PRM4RL integrates both in a single pipeline: subtask decomposition plus subtask priors at the high level, and dense reward at the low level.

2. A subtask determination function Ψ(s) that avoids repeated LLM calls. Rather than querying the LLM at every timestep to replan or to check subtask completion (as prior work does), the LLM generates a low-cost logical function that maps the current state directly to the index of the active subtask in plan_list.

3. A Progress Reward Model (PRM) built from a progress function via potential-based reward shaping. The reward is R_t^PRM = γ·Φ(s_{t+1}) − Φ(s_t) + I(s_{t+1})·r_bonus, where Φ is the LLM-written progress function. The paper proves policy invariance (Theorem 4.1), a convergence guarantee (Corollary 4.2), and a convergence-efficiency result (Theorem 4.3) showing equivalence to a Q-value initialization offset by Φ(s).

4. A "verb + noun" subtask-prior format embedded with SimCSE. Subtask priors are constrained to a fixed linguistic pattern (e.g., "Move to object", "Grasp object") and embedded into the state space, which the authors credit for improved generalization to unseen tasks.

Main Findings

• Outperforms LLM-planner and LLM-rewarder baselines: Across nearly all MetaWorld and Maniskill environments, PRM4RL converges faster and reaches higher success rates than ELLM (an LLM-planner baseline) and T2R/Text2Reward (an LLM-rewarder baseline), as well as the expert-written Oracle dense reward.

• Large margins on hard long-horizon tasks: On MW-pick-out-of-hole and MS-LiftPegUpright, the baseline methods achieve less than 15% success while PRM4RL achieves nearly 100%.

• Faster convergence than the planner baseline: Compared to ELLM, the improvement in convergence speed is attributed to the PRM providing dense rewards with theoretical guarantees for low-level control.

• Subtask decomposition helps the LLM write better reward code: Compared to T2R, decomposition reduces the difficulty of understanding environment dynamics and designing useful metrics, letting the method solve complex tasks that T2R fails on.

• Removing the PRM reward hurts badly: The "our w.o. prm" ablation (PRM reward replaced by the oracle reward, subtask prior retained) shows a significant performance drop; in most cases its curves resemble the oracle, and adding the subtask prior sometimes even caused a slight performance decrease.

• Removing the subtask prior slows convergence: "Our w.o. prior" converges more slowly, yet still outperforms T2R, which the authors read as evidence that the PRM reward alone is effective.

• Using progress directly as reward fails: "Our w.o. potential reward" (progress function used directly as the reward, prior retained) shows a significant drop, with success rate rising then declining while episode reward continues to increase — high reward despite task failure, leading the policy in an undesirable direction.

• Better generalization to unseen MetaWorld tasks: Trained policies were evaluated on unseen tasks (e.g., original task Metaworld-faucet-open tested on drawer-close, button-press-wall, button-press, coffee-button; separately, Metaworld-pick-place tested on drawer-close, pick-place-wall, sweep-into, handle-press). PRM4RL shows the least performance loss; "our w.o. prior" drops significantly, which the authors attribute to the "verb + noun" prior capturing cross-task commonalities.

• Reduced LLM calls: Figure 1(d) reports average results on MetaWorld and Maniskill showing improved success rate alongside reduced LLM calls relative to existing methods.

Methodology in Plain English

The pipeline has three stages, and the LLM is queried essentially once.

1. Ask the LLM for a plan plus code. The authors give GPT-4o a "Pythonic" prompt containing the state space, an environment description, and the task description. Following the PAL idea, the LLM writes its reasoning as Python comments and its actual outputs as executable Python. The output includes a plan_list of subtasks, each phrased as "verb + noun".

2. Ask the same LLM for two functions. The first is a determination function Ψ(s) that looks at the current state and returns which subtask index the robot is currently in, based on completion conditions and relevant state features the LLM identifies. The second is a progress function Φ(s) that combines fine-grained per-subtask progress indicators into a single overall progress score.

3. Convert progress into reward and train. Instead of rewarding progress directly, the framework takes the change in progress between consecutive states, discounted: γ·Φ(s_{t+1}) − Φ(s_t), plus a bonus when the success condition is met. This is potential-based reward shaping, so the optimal policy is unchanged. The current subtask prior is embedded with SimCSE and appended to the state, giving an augmented MDP that is optimized with off-the-shelf SAC or PPO from Stable-Baselines3.

The reason for using the difference rather than the raw progress score is demonstrated empirically: raw progress as reward lets the agent collect high reward while failing the task.

Why This Matters

Impact on research. The paper offers a principled bridge between two previously separate uses of LLMs in RL — planning and reward design — and backs the reward component with policy-invariance and convergence proofs rather than treating shaped rewards as a heuristic. It also removes a practical cost of prior LLM-planner methods: repeated LLM calls during execution, replaced by a generated state-based classifier.

Real-world applications (as suggested by the robotics focus):

• Robot manipulation in tabletop settings (the MetaWorld benchmark uses a Sawyer robot with drawers, windows, faucets, doors, and balls).
• Multi-stage assembly or pick-and-place in warehouses, where success is only observed at the end of a long sequence.
• Household service robots performing compound instructions ("open the drawer and place the object inside").
• Any sparse-reward control pipeline where an engineer wants dense guidance without hand-engineering a reward function per task.

Industry relevance. The approach requires no expert demonstrations and no large-scale human reward annotation — only an LLM call and an environment description — which lowers the cost of setting up RL training for new tasks. Training cost is reported as roughly 1 hour per task (5 seeds simultaneously) on MetaWorld and 1 to 10 hours on ManiSkill on a 4× RTX-3090, 128-core, 256 GiB machine. Code is released at github.com/deng-ai-lab/PRM4RL.

Future Directions

• Mitigating LLM hallucination. The authors list hallucination during reasoning as an explicit limitation, since the entire plan, determination function, and progress function come from a single LLM generation.
• Extending beyond robotics. The paper states that applicability to domains outside robotics remains unexplored and that further optimization and adaptation will be needed for diverse environments and tasks.
• Reducing reliance on state-feature identification. Ψ(s) depends on the LLM correctly locating relevant features in the state space from the prompt; how robust this is to high-dimensional or image-based observations is not reported.
• Comparing LLM backbones and prompt sensitivity. All LLM-related components use GPT-4o with a sampling temperature of 0.7; the paper does not report how performance varies across models, prompts, or repeated generations.

Target Audience

Researchers and graduate students working on LLM-augmented reinforcement learning, reward shaping, and hierarchical RL; robotics engineers who need dense reward signals for sparse long-horizon manipulation tasks without expert demonstrations; and practitioners already using LLM planners (SayCan-style) or LLM reward generators (Text2Reward/Eureka-style) who want a single pipeline that covers both levels. Readers without a background in MDPs and potential-based reward shaping will find the theory section (Section 4.2 and Appendix A) demanding, but the pipeline description and ablation results are accessible on their own.

Authors’ abstract

Traditional reinforcement learning (RL) algorithms face significant limitations in handling long-term tasks with sparse rewards. Recent advancements have leveraged large language models (LLMs) to enhance RL by utilizing their world knowledge for task planning and reward generation. However, planning-based approaches often depend on pre-defined skill libraries and fail to optimize low-level control policies, while reward-based methods require extensive human feedback or exhaustive searching due to the complexity of tasks. In this paper, we propose the Progress Reward Model for RL (PRM4RL), a novel framework that integrates task planning and dense reward to enhance RL. For high-level planning, a complex task is decomposed into a series of simple manageable subtasks, with a subtask-oriented, fine-grained progress function designed to monitor task execution progress. For low-level reward generation, inspired by potential-based reward shaping, we use the progress function to construct a Progress Reward Model (PRM), providing theoretically grounded optimality and convergence guarantees, thereby enabling effective policy optimization. Experimental results on robotics control tasks demonstrate that our approach outperforms both LLM-based planning and reward methods, achieving state-of-the-art performance. Code is available at https://github.com/deng-ai-lab/PRM4RL.