Stable-Baselines3 1.8.0中SubprocEnv reset()未触发及set_attr异常问题

🐛 问题详情 #

自定义gym.Env已通过Stable-Baselines3的check_env所有检查，当前使用Gym 0.21.0和SB3 1.8.0版本。

自定义环境包含device_freeze属性，用于在reset()方法中触发环境动态变化。需求是通过回调每经过n_timesteps就将device_freeze设为False，让后续reset()调用时改变环境动态（必须在episode结束时修改，符合MDP公式要求）。

自定义环境代码 #

import numpy as np
from commons import BaseInterface
from .oscar import OscarEnv
from typing import Iterable, Text, Dict
from numpy.typing import NDArray

class MarcellaEnv(OscarEnv): 
    def __init__(self, 
                 [..., parent class args],
                 devices_and_bounds:Dict[Text, float]={"device1": 50., "device2": 6., "device3": 7.}):
        """
        存储多任务训练使用的设备列表及对应性能边界。
        """
        self.device_freeze = True # 存在问题的属性
        self.devices_and_bounds = devices_and_bounds

        super().__init__(
           [...parent args]
        )

    @property
    def name(self): 
        return "marcella"
    
    def change_device(self):
        """
        如果未启用device_freeze，随机选择目标设备并更新。
        
        返回:
            None
        
        说明:
            仅当device_freeze未启用时，从可用设备中随机选择新设备，更新目标设备和边界，之后将device_freeze设为True。
        """
        if not self.device_freeze:
            new_device = np.random.choice(list(self.devices_and_bounds.keys()))
            self.target_device = new_device
            self.new_bound = self.devices_and_bounds[new_device]
            # 因为device_freeze为False进入逻辑，执行后将其设为True
            self.device_freeze = True

    def reset(self)->NDArray:
        """重置自定义环境属性。"""
        self._observation = self.observation_space.sample()
        self.change_device()
        self.update_current_net()

        self.timestep_counter= 0

        return self._get_obs()

回调代码 #

"""训练过程中用于记录的自定义回调。"""
from stable_baselines3.common.callbacks import BaseCallback
import numpy as np
from stable_baselines3.common.vec_env import VecEnv
from typing import Text, Iterable

class MultiTask_Callback(BaseCallback): 
    """继承自`BaseCallback`的自定义回调。

    :param verbose: (int) 日志级别 0: 无输出 1: 信息 2: 调试。

    触发时执行以下操作（作为EventCallback的子类）: 
        1. 评估当前策略（执行n_eval_episodes次）
        2. 更新当前最佳策略变量
        3. 在wandb上记录信息，具体内容见:_on_step方法。
    """
    def __init__(self, verbose:int=0):
        """初始化回调上下文。"""
        super().__init__(verbose)
        
        self.devices_history = []

    def _on_step(self) -> bool:
        """
        模型每次调用`_env.step()`后会执行此方法。
        作为EventCallback的子类，事件触发时会调用此方法。
        :return: (bool) 返回False会提前终止训练。
        """
        # 记录当前训练使用的硬件设备
        current_device = self.model.env.get_attr("target_device")

        # 存储模型当前训练的目标硬件
        self.devices_history.append(current_device)

        print(self.model.env.get_attr("device_freeze"), self.model.env.get_attr("target_device"))
        
        # 修改开关，允许episode初始化时选择不同设备
        self.model.env.set_attr("device_freeze", False)
 
        return True
    
    def get_devices_history(self):
        """返回所有硬件设备的历史记录"""
        return self.devices_history

现象与问题 #

• 使用DummyVecEnv或SubprocEnv时，训练脚本的回调输出如下（以2个环境为例）：

[True], [True] ['device1'], ['device1']
[False],[False] ['device1'], ['device1']
[False], [False]  ['device1'], ['device1']
[False], [False]  ['device1'], ['device1']
...

打印每30步触发一次，单episode最大步数设为50，此行为不符合预期——device_freeze被设为False后，后续reset()应该切换设备，但实际设备未变化，且set_attr仅第一次生效，之后即使_on_step持续执行，device_freeze也保持False不再响应修改。

• 若将回调中的self.model.env.set_attr("device_freeze", False)替换为以下代码，DummyVecEnv能正常工作，设备会按预期切换：

for env_idx in range(self.model.env.num_envs):
    # 手动修改，因为setattr似乎无效
    self.model.env.envs[env_idx].unwrapped.device_freeze = False

输出示例（1个环境）：

[True] ['device1']
[False]  ['device1']
[True]  ['device2']
[True]  ['device3']
[True]  ['device3']
[False]  ['device3']
[True] ['device1']
...

但此方法在SubprocEnv中无法使用，会报错：

AttributeError: SubprocEnv does not have the envs attribute

系统信息 #

• OS: Linux-5.13.0-52-generic-x86_64-with-glibc2.31 # 59~20.04.1-Ubuntu SMP Thu Jun 16 21:21:28 UTC 2022
• Python: 3.10.8
• Stable-Baselines3: 1.8.0
• PyTorch: 2.0.1+cu117
• GPU Enabled: True
• Numpy: 1.24.3
• Gym: 0.21.0

内容的提问来源于stack exchange，提问作者fracapuano