Fix VecNormalization bug for Dict obs (#768)

* fix #724 VecNormalization bug for Dict obs

* update test and changelog

* Update changelog

Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.4.1a1 (WIP)
+Release 1.4.1a2 (WIP)
 ---------------------------
 
 
@@ -27,7 +27,8 @@ Bug Fixes:
 - Fixed a bug in ``HumanOutputFormat``. Distinct keys truncated to the same prefix would overwrite each others value,
   resulting in only one being output. This now raises an error (this should only affect a small fraction of use cases
   with very long keys.)
-- Routing all the nn.Module calls through implicit rather than explict forward as per pytorch guidelines (@manuel-delverme)
+- Routing all the ``nn.Module`` calls through implicit rather than explict forward as per pytorch guidelines (@manuel-delverme)
+- Fixed a bug in ``VecNormalize`` where error occurs when ``norm_obs`` is set to False for environment with dictionary observation  (@buoyancy99)
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -921,4 +922,4 @@ And all the contributors:
 @benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc
 @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev @cleversonahum
 @eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP
-@simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola
+@simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99

stable_baselines3/common/vec_env/vec_normalize.py

@@ -150,7 +150,7 @@ class VecNormalize(VecEnvWrapper):
         self.old_obs = obs
         self.old_reward = rewards
 
-        if self.training:
+        if self.training and self.norm_obs:
             if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
                 for key in self.obs_rms.keys():
                     self.obs_rms[key].update(obs[key])
@@ -258,7 +258,7 @@ class VecNormalize(VecEnvWrapper):
         obs = self.venv.reset()
         self.old_obs = obs
         self.returns = np.zeros(self.num_envs)
-        if self.training:
+        if self.training and self.norm_obs:
             if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
                 for key in self.obs_rms.keys():
                     self.obs_rms[key].update(obs[key])

stable_baselines3/version.txt

@@ -1 +1 @@
-1.4.1a1
+1.4.1a2

tests/test_vec_normalize.py

@@ -451,3 +451,6 @@ def test_non_dict_obs_keys():
 
     # Ignore Discrete observation key
     _make_warmstart(lambda: DummyMixedDictEnv(), norm_obs_keys=["obs1", "obs3"])
+
+    # Test dict obs with norm_obs set to False
+    _make_warmstart(lambda: DummyMixedDictEnv(), norm_obs=False)