Stable-Baselines 3 部分源代码解读 1 base_class.py
阅读PPO相关的源码，了解一下标准库是如何建立PPO算法以及各种tricks的，以便于自己的复现。在Pycharm里面一直跳转，可以看到PPO类是最终继承于基类，也就是这个py文件的内容。所以阅读源码就先从这里开始。: )
文章共12,717字 · 阅读需要大约43分钟
一键AI生成摘要，助你高效阅读
问答
·
Stable-Baselines 3 部分源代码解读 ./common/base_class.py
前言
阅读PPO相关的源码，了解一下标准库是如何建立PPO算法以及各种tricks的，以便于自己的复现。
在Pycharm里面一直跳转，可以看到PPO类是最终继承于基类，也就是这个py文件的内容。
所以阅读源码就先从这里开始。: )
import 包
从引用的文件可以推知出，这个基类包含了进度条显示、环境包装、噪声、输入是图片信息的预处理、保存zip文件和环境修改等文件代码。
import io
import pathlib
import time
import warnings
from abc import ABC, abstractmethod
from collections import deque
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
import gym
import numpy as np
import torch as th
from gym import spaces
from stable_baselines3.common import utils
from stable_baselines3.common.callbacks import BaseCallback, CallbackList, ConvertCallback, ProgressBarCallback
from stable_baselines3.common.env_util import is_wrapped
from stable_baselines3.common.logger import Logger
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.noise import ActionNoise
from stable_baselines3.common.policies import BasePolicy
from stable_baselines3.common.preprocessing import check_for_nested_spaces, is_image_space, is_image_space_channels_first
from stable_baselines3.common.save_util import load_from_zip_file, recursive_getattr, recursive_setattr, save_to_zip_file
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
from stable_baselines3.common.utils import (
check_for_correct_spaces,
get_device,
get_schedule_fn,
get_system_info,
set_random_seed,
update_learning_rate,
)
from stable_baselines3.common.vec_env import (
DummyVecEnv,
VecEnv,
VecNormalize,
VecTransposeImage,
is_vecenv_wrapped,
unwrap_vec_normalize,
)
变量
目前不太明白这个变量的具体作用，通过后续的代码上下文再做推断。
SelfBaseAlgorithm = TypeVar("SelfBaseAlgorithm", bound="BaseAlgorithm")
函数
Stable-Baselines 3 在文档中给出，如果环境是Gym库登记好的，那么可以直接通过输入字符串的方式，创建新的环境。这个函数maybe_make_env(...)便是用来处理字符串的输入。如果不是字符串，而是一个环境的实例的话，那么就直接跳过这个函数。
def maybe_make_env(env: Union[GymEnv, str, None], verbose: int) -> Optional[GymEnv]:
# verbose 提示信息的等级，数值为0表示不输出，数值为1来提示环境是否创建成功了。
"""If env is a string, make the environment; otherwise, return env.
:param env: The environment to learn from.
:param verbose: Verbosity level: 0 for no output, 1 for indicating if envrironment is created
:return A Gym (vector) environment.
"""
# if判断 env 位置变量是否是字符串类型，如果是则执行；反之不是，则跳过条件判断。
if isinstance(env, str):
# verbose 信息等级至少1级以上，在控制台输出创建信息。“正在创建给定名字的环境...”
if verbose >= 1:
print(f"Creating environment from the given name '{env}'")
# 创建环境实例
env = gym.make(env)
# 返回创建好的环境实例
return env
类
这个就是基类了，其他算法继承于这个基类中。
policy: 使用的策略模型，必填。可以是多层感知机模型MlpPolicy，也可以是卷积神经网络模型CnnPolicy，抑或是其他的模型。env: 需要进行学习的环境，智能体处在的环境，必填。learning_rate: 优化器的学习率，必填。这个参数也可以是一个与当前剩余进度有关的函数。(from 1 to 0)的意思是学习率必须随着智能体的交互而不断递减。policy_kwargs: 策略模型可以是自己创建的。自己创建时需要将模型的每层参数和激活函数以及其他信息写成字典，并由这个参数传进去。tensorboard_log: 使用tensorboard来记录日志信息。这个位置参数输入的是记录日志信息的位置/文件。默认值为None表示不输出日志。verbose: 记录运行过程的信息的详细程度。0不输出任何信息，1输出设备和装饰器使用信息，2输出调试信息。可以推断出verbose数值越大等级越高那么输出的信息越详细。device: 这个位置参数用来指定代码在那个器件上运行。缺省时，字符串“auto”，先检测代码是否能在gpu环境下运行，能的话就在gpu下；要是不能的话，就转回到cpu下。support_multi_env: 默认False，部分算法（A2C）需要在多环境下进行，那么通过这个位置参数标识出是否要多个环境进行训练。monitor_wrapper: 默认True，创建环境时选择是否需要监视器包装器。seed: 默认是None，是否需要通过伪随机生成器生成伪随机数。use_sde: 默认是False，不使用这个探索模式。对智能体的探索方式做处理。选择生成状态独立性探索，还是通过增加噪声方式进行探索。（需要再进一步了解下～）sde_sample_freq: 当使用生成状态独立性探索时候，选择几步进行一次噪声矩阵的生成。缺省数值是-，也就是在每个rollout开始时生成一个新的噪声矩阵。supported_action_spaces: 算法支持的动作空间是如何的。
class BaseAlgorithm(ABC):
"""
The base of RL algorithms
:param policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
:param env: The environment to learn from
(if registered in Gym, can be str. Can be None for loading trained models)
:param learning_rate: learning rate for the optimizer,
it can be a function of the current progress remaining (from 1 to 0)
:param policy_kwargs: Additional arguments to be passed to the policy on creation
:param tensorboard_log: the log location for tensorboard (if None, no logging)
:param verbose: Verbosity level: 0 for no output, 1 for info messages (such as device or wrappers used), 2 for debug messages
:param device: Device on which the code should run. By default, it will try to use a Cuda compatible device and fallback to cpu if it is not possible.
:param support_multi_env: Whether the algorithm supports training with multiple environments (as in A2C)
:param monitor_wrapper: When creating an environment, whether to wrap it
or not in a Monitor wrapper.
:param seed: Seed for the pseudo random generators
:param use_sde: Whether to use generalized State Dependent Exploration (gSDE)
instead of action noise exploration (default: False)
:param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
Default: -1 (only sample at the beginning of the rollout)
:param supported_action_spaces: The action spaces supported by the algorithm.
"""
# Policy aliases (see _get_policy_from_name())
policy_aliases: Dict[str, Type[BasePolicy]] = {}
def __init__(
self,
policy: Union[str, Type[BasePolicy]],
env: Union[GymEnv, str, None],
learning_rate: Union[float, Schedule],
policy_kwargs: Optional[Dict[str, Any]] = None,
tensorboard_log: Optional[str] = None,
verbose: int = 0,
device: Union[th.device, str] = "auto",
support_multi_env: bool = False,
monitor_wrapper: bool = True,
seed: Optional[int] = None,
use_sde: bool = False,
sde_sample_freq: int = -1,
supported_action_spaces: Optional[Tuple[spaces.Space, ...]] = None,
):
## 如果策略是字符串类型时，self.policy_class调用“self._get_policy_from_name()”获得实际策略。不是字符串类型时，就直接把policy赋值给self.policy_class。
if isinstance(policy, str):
self.policy_class = self._get_policy_from_name(policy)
else:
self.policy_class = policy
## 对gpu/cpu做处理。在verbose大于等于1级的时候输出硬件信息。
self.device = get_device(device)
if verbose >= 1:
print(f"Using {self.device} device")
# self.env初始化环境，self._vec_normalize_env应该是类似于环境是否归一化。
self.env = None # type: Optional[GymEnv]
# get VecNormalize object if needed
self._vec_normalize_env = unwrap_vec_normalize(env)
# self.verbose初始化输出信息等级
# self.policy_kwargs客制化策略模型的参数（神经元数量和激活函数）
self.verbose = verbose
self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
# self.observation_space初始化观测空间
# self.action_space初始化动作空间
self.observation_space = None # type: Optional[spaces.Space]
self.action_space = None # type: Optional[spaces.Space]
# self.n_envs初始化环境的数量，可能存在多的环境
# self.num_timesteps初始化步骤数量，具体含义从下文推断
self.n_envs = None
self.num_timesteps = 0
# Used for updating schedules
# self._total_timesteps用于初始化训练次数，待定
self._total_timesteps = 0
# Used for computing fps, it is updated at each call of learn()
# self._num_timesteps_at_start用于计算帧率，在调用learn()时候做更新。
self._num_timesteps_at_start = 0
# 初始化随机数种子、动作噪声、开始时间、策略、学习率、tensorboard日志和学习率的函数。
self.seed = seed
self.action_noise: Optional[ActionNoise] = None
self.start_time = None
self.policy = None
self.learning_rate = learning_rate
self.tensorboard_log = tensorboard_log
self.lr_schedule = None # type: Optional[Schedule]
# 初始化上一回合最后的观测、开始状态，以及上一回合的初始观测数据，和整个回合的数量。
self._last_obs = None # type: Optional[Union[np.ndarray, Dict[str, np.ndarray]]]
self._last_episode_starts = None # type: Optional[np.ndarray]
# When using VecNormalize:
self._last_original_obs = None # type: Optional[Union[np.ndarray, Dict[str, np.ndarray]]]
self._episode_num = 0
# 使用状态独立性探索的初始化信息
# Used for gSDE only
self.use_sde = use_sde
self.sde_sample_freq = sde_sample_freq
# 初始化进度条信息，帮助学习率的动态调节
# Track the training progress remaining (from 1 to 0)
# this is used to update the learning rate
self._current_progress_remaining = 1
# self.ep_info_buffer和self.ep_success_buffer用于记录智能体与环境交互时候的成功率和info信息，也就是step()方法下得到的结果。
# Buffers for logging
self.ep_info_buffer = None # type: Optional[deque]
self.ep_success_buffer = None # type: Optional[deque]
# self._n_updates初始化更新次数
# For logging (and TD3 delayed updates)
self._n_updates = 0 # type: int
# 初始化日志器和客制化日志器
# The logger object
self._logger = None # type: Logger
# Whether the user passed a custom logger or not
self._custom_logger = False
# 对环境做基本的包装，并记录状态空间、动作空间、环境的数量以及环境实例信息。
# Create and wrap the env if needed
if env is not None:
env = maybe_make_env(env, self.verbose)
env = self._wrap_env(env, self.verbose, monitor_wrapper)
self.observation_space = env.observation_space
self.action_space = env.action_space
self.n_envs = env.num_envs
self.env = env
# 判断动作空间是否合法，也就是是否是spaces.Space类型的。
# 如果不是，则发生异常并报错。
if supported_action_spaces is not None:
assert isinstance(self.action_space, supported_action_spaces), (
f"The algorithm only supports {supported_action_spaces} as action spaces "
f"but {self.action_space} was provided"
)
# 判断传进来的环境是不是多个环境，如果是则报错
# 提示用户需要进行“vectorized wrapper”，做统一整理
if not support_multi_env and self.n_envs > 1:
raise ValueError(
"Error: the model does not support multiple envs; it requires " "a single vectorized environment."
)
# 如果输入的策略字符串不是"MlpPolicy"和"CnnPolicy"，则报错
# Catch common mistake: using MlpPolicy/CnnPolicy instead of MultiInputPolicy
if policy in ["MlpPolicy", "CnnPolicy"] and isinstance(self.observation_space, spaces.Dict):
raise ValueError(f"You must use `MultiInputPolicy` when working with dict observation space, not {policy}")
# 当使用状态独立性探索且动作空间不是spaces.Box的实例的时候，报错“状态独立性探索只能在连续空 # 间中存在”
if self.use_sde and not isinstance(self.action_space, spaces.Box):
raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.")
# 在连续空间中，需要对动作空间的上界和下界做无穷性检验。
if isinstance(self.action_space, spaces.Box):
assert np.all(
np.isfinite(np.array([self.action_space.low, self.action_space.high]))
), "Continuous action space must have a finite lower and upper bound"
@staticmethod
def _wrap_env(env: GymEnv, verbose: int = 0, monitor_wrapper: bool = True) -> VecEnv:
# 如果需要的话，用恰当的包装器包装环境。
# 例如，需要一个向量化的环境；对于有图像的环境，需要对图像通道重新排序。
# 输入：env环境；verbose：输出控制台的信息简略程度；
# monitor_wrapper：是否需要增加“monitor_wrapper”类似的包装器。
""" "
Wrap environment with the appropriate wrappers if needed.
For instance, to have a vectorized environment
or to re-order the image channels.
:param env:
:param verbose: Verbosity level: 0 for no output, 1 for indicating wrappers used
:param monitor_wrapper: Whether to wrap the env in a ``Monitor`` when possible.
:return: The wrapped environment.
"""
if not isinstance(env, VecEnv):
# 判断这个环境是否包装了Monitor包装器。如果没有被包装且有包装意向的话，执行接下来的语句。
# A monitor wrapper for Gym environments, it is used to know the episode # reward, length, time and other data.
# Monitor包装器可以知道回合奖励、回合长度、运行时间和其他数据。
if not is_wrapped(env, Monitor) and monitor_wrapper:
# 输出详细信息。
if verbose >= 1:
print("Wrapping the env with a `Monitor` wrapper")
# 进行包装，用Monitor的类嵌套，主要是对重置环境、执行动作时数据的收集、简单处理以及写入 # 日志中。
env = Monitor(env)
if verbose >= 1:
print("Wrapping the env in a DummyVecEnv.")
# 如果是多任务的话，就包装DummyVecEnv()类，使用多线程来训练强化学习。
env = DummyVecEnv([lambda: env])
# Make sure that dict-spaces are not nested (not supported)
# 判断是否有嵌套的观测空间，tuple/dict里面还有tuple/dict
check_for_nested_spaces(env.observation_space)
# 对图像的状态信息是否需要重整通道进行判断并调整。
if not is_vecenv_wrapped(env, VecTransposeImage):
wrap_with_vectranspose = False
if isinstance(env.observation_space, spaces.Dict):
# If even one of the keys is a image-space in need of transpose, apply transpose
# If the image spaces are not consistent (for instance one is channel first,
# the other channel last), VecTransposeImage will throw an error
for space in env.observation_space.spaces.values():
wrap_with_vectranspose = wrap_with_vectranspose or (
is_image_space(space) and not is_image_space_channels_first(space)
)
else:
wrap_with_vectranspose = is_image_space(env.observation_space) and not is_image_space_channels_first(
env.observation_space
)
if wrap_with_vectranspose:
if verbose >= 1:
print("Wrapping the env in a VecTransposeImage.")
env = VecTransposeImage(env)
# 最后返回包装后的环境
return env
@abstractmethod
def _setup_model(self) -> None:
# 建立模型
"""Create networks, buffer and optimizers."""
def set_logger(self, logger: Logger) -> None:
# 设置日志记录器
"""
Setter for for logger object.
.. warning::
When passing a custom logger object,
this will overwrite ``tensorboard_log`` and ``verbose`` settings
passed to the constructor.
"""
self._logger = logger
# User defined logger
self._custom_logger = True
@property
def logger(self) -> Logger:
"""Getter for the logger object."""
return self._logger
def _setup_lr_schedule(self) -> None:
"""Transform to callable if needed."""
self.lr_schedule = get_schedule_fn(self.learning_rate)
def _update_current_progress_remaining(self, num_timesteps: int, total_timesteps: int) -> None:
# 计算当前的进度信息
# 进度 = 1 - （ 已经运行的时间步 / 总的时间步 ）
"""
Compute current progress remaining (starts from 1 and ends to 0)
:param num_timesteps: current number of timesteps
:param total_timesteps:
"""
self._current_progress_remaining = 1.0 - float(num_timesteps) / float(total_timesteps)
def _update_learning_rate(self, optimizers: Union[List[th.optim.Optimizer], th.optim.Optimizer]) -> None:
# 用当前学习率常数和当前剩余进度更新当前的更新器的学习率
"""
Update the optimizers learning rate using the current learning rate schedule
and the current progress remaining (from 1 to 0).
:param optimizers:
An optimizer or a list of optimizers.
"""
# Log the current learning rate
# 对学习率先进行记录
self.logger.record("train/learning_rate", self.lr_schedule(self._current_progress_remaining))
# 如果优化器不是list类，那么就包装成list类，然后遍历这个list来更新优化器的学习率。
if not isinstance(optimizers, list):
optimizers = [optimizers]
for optimizer in optimizers:
update_learning_rate(optimizer, self.lr_schedule(self._current_progress_remaining))
def _excluded_save_params(self) -> List[str]:
# 这个函数是输出一些变量的名字
# 在模型数据保存的时候，这些名字的变量被排除，不进行保存
# 例如经验池子，占据的空间太大了
"""
Returns the names of the parameters that should be excluded from being
saved by pickling. E.g. replay buffers are skipped by default
as they take up a lot of space. PyTorch variables should be excluded
with this so they can be stored with ``th.save``.
:return: List of parameters that should be excluded from being saved with pickle.
"""
return [
"policy",
"device",
"env",
"replay_buffer",
"rollout_buffer",
"_vec_normalize_env",
"_episode_storage",
"_logger",
"_custom_logger",
]
def _get_policy_from_name(self, policy_name: str) -> Type[BasePolicy]:
## 从名字表达式中获得策略的类。这样做目标是标准化策略命名，例如，所有策略都是调用"MlpPolicy"或 ## 者"CnnPolicy"，并收到对应的策略来执行。
"""
Get a policy class from its name representation.
The goal here is to standardize policy naming, e.g.
all algorithms can call upon "MlpPolicy" or "CnnPolicy",
and they receive respective policies that work for them.
:param policy_name: Alias of the policy
:return: A policy class (type)
"""
# 如果策略的名字字符串在self.policy_aliases的字典内，那么就以这个字符串为key，返回这个key对应的value。在此文件中，self.policy_aliases是空字典，这是因为这是基类暂时不做具体的定义。如果字典内没有相应的字符串key，那么输出报错。
if policy_name in self.policy_aliases:
return self.policy_aliases[policy_name]
else:
raise ValueError(f"Policy {policy_name} unknown")
def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
# 这个方法制定那些变量用torch.state_dict，那些用的是torch.save
# 这样便于处理设备转移过程中保存时遇到的不必要的麻烦
"""
Get the name of the torch variables that will be saved with
PyTorch ``th.save``, ``th.load`` and ``state_dicts`` instead of the default
pickling strategy. This is to handle device placement correctly.
Names can point to specific variables under classes, e.g.
"policy.optimizer" would point to ``optimizer`` object of ``self.policy``
if this object.
:return:
List of Torch variables whose state dicts to save (e.g. th.nn.Modules),
and list of other Torch variables to store with ``th.save``.
"""
state_dicts = ["policy"]
return state_dicts, []
def _init_callback(
self,
callback: MaybeCallback,
progress_bar: bool = False,
) -> BaseCallback:
# 对召回函数做一些处理，以及进度条的显示（会导入tqdm和rich）
"""
:param callback: Callback(s) called at every step with state of the algorithm.
:param progress_bar: Display a progress bar using tqdm and rich.
:return: A hybrid callback calling `callback` and performing evaluation.
"""
# Convert a list of callbacks into a callback
if isinstance(callback, list):
callback = CallbackList(callback)
# Convert functional callback to object
if not isinstance(callback, BaseCallback):
callback = ConvertCallback(callback)
# Add progress bar callback
if progress_bar:
callback = CallbackList([callback, ProgressBarCallback()])
callback.init_callback(self)
return callback
def _setup_learn(
self,
total_timesteps: int,
callback: MaybeCallback = None,
reset_num_timesteps: bool = True,
tb_log_name: str = "run",
progress_bar: bool = False,
) -> Tuple[int, BaseCallback]:
# 这个函数是对强化学习训练过程做了初步的建立
"""
Initialize different variables needed for training.
:param total_timesteps: The total number of samples (env steps) to train on
:param callback: Callback(s) called at every step with state of the algorithm.
:param reset_num_timesteps: Whether to reset or not the ``num_timesteps`` attribute
:param tb_log_name: the name of the run for tensorboard log
:param progress_bar: Display a progress bar using tqdm and rich.
:return: Total timesteps and callback(s)
"""
# 设置强化学习训练开始的时间
self.start_time = time.time_ns()
# 如果self.ep_info_buffer这个记录器是空的，且reset_num_timesteps为true
# 也就是需要进行重置时间步的时候
# 执行如下代码
if self.ep_info_buffer is None or reset_num_timesteps:
# Initialize buffers if they don't exist, or reinitialize if resetting counters
self.ep_info_buffer = deque(maxlen=100)
self.ep_success_buffer = deque(maxlen=100)
# 如果噪声不存在的话，重置噪声
if self.action_noise is not None:
self.action_noise.reset()
# reset_num_timesteps为true
# 也就是需要进行重置时间步的时候
# 执行如下代码
# else表示增加当前的训练步长
if reset_num_timesteps:
self.num_timesteps = 0
self._episode_num = 0
else:
# Make sure training timesteps are ahead of the internal counter
total_timesteps += self.num_timesteps
self._total_timesteps = total_timesteps
self._num_timesteps_at_start = self.num_timesteps
# 避免连续调用.learn()函数而做的一些处理
# Avoid resetting the environment when calling ``.learn()`` consecutive times
if reset_num_timesteps or self._last_obs is None:
self._last_obs = self.env.reset() # pytype: disable=annotation-type-mismatch
self._last_episode_starts = np.ones((self.env.num_envs,), dtype=bool)
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
self._last_original_obs = self._vec_normalize_env.get_original_obs()
# 如果用户没有自己设置日志信息的话，自己配置日志信息和输出
# Configure logger's outputs if no logger was passed
if not self._custom_logger:
self._logger = utils.configure_logger(self.verbose, self.tensorboard_log, tb_log_name, reset_num_timesteps)
# 对召回函数做处理
# Create eval callback if needed
callback = self._init_callback(callback, progress_bar)
return total_timesteps, callback
def _update_info_buffer(self, infos: List[Dict[str, Any]], dones: Optional[np.ndarray] = None) -> None:
# 如果使用了Gym库的Monitor的包装器，就取回奖励、回合长度、回合内的成功率以及更新经验池
"""
Retrieve reward, episode length, episode success and update the buffer
if using Monitor wrapper or a GoalEnv.
:param infos: List of additional information about the transition.
:param dones: Termination signals
"""
# 任务没有执行成功的话，done输出一个np.array的矩阵
if dones is None:
dones = np.array([False] * len(infos))
# 从任务执行的一系列info中取出每一个信息
# 获得可能是"episode"回合长度和"is_success"是否成功的信息
for idx, info in enumerate(infos):
maybe_ep_info = info.get("episode")
maybe_is_success = info.get("is_success")
if maybe_ep_info is not None:
self.ep_info_buffer.extend([maybe_ep_info])
if maybe_is_success is not None and dones[idx]:
self.ep_success_buffer.append(maybe_is_success)
def get_env(self) -> Optional[VecEnv]:
# 单任务而言，获取当前的环境实例
"""
Returns the current environment (can be None if not defined).
:return: The current environment
"""
return self.env
def get_vec_normalize_env(self) -> Optional[VecNormalize]:
# 多任务而言，获取当前的批量环境
"""
Return the ``VecNormalize`` wrapper of the training env
if it exists.
:return: The ``VecNormalize`` env.
"""
return self._vec_normalize_env
def set_env(self, env: GymEnv, force_reset: bool = True) -> None:
"""
Checks the validity of the environment, and if it is coherent, set it as the current environment.
Furthermore wrap any non vectorized env into a vectorized
checked parameters:
- observation_space
- action_space
:param env: The environment for learning a policy
:param force_reset: Force call to ``reset()`` before training
to avoid unexpected behavior.
See issue https://github.com/DLR-RM/stable-baselines3/issues/597
"""
# 检查环境是否有效，如果无效则报错
# 检查环境是否包装成向量环境了，没有的话就包装
# if it is not a VecEnv, make it a VecEnv
# and do other transformations (dict obs, image transpose) if needed
env = self._wrap_env(env, self.verbose)
assert env.num_envs == self.n_envs, (
"The number of environments to be set is different from the number of environments in the model: "
f"({env.num_envs} != {self.n_envs}), whereas `set_env` requires them to be the same. To load a model with "
f"a different number of environments, you must use `{self.__class__.__name__}.load(path, env)` instead"
)
# Check that the observation spaces match
check_for_correct_spaces(env, self.observation_space, self.action_space)
# Update VecNormalize object
# otherwise the wrong env may be used, see https://github.com/DLR-RM/stable-baselines3/issues/637
self._vec_normalize_env = unwrap_vec_normalize(env)
# Discard `_last_obs`, this will force the env to reset before training
# See issue https://github.com/DLR-RM/stable-baselines3/issues/597
# 强制重置，避免意外发生
if force_reset:
self._last_obs = None
self.n_envs = env.num_envs
self.env = env
@abstractmethod
def learn(
self: SelfBaseAlgorithm,
total_timesteps: int,
callback: MaybeCallback = None,
log_interval: int = 100,
tb_log_name: str = "run",
reset_num_timesteps: bool = True,
progress_bar: bool = False,
) -> SelfBaseAlgorithm:
# 这是一个抽象的方法
# 主要是给后面的子类提供一些输入
# 除了基本算法和总时间步长以外其他都默认数值了
"""
Return a trained model.
:param total_timesteps: The total number of samples (env steps) to train on
:param callback: callback(s) called at every step with state of the algorithm.
:param log_interval: The number of episodes before logging.
:param tb_log_name: the name of the run for TensorBoard logging
:param reset_num_timesteps: whether or not to reset the current timestep number (used in logging)
:param progress_bar: Display a progress bar using tqdm and rich.
:return: the trained model
"""
def predict(
self,
observation: Union[np.ndarray, Dict[str, np.ndarray]],
state: Optional[Tuple[np.ndarray, ...]] = None,
episode_start: Optional[np.ndarray] = None,
deterministic: bool = False,
) -> Tuple[np.ndarray, Optional[Tuple[np.ndarray, ...]]]:
# 从观测信息中获得动作信息
# 引入一些包装器对图片、批量状态做一些包装处理
# observation是当前的观测，state是隐藏状态（用于循环神经网络）
# episode_start是回合的开始信息（用于循环神经网络，回合开始时重置潜在状态）
# deterministic标志是否输出确定性的动作
"""
Get the policy action from an observation (and optional hidden state).
Includes sugar-coating to handle different observations (e.g. normalizing images).
:param observation: the input observation
:param state: The last hidden states (can be None, used in recurrent policies)
:param episode_start: The last masks (can be None, used in recurrent policies)
this correspond to beginning of episodes,
where the hidden states of the RNN must be reset.
:param deterministic: Whether or not to return deterministic actions.
:return: the model's action and the next hidden state
(used in recurrent policies)
"""
return self.policy.predict(observation, state, episode_start, deterministic)
def set_random_seed(self, seed: Optional[int] = None) -> None:
# 设置随机数种子，并设置在环境中
"""
Set the seed of the pseudo-random generators
(python, numpy, pytorch, gym, action_space)
:param seed:
"""
if seed is None:
return
set_random_seed(seed, using_cuda=self.device.type == th.device("cuda").type)
self.action_space.seed(seed)
# self.env is always a VecEnv
if self.env is not None:
self.env.seed(seed)
def set_parameters(
self,
load_path_or_dict: Union[str, Dict[str, Dict]],
exact_match: bool = True,
device: Union[th.device, str] = "auto",
) -> None:
# 如果模型是预训练的，那么就从zip文件中导入网络模型
# load_path_or_dict，指定模型保存的位置。这个文件可能是字典、也可能是其他的
# 由torch.nn.Module.state_dict()方法得到的
# exact_match表示精确的匹配，如果导入的参数没有符合预先设置的参数，那么就报错
# exact_match默认为True，表示进行匹配
"""
Load parameters from a given zip-file or a nested dictionary containing parameters for
different modules (see ``get_parameters``).
:param load_path_or_iter: Location of the saved data (path or file-like, see ``save``), or a nested dictionary containing nn.Module parameters used by the policy. The dictionary maps object names to a state-dictionary returned by ``torch.nn.Module.state_dict()``.
:param exact_match: If True, the given parameters should include parameters for each module and each of their parameters, otherwise raises an Exception. If set to False, this can be used to update only specific parameters.
:param device: Device on which the code should run.
"""
params = None
# 如果是字典类型，就直接复制给param
if isinstance(load_path_or_dict, dict):
params = load_path_or_dict
# 如果不是字典类型，那么就是zip文件，那么就直接输出第二个参数
else:
_, params, _ = load_from_zip_file(load_path_or_dict, device=device)
# 更新一下参数更新的状态
# Keep track which objects were updated.
# `_get_torch_save_params` returns [params, other_pytorch_variables].
# We are only interested in former here.
objects_needing_update = set(self._get_torch_save_params()[0])
updated_objects = set()
for name in params:
attr = None
try:
attr = recursive_getattr(self, name)
except Exception as e:
# What errors recursive_getattr could throw? KeyError, but
# possible something else too (e.g. if key is an int?).
# Catch anything for now.
raise ValueError(f"Key {name} is an invalid object name.") from e
if isinstance(attr, th.optim.Optimizer):
# Optimizers do not support "strict" keyword...
# Seems like they will just replace the whole
# optimizer state with the given one.
# On top of this, optimizer state-dict
# seems to change (e.g. first ``optim.step()``),
# which makes comparing state dictionary keys
# invalid (there is also a nesting of dictionaries
# with lists with dictionaries with ...), adding to the
# mess.
#
# TL;DR: We might not be able to reliably say
# if given state-dict is missing keys.
#
# Solution: Just load the state-dict as is, and trust
# the user has provided a sensible state dictionary.
attr.load_state_dict(params[name])
else:
# Assume attr is th.nn.Module
attr.load_state_dict(params[name], strict=exact_match)
updated_objects.add(name)
if exact_match and updated_objects != objects_needing_update:
raise ValueError(
"Names of parameters do not match agents' parameters: "
f"expected {objects_needing_update}, got {updated_objects}"
)
@classmethod # noqa: C901
def load(
cls: Type[SelfBaseAlgorithm],
path: Union[str, pathlib.Path, io.BufferedIOBase],
env: Optional[GymEnv] = None,
device: Union[th.device, str] = "auto",
custom_objects: Optional[Dict[str, Any]] = None,
print_system_info: bool = False,
force_reset: bool = True,
**kwargs,
) -> SelfBaseAlgorithm:
# 从zip文件中导入模型
"""
Load the model from a zip-file.
Warning: ``load`` re-creates the model from scratch, it does not update it in-place!
For an in-place load use ``set_parameters`` instead.
:param path: path to the file (or a file-like) where to
load the agent from
:param env: the new environment to run the loaded model on
(can be None if you only need prediction from a trained model) has priority over any saved environment
:param device: Device on which the code should run.
:param custom_objects: Dictionary of objects to replace
upon loading. If a variable is present in this dictionary as a
key, it will not be deserialized and the corresponding item
will be used instead. Similar to custom_objects in
``keras.models.load_model``. Useful when you have an object in
file that can not be deserialized.
:param print_system_info: Whether to print system info from the saved model
and the current system info (useful to debug loading issues)
:param force_reset: Force call to ``reset()`` before training
to avoid unexpected behavior.
See https://github.com/DLR-RM/stable-baselines3/issues/597
:param kwargs: extra arguments to change the model when loading
:return: new model instance with loaded parameters
"""
if print_system_info:
print("== CURRENT SYSTEM INFO ==")
get_system_info()
data, params, pytorch_variables = load_from_zip_file(
path,
device=device,
custom_objects=custom_objects,
print_system_info=print_system_info,
)
# Remove stored device information and replace with ours
if "policy_kwargs" in data:
if "device" in data["policy_kwargs"]:
del data["policy_kwargs"]["device"]
# backward compatibility, convert to new format
if "net_arch" in data["policy_kwargs"] and len(data["policy_kwargs"]["net_arch"]) > 0:
saved_net_arch = data["policy_kwargs"]["net_arch"]
if isinstance(saved_net_arch, list) and isinstance(saved_net_arch[0], dict):
data["policy_kwargs"]["net_arch"] = saved_net_arch[0]
if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]:
raise ValueError(
f"The specified policy kwargs do not equal the stored policy kwargs."
f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}"
)
if "observation_space" not in data or "action_space" not in data:
raise KeyError("The observation_space and action_space were not given, can't verify new environments")
if env is not None:
# Wrap first if needed
env = cls._wrap_env(env, data["verbose"])
# Check if given env is valid
check_for_correct_spaces(env, data["observation_space"], data["action_space"])
# Discard `_last_obs`, this will force the env to reset before training
# See issue https://github.com/DLR-RM/stable-baselines3/issues/597
if force_reset and data is not None:
data["_last_obs"] = None
# `n_envs` must be updated. See issue https://github.com/DLR-RM/stable-baselines3/issues/1018
if data is not None:
data["n_envs"] = env.num_envs
else:
# Use stored env, if one exists. If not, continue as is (can be used for predict)
if "env" in data:
env = data["env"]
# noinspection PyArgumentList
model = cls( # pytype: disable=not-instantiable,wrong-keyword-args
policy=data["policy_class"],
env=env,
device=device,
_init_setup_model=False, # pytype: disable=not-instantiable,wrong-keyword-args
)
# load parameters
model.__dict__.update(data)
model.__dict__.update(kwargs)
model._setup_model()
try:
# put state_dicts back in place
model.set_parameters(params, exact_match=True, device=device)
except RuntimeError as e:
# Patch to load Policy saved using SB3 < 1.7.0
# the error is probably due to old policy being loaded
# See https://github.com/DLR-RM/stable-baselines3/issues/1233
if "pi_features_extractor" in str(e) and "Missing key(s) in state_dict" in str(e):
model.set_parameters(params, exact_match=False, device=device)
warnings.warn(
"You are probably loading a model saved with SB3 < 1.7.0, "
"we deactivated exact_match so you can save the model "
"again to avoid issues in the future "
"(see https://github.com/DLR-RM/stable-baselines3/issues/1233 for more info). "
f"Original error: {e} \n"
"Note: the model should still work fine, this only a warning."
)
else:
raise e
# put other pytorch variables back in place
if pytorch_variables is not None:
for name in pytorch_variables:
# Skip if PyTorch variable was not defined (to ensure backward compatibility).
# This happens when using SAC/TQC.
# SAC has an entropy coefficient which can be fixed or optimized.
# If it is optimized, an additional PyTorch variable `log_ent_coef` is defined,
# otherwise it is initialized to `None`.
if pytorch_variables[name] is None:
continue
# Set the data attribute directly to avoid issue when using optimizers
# See https://github.com/DLR-RM/stable-baselines3/issues/391
recursive_setattr(model, name + ".data", pytorch_variables[name].data)
# Sample gSDE exploration matrix, so it uses the right device
# see issue #44
if model.use_sde:
model.policy.reset_noise() # pytype: disable=attribute-error
return model
# 获得参数
def get_parameters(self) -> Dict[str, Dict]:
"""
Return the parameters of the agent. This includes parameters from different networks, e.g.
critics (value functions) and policies (pi functions).
:return: Mapping of from names of the objects to PyTorch state-dicts.
"""
state_dicts_names, _ = self._get_torch_save_params()
params = {}
for name in state_dicts_names:
attr = recursive_getattr(self, name)
# Retrieve state dict
params[name] = attr.state_dict()
return params
def save(
self,
path: Union[str, pathlib.Path, io.BufferedIOBase],
exclude: Optional[Iterable[str]] = None,
include: Optional[Iterable[str]] = None,
) -> None:
# 保存模型
"""
Save all the attributes of the object and the model parameters in a zip-file.
:param path: path to the file where the rl agent should be saved
:param exclude: name of parameters that should be excluded in addition to the default ones
:param include: name of parameters that might be excluded but should be included anyway
"""
# Copy parameter list so we don't mutate the original dict
data = self.__dict__.copy()
# Exclude is union of specified parameters (if any) and standard exclusions
if exclude is None:
exclude = []
exclude = set(exclude).union(self._excluded_save_params())
# Do not exclude params if they are specifically included
if include is not None:
exclude = exclude.difference(include)
state_dicts_names, torch_variable_names = self._get_torch_save_params()
all_pytorch_variables = state_dicts_names + torch_variable_names
for torch_var in all_pytorch_variables:
# We need to get only the name of the top most module as we'll remove that
var_name = torch_var.split(".")[0]
# Any params that are in the save vars must not be saved by data
exclude.add(var_name)
# Remove parameter entries of parameters which are to be excluded
for param_name in exclude:
data.pop(param_name, None)
# Build dict of torch variables
pytorch_variables = None
if torch_variable_names is not None:
pytorch_variables = {}
for name in torch_variable_names:
attr = recursive_getattr(self, name)
pytorch_variables[name] = attr
# Build dict of state_dicts
params_to_save = self.get_parameters()
save_to_zip_file(path, data=data, params=params_to_save, pytorch_variables=pytorch_variables)
旨在为数千万中国开发者提供一个无缝且高效的云端环境，以支持学习、使用和贡献开源项目。
更多推荐
1
0- 0
已为社区贡献1条内容
所有评论(0)