Stable Baselines3中实现仅保留最近K个检查点的技术方案

背景介绍

在使用Stable Baselines3进行强化学习模型训练时，CheckpointCallback是一个常用的回调函数，它能够定期保存训练过程中的模型检查点。然而，随着训练时间的增长，这些检查点会不断累积，占用大量存储空间。许多情况下，我们只需要保留最近的几个检查点即可。

现有问题分析

Stable Baselines3内置的CheckpointCallback目前不支持自动删除旧检查点的功能。这意味着：

1. 长期训练会产生大量检查点文件
2. 手动删除旧检查点既麻烦又容易出错
3. 存储空间可能被不必要的旧检查点占满

解决方案实现

我们可以通过自定义回调函数来实现仅保留最近K个检查点的功能。以下是实现这一功能的技术细节：

核心思路

1. 使用最小堆(Min-Heap)数据结构来管理检查点
2. 每次保存新检查点时将其加入堆
3. 当堆大小超过K时，弹出并删除最早的检查点

代码实现

from heapq import heappop, heappush

class KeepLastKCheckpoints(BaseCallback):
    def __init__(self, k: int, save_freq: int, save_path: str, name_prefix: str = "rl_model", verbose: int = 0):
        super().__init__(verbose)
        self.k = k
        self.save_freq = save_freq
        self.save_path = save_path
        self.name_prefix = name_prefix
        # 使用最小堆存储检查点，按时间步排序
        self._ckpt_heap = []
    
    def _checkpoint_path(self, checkpoint_type: str = "", extension: str = "") -> str:
        # 生成检查点路径
        return os.path.join(self.save_path, f"{self.name_prefix}_{checkpoint_type}{self.num_timesteps}_steps.{extension}")

    def _on_step(self) -> bool:
        # 跟踪CheckpointCallback的触发
        if self.n_calls % self.save_freq == 0:
            heappush(self._ckpt_heap, (self.num_timesteps, self._checkpoint_path(extension="zip")))
            if len(self._ckpt_heap) > self.k:
                _, ckpt_name = heappop(self._ckpt_heap)
                os.remove(ckpt_name)
            
        return super()._on_step()

实现要点

1. 堆数据结构选择：使用最小堆可以高效地获取最早的检查点
2. 路径生成：复用CheckpointCallback的路径生成逻辑
3. 同步问题：需要确保该回调在CheckpointCallback之后执行

使用建议

1. 将此回调添加到回调列表中时，确保它位于CheckpointCallback之后
2. 根据实际需求调整K值，平衡存储空间和检查点保留需求
3. 对于长时间训练，建议设置合理的K值(如5-10)

潜在改进方向

虽然当前实现可以满足基本需求，但仍有改进空间：

1. 与CheckpointCallback集成：可以考虑将这一功能直接整合到CheckpointCallback中
2. 更灵活的删除策略：不仅基于数量，还可以考虑基于时间或性能指标
3. 错误处理：增加对文件删除失败的处理逻辑

总结

通过自定义回调函数实现仅保留最近K个检查点的功能，可以有效解决长期训练中检查点累积的问题。这种方案实现简单，对现有代码侵入性小，是当前Stable Baselines3框架下较为理想的解决方案。未来如果社区需求增加，可以考虑将其作为标准功能集成到CheckpointCallback中。