Stable Baselines3 中如何避免保存外部LLM模型的技术方案

问题背景

在使用Stable Baselines3框架开发强化学习模型时，开发者经常会遇到需要集成外部模型（如大型语言模型LLM）来增强智能体决策能力的情况。本文以一个实际案例为例，展示了如何在自定义DQN策略中使用LLM模型来调整Q值，同时解决模型保存时遇到的技术难题。

技术实现方案

自定义Q网络架构

开发者创建了一个名为LangQNetwork的自定义Q网络类，继承自标准的QNetwork。这个网络的核心创新点在于：

1. LLM集成：通过构造函数接收外部LLM模型和模型名称参数
2. 专家查询功能：实现了query_expert方法，利用LLM模型根据当前屏幕图像生成专家建议
3. Q值调整：在预测方法中，可以选择性地将LLM生成的专家建议概率与原始Q值结合

class LangQNetwork(QNetwork):
    def __init__(self, *args, **kwargs):
        self.vlm_model = kwargs.pop("vlm_model", None)
        self.vlm_model_name = kwargs.pop("vlm_model_name",None)
        super().__init__(*args, **kwargs)
        
    @th.no_grad()
    def query_expert(self, screen_image: np.array, device: str, vlm_model, vlm_model_name) -> th.tensor:
        p = LMExpert(vlm_model,vlm_model_name,screen_image)
        p = th.tensor(p).to(device)
        return p
    
    def _predict(self, observation: PyTorchObs, deterministic: bool = True, screen_image: np.array = None, use_expert: bool = False) -> th.Tensor:
        q_values = self(observation)
        if use_expert:
            log_p = self.query_expert(screen_image, q_values.device,self.vlm_model,self.vlm_model_name)
            q_values = q_values + log_p
        action = q_values.argmax(dim=1).reshape(-1)
        return action

策略与算法层的扩展

为了将自定义Q网络集成到完整的强化学习流程中，开发者还实现了：

1. 自定义策略类：CustomDQNPolicy负责管理Q网络和目标Q网络
2. 自定义算法类：CustomDQN扩展了标准DQN算法，添加了专家模式切换逻辑

关键技术挑战与解决方案

模型保存问题

当使用Stable Baselines3的标准保存机制时，系统会尝试序列化整个策略对象，包括其中的LLM模型。这会导致以下问题：

1. LLM模型通常体积庞大，不适合直接序列化
2. 某些LLM框架（如vLLM）明确禁止pickle序列化
3. 保存的模型文件会不必要地包含LLM部分，影响存储和加载效率

解决方案：参数排除机制

Stable Baselines3提供了_excluded_save_params机制，允许开发者指定哪些参数不应该被保存。通过重写这个方法，可以排除包含LLM模型的policy_kwargs：

class CustomDQN(DQN):
    def _excluded_save_params(self):
        return super()._excluded_save_params() + ['policy_kwargs']

这种方法的关键优势在于：

1. 简洁高效：只需少量代码修改
2. 不影响功能：训练和推理时仍可使用LLM模型
3. 兼容性好：与Stable Baselines3的现有保存/加载机制无缝集成

实际应用建议

1. 初始化注意事项：在创建模型实例时，通过policy_kwargs传递LLM模型参数
2. 训练过程控制：可以在收集经验时周期性地启用专家建议
3. 调试技巧：在保存前检查将被序列化的数据，确保没有包含不必要的大对象

总结

本文介绍的技术方案展示了如何在Stable Baselines3框架中优雅地集成外部LLM模型，同时避免了模型保存时的技术障碍。这种方法不仅适用于DQN算法，也可以推广到其他基于策略或价值的强化学习算法中，为开发更智能的强化学习系统提供了灵活的技术支持。