LLM-Driven Policy Diffusion: Enhancing Generalization in Offline Reinforcement Learning

2025-08-30

AI Summary

To address the limited generalization capability of policies in offline reinforcement learning (RL), this paper proposes a task-trajectory dual-prompted generalization framework. Methodologically, it introduces the first integration of large language models (LLMs) with trajectory prompting in offline RL: an LLM parses natural-language task descriptions to extract semantic priors, a Transformer encodes historical trajectories to model behavioral patterns, and a conditional diffusion model generates context-aware generalized policies. Crucially, the framework operates entirely without online interaction, enabling effective adaptation to unseen tasks using only static offline datasets. It achieves significant improvements over existing state-of-the-art methods across multiple cross-task benchmarks, demonstrating strong generalization and environmental robustness. The core contribution is the establishment of the first LLM-driven dual-prompt diffusion policy learning paradigm, offering a novel conceptual and technical pathway for generalization in offline RL.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Natural Language Processing: (Large) Language Models
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Reinforcement Learning (RL) is known for its strong decision-making capabilities and has been widely applied in various real-world scenarios. However, with the increasing availability of offline datasets and the lack of well-designed online environments from human experts, the challenge of generalization in offline RL has become more prominent. Due to the limitations of offline data, RL agents trained solely on collected experiences often struggle to generalize to new tasks or environments. To address this challenge, we propose LLM-Driven Policy Diffusion (LLMDPD), a novel approach that enhances generalization in offline RL using task-specific prompts. Our method incorporates both text-based task descriptions and trajectory prompts to guide policy learning. We leverage a large language model (LLM) to process text-based prompts, utilizing its natural language understanding and extensive knowledge base to provide rich task-relevant context. Simultaneously, we encode trajectory prompts using a transformer model, capturing structured behavioral patterns within the underlying transition dynamics. These prompts serve as conditional inputs to a context-aware policy-level diffusion model, enabling the RL agent to generalize effectively to unseen tasks. Our experimental results demonstrate that LLMDPD outperforms state-of-the-art offline RL methods on unseen tasks, highlighting its effectiveness in improving generalization and adaptability in diverse settings.

Problem

• Enhancing generalization in offline reinforcement learning
• Overcoming limitations of offline data for new tasks
• Improving RL agent adaptability to unseen environments

Innovation

• LLM-Driven Policy Diffusion for offline RL generalization
• Combines text and trajectory prompts for policy learning
• Uses context-aware diffusion model for unseen tasks