diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 30f5a60..0bb3910 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -19,6 +19,7 @@ New Features: - Add support for Callback (cf https://github.com/hill-a/stable-baselines/pull/644) - Add methods for saving and loading replay buffer - Add `extend()` method to the buffers +- Add `get_vec_normalize_env()` to `BaseRLModel` to retrieve `VecNormalize` wrapper when it exists Bug Fixes: ^^^^^^^^^^ diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index ab0f7f1..3c21f69 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -123,6 +123,8 @@ def test_offpolicy_normalization(model_class): model = model_class('MlpPolicy', env, verbose=1) model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500) + # Check getter + assert isinstance(model.get_vec_normalize_env(), VecNormalize) def test_sync_vec_normalize(): diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 68c9025..d4e6810 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -14,7 +14,7 @@ import numpy as np from torchy_baselines.common import logger from torchy_baselines.common.policies import BasePolicy, get_policy_from_name from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate -from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize +from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback @@ -212,10 +212,18 @@ class BaseRLModel(ABC): """ Returns the current environment (can be None if not defined). - :return: The current environment + :return: (Optional[VecEnv]) 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: Optional[VecNormalize] The `VecNormalize` env. + """ + return self._vec_normalize_env + @staticmethod def check_env(env, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> bool: """