diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 4482e99..c7a6443 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,6 +4,29 @@ Changelog ========== +Release 1.3.1a0 (WIP) +--------------------------- + +Breaking Changes: +^^^^^^^^^^^^^^^^^ + +New Features: +^^^^^^^^^^^^^ +- Added ``norm_obs_keys`` param for ``VecNormalize`` wrapper to configure which observation keys to normalize (@kachayev) + +Bug Fixes: +^^^^^^^^^^ + +Deprecations: +^^^^^^^^^^^^^ + +Others: +^^^^^^^ + +Documentation: +^^^^^^^^^^^^^^ + + Release 1.3.0 (2021-10-23) --------------------------- @@ -30,13 +53,13 @@ New Features: - Added methods ``get_distribution`` and ``predict_values`` for ``ActorCriticPolicy`` for A2C/PPO/TRPO (@cyprienc) - Added methods ``forward_actor`` and ``forward_critic`` for ``MlpExtractor`` - Added ``sb3.get_system_info()`` helper function to gather version information relevant to SB3 (e.g., Python and PyTorch version) -- Saved models now store system information where agent was trained, and load functions have ``print_system_info`` parameter to help debugging load issues. +- Saved models now store system information where agent was trained, and load functions have ``print_system_info`` parameter to help debugging load issues Bug Fixes: ^^^^^^^^^^ - Fixed ``dtype`` of observations for ``SimpleMultiObsEnv`` - Allow `VecNormalize` to wrap discrete-observation environments to normalize reward - when observation normalization is disabled. + when observation normalization is disabled - Fixed a bug where ``DQN`` would throw an error when using ``Discrete`` observation and stochastic actions - Fixed a bug where sub-classed observation spaces could not be used - Added ``force_reset`` argument to ``load()`` and ``set_env()`` in order to be able to call ``learn(reset_num_timesteps=False)`` with a new environment @@ -804,4 +827,4 @@ And all the contributors: @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray @tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @JadenTravnik @decodyng @ardabbour @lorenz-h @mschweizer @lorepieri8 @vwxyzjn -@ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro @amy12xx @juancroldan @benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau +@ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro @amy12xx @juancroldan @benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev diff --git a/stable_baselines3/common/vec_env/vec_normalize.py b/stable_baselines3/common/vec_env/vec_normalize.py index 5eae7f5..0448adf 100644 --- a/stable_baselines3/common/vec_env/vec_normalize.py +++ b/stable_baselines3/common/vec_env/vec_normalize.py @@ -1,7 +1,7 @@ import pickle import warnings from copy import deepcopy -from typing import Any, Dict, Union +from typing import Any, Dict, List, Optional, Union import gym import numpy as np @@ -24,6 +24,8 @@ class VecNormalize(VecEnvWrapper): :param clip_reward: Max value absolute for discounted reward :param gamma: discount factor :param epsilon: To avoid division by zero + :param norm_obs_keys: Which keys from observation dict to normalize. + If not specified, all keys will be normalized. """ def __init__( @@ -36,19 +38,21 @@ class VecNormalize(VecEnvWrapper): clip_reward: float = 10.0, gamma: float = 0.99, epsilon: float = 1e-8, + norm_obs_keys: Optional[List[str]] = None, ): VecEnvWrapper.__init__(self, venv) - if norm_obs: - if not isinstance(self.observation_space, (gym.spaces.Box, gym.spaces.Dict)): - raise ValueError("VecNormalize only supports `gym.spaces.Box` and `gym.spaces.Dict` observation spaces") + self.norm_obs = norm_obs + self.norm_obs_keys = norm_obs_keys + # Check observation spaces + if self.norm_obs: + self._sanity_checks() if isinstance(self.observation_space, gym.spaces.Dict): - self.obs_keys = set(self.observation_space.spaces.keys()) self.obs_spaces = self.observation_space.spaces - self.obs_rms = {key: RunningMeanStd(shape=space.shape) for key, space in self.obs_spaces.items()} + self.obs_rms = {key: RunningMeanStd(shape=self.obs_spaces[key].shape) for key in self.norm_obs_keys} else: - self.obs_keys, self.obs_spaces = None, None + self.obs_spaces = None self.obs_rms = RunningMeanStd(shape=self.observation_space.shape) self.ret_rms = RunningMeanStd(shape=()) @@ -64,6 +68,34 @@ class VecNormalize(VecEnvWrapper): self.old_obs = np.array([]) self.old_reward = np.array([]) + def _sanity_checks(self) -> None: + """ + Check the observations that are going to be normalized are of the correct type (spaces.Box). + """ + if isinstance(self.observation_space, gym.spaces.Dict): + # By default, we normalize all keys + if self.norm_obs_keys is None: + self.norm_obs_keys = list(self.observation_space.spaces.keys()) + # Check that all keys are of type Box + for obs_key in self.norm_obs_keys: + if not isinstance(self.observation_space.spaces[obs_key], gym.spaces.Box): + raise ValueError( + f"VecNormalize only supports `gym.spaces.Box` observation spaces but {obs_key} " + f"is of type {self.observation_space.spaces[obs_key]}. " + "You should probably explicitely pass the observation keys " + " that should be normalized via the `norm_obs_keys` parameter." + ) + + elif isinstance(self.observation_space, gym.spaces.Box): + if self.norm_obs_keys is not None: + raise ValueError("`norm_obs_keys` param is applicable only with `gym.spaces.Dict` observation spaces") + + else: + raise ValueError( + "VecNormalize only supports `gym.spaces.Box` and `gym.spaces.Dict` observation spaces, " + f"not {self.observation_space}" + ) + def __getstate__(self) -> Dict[str, Any]: """ Gets state for pickling. @@ -84,6 +116,9 @@ class VecNormalize(VecEnvWrapper): User must call set_venv() after unpickling before using. :param state:""" + # Backward compatibility + if "norm_obs_keys" not in state: + state["norm_obs_keys"] = list(state["observation_space"].spaces.keys()) self.__dict__.update(state) assert "venv" not in state self.venv = None @@ -170,7 +205,8 @@ class VecNormalize(VecEnvWrapper): obs_ = deepcopy(obs) if self.norm_obs: if isinstance(obs, dict) and isinstance(self.obs_rms, dict): - for key in self.obs_rms.keys(): + # Only normalize the specified keys + for key in self.norm_obs_keys: obs_[key] = self._normalize_obs(obs[key], self.obs_rms[key]).astype(np.float32) else: obs_ = self._normalize_obs(obs, self.obs_rms).astype(np.float32) @@ -190,7 +226,7 @@ class VecNormalize(VecEnvWrapper): obs_ = deepcopy(obs) if self.norm_obs: if isinstance(obs, dict) and isinstance(self.obs_rms, dict): - for key in self.obs_rms.keys(): + for key in self.norm_obs_keys: obs_[key] = self._unnormalize_obs(obs[key], self.obs_rms[key]) else: obs_ = self._unnormalize_obs(obs, self.obs_rms) diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index f0bb29e..e18a0e5 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.3.0 +1.3.1a0 diff --git a/tests/test_dict_env.py b/tests/test_dict_env.py index 7d936c5..34d9dbd 100644 --- a/tests/test_dict_env.py +++ b/tests/test_dict_env.py @@ -66,7 +66,7 @@ class DummyDictEnv(gym.Env): def step(self, action): reward = 0.0 - done = False + done = np.random.rand() > 0.8 return self.observation_space.sample(), reward, done, {} def compute_reward(self, achieved_goal, desired_goal, info): @@ -266,8 +266,8 @@ def test_vec_normalize(model_class): Additional tests for PPO/A2C/SAC/DDPG/TD3/DQN to check observation space support for GoalEnv and VecNormalize using MultiInputPolicy. """ - env = DummyVecEnv([lambda: BitFlippingEnv(n_bits=4, continuous=not (model_class == DQN))]) - env = VecNormalize(env) + env = DummyVecEnv([lambda: DummyDictEnv(use_discrete_actions=model_class == DQN)]) + env = VecNormalize(env, norm_obs_keys=["vec"]) kwargs = {} n_steps = 256 diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 659174b..b002928 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -1,3 +1,5 @@ +import operator + import gym import numpy as np import pytest @@ -68,6 +70,31 @@ class DummyDictEnv(gym.GoalEnv): return -(distance > 0).astype(np.float32) +class DummyMixedDictEnv(gym.Env): + """ + Dummy mixed gym env for testing purposes + """ + + def __init__(self): + super().__init__() + self.observation_space = spaces.Dict( + { + "obs1": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32), + "obs2": spaces.Discrete(1), + "obs3": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32), + } + ) + self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32) + + def reset(self): + return self.observation_space.sample() + + def step(self, action): + obs = self.observation_space.sample() + done = np.random.rand() > 0.8 + return obs, 0.0, done, {} + + def allclose(obs_1, obs_2): """ Generalized np.allclose() to work with dict spaces. @@ -152,9 +179,9 @@ def _make_warmstart_cartpole(): return _make_warmstart(lambda: gym.make("CartPole-v1")) -def _make_warmstart_dict_env(): - """Warm-start VecNormalize by stepping through BitFlippingEnv""" - return _make_warmstart(make_dict_env) +def _make_warmstart_dict_env(**kwargs): + """Warm-start VecNormalize by stepping through DummyDictEnv""" + return _make_warmstart(make_dict_env, **kwargs) def test_runningmeanstd(): @@ -267,6 +294,21 @@ def test_normalize_external(): assert np.all(norm_rewards < 1) +def test_normalize_dict_selected_keys(): + venv = _make_warmstart_dict_env(norm_obs=True, norm_obs_keys=["observation"]) + for _ in range(3): + actions = [venv.action_space.sample()] + obs, rewards, _, _ = venv.step(actions) + orig_obs = venv.get_original_obs() + + # "observation" is expected to be normalized + np.testing.assert_array_compare(operator.__ne__, obs["observation"], orig_obs["observation"]) + assert allclose(venv.normalize_obs(orig_obs), obs) + + # other keys are expected to be presented "as is" + np.testing.assert_array_equal(obs["achieved_goal"], orig_obs["achieved_goal"]) + + @pytest.mark.parametrize("model_class", [SAC, TD3, HerReplayBuffer]) @pytest.mark.parametrize("online_sampling", [False, True]) def test_offpolicy_normalization(model_class, online_sampling): @@ -358,3 +400,14 @@ def test_discrete_obs(): # Smoke test that it runs with norm_obs False _make_warmstart_cliffwalking(norm_obs=False) + + +def test_non_dict_obs_keys(): + with pytest.raises(ValueError, match=".*is applicable only.*"): + _make_warmstart(lambda: DummyRewardEnv(), norm_obs_keys=["key"]) + + with pytest.raises(ValueError, match=".* explicitely pass the observation keys.*"): + _make_warmstart(lambda: DummyMixedDictEnv()) + + # Ignore Discrete observation key + _make_warmstart(lambda: DummyMixedDictEnv(), norm_obs_keys=["obs1", "obs3"])