CDCP: Conditional Diffusion Model with Contextual Prompts for Multi-task Offline Safe Reinforcement Learning

2026-07-04

AI Summary

This work addresses the challenge of balancing safety and performance in multi-task offline safe reinforcement learning by proposing a policy learning framework based on conditional diffusion models. The approach enhances multi-task representation through a context prompting mechanism and enables flexible cost constraint control via classifier-free guidance. To mitigate interference across tasks, a gradient loss synchronization strategy is introduced, allowing the model to adapt to varying safety constraints without retraining. Experimental results demonstrate that the proposed method significantly outperforms current state-of-the-art approaches across diverse tasks and constraint settings, achieving higher returns while maintaining strong safety guarantees.

Abstract

Multi-task offline safe reinforcement learning (RL) promises to learn a shared optimal safe policy from offline data across multiple tasks. This paradigm provides an effective means for the widespread application of RL in multi-task scenarios with high risk and interaction costs. However, the triple challenges of multi-tasking, safety constraints, and out-of-distribution (OOD) actions pose a significant hurdle for existing methods to ensure safety while maximizing reward returns. In this work, we propose a Conditional Diffusion model with Contextual Prompts (CDCP) to address these challenges. Concretely, we first rethink the requirements and challenges in current multi-task decision-making and control scenarios and establish the objectives of multi-task offline safe RL. Subsequently, we transform the multi-task constrained optimization problem into a conditional generation problem using the diffusion model. Based on this, we design a classifier-free guided cost-constraint strategy to provide flexible cost constraints and eliminate extrapolation errors from OOD actions via supervised learning. Additionally, we introduce a novel contextual prompting method to enhance multi-task representation accuracy and adaptability to unseen tasks. A gradient loss synchronization strategy is also introduced to eliminate gradient interference, improving training stability. Finally, extensive experiments demonstrate that the CDCP algorithm exhibits higher performance and safety in multi-task scenarios than the current state-of-the-art baseline methods. It meets different cost constraints without further training, providing a more flexible cost-constraint solution for the multi-task safe RL.

Problem

• multi-task offline reinforcement learning
• safe reinforcement learning
• out-of-distribution actions
• safety constraints
• conditional generation

Innovation

• Conditional Diffusion Model
• Contextual Prompts
• Offline Safe Reinforcement Learning
• Multi-task RL
• Out-of-Distribution Generalization