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>
This commit is contained in:
Boyuan Chen 2022-02-23 06:33:41 -05:00 committed by GitHub
parent d2ebd2eeaa
commit 7a01637128
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 16 additions and 12 deletions

View file

@ -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, - 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 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.) 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: Deprecations:
^^^^^^^^^^^^^ ^^^^^^^^^^^^^
@ -921,4 +922,4 @@ And all the contributors:
@benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc @benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc
@wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev @cleversonahum @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev @cleversonahum
@eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP @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

View file

@ -48,12 +48,12 @@ class VecNormalize(VecEnvWrapper):
if self.norm_obs: if self.norm_obs:
self._sanity_checks() self._sanity_checks()
if isinstance(self.observation_space, gym.spaces.Dict): if isinstance(self.observation_space, gym.spaces.Dict):
self.obs_spaces = self.observation_space.spaces self.obs_spaces = self.observation_space.spaces
self.obs_rms = {key: RunningMeanStd(shape=self.obs_spaces[key].shape) for key in self.norm_obs_keys} self.obs_rms = {key: RunningMeanStd(shape=self.obs_spaces[key].shape) for key in self.norm_obs_keys}
else: else:
self.obs_spaces = None self.obs_spaces = None
self.obs_rms = RunningMeanStd(shape=self.observation_space.shape) self.obs_rms = RunningMeanStd(shape=self.observation_space.shape)
self.ret_rms = RunningMeanStd(shape=()) self.ret_rms = RunningMeanStd(shape=())
self.clip_obs = clip_obs self.clip_obs = clip_obs
@ -150,7 +150,7 @@ class VecNormalize(VecEnvWrapper):
self.old_obs = obs self.old_obs = obs
self.old_reward = rewards self.old_reward = rewards
if self.training: if self.training and self.norm_obs:
if isinstance(obs, dict) and isinstance(self.obs_rms, dict): if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
for key in self.obs_rms.keys(): for key in self.obs_rms.keys():
self.obs_rms[key].update(obs[key]) self.obs_rms[key].update(obs[key])
@ -258,7 +258,7 @@ class VecNormalize(VecEnvWrapper):
obs = self.venv.reset() obs = self.venv.reset()
self.old_obs = obs self.old_obs = obs
self.returns = np.zeros(self.num_envs) 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): if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
for key in self.obs_rms.keys(): for key in self.obs_rms.keys():
self.obs_rms[key].update(obs[key]) self.obs_rms[key].update(obs[key])

View file

@ -1 +1 @@
1.4.1a1 1.4.1a2

View file

@ -451,3 +451,6 @@ def test_non_dict_obs_keys():
# Ignore Discrete observation key # Ignore Discrete observation key
_make_warmstart(lambda: DummyMixedDictEnv(), norm_obs_keys=["obs1", "obs3"]) _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)