Add get_vec_normalize_env()

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:
 ^^^^^^^^^^

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():

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:
         """