From a60b0179e0904491298ec89c0ebf0ab7211a2299 Mon Sep 17 00:00:00 2001 From: Omar Younis <42100908+younik@users.noreply.github.com> Date: Wed, 29 Mar 2023 16:25:05 +0200 Subject: [PATCH] Fix: Reshape action in DictRolloutBuffer (#1395) * reshape action in DictRolloutBuffer * improve buffer test * update changelog * add comment * Update comments and version --------- Co-authored-by: Antonin Raffin --- docs/misc/changelog.rst | 3 ++- stable_baselines3/common/buffers.py | 9 ++++++--- stable_baselines3/version.txt | 2 +- tests/test_buffers.py | 3 ++- 4 files changed, 11 insertions(+), 6 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e96e26e..5c4dc42 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Release 1.8.0a11 (WIP) +Release 1.8.0a12 (WIP) -------------------------- .. warning:: @@ -46,6 +46,7 @@ Bug Fixes: - Added the argument ``dtype`` (default to ``float32``) to the noise for consistency with gym action (@sidney-tio) - Fixed PPO train/n_updates metric not accounting for early stopping (@adamfrly) - Fixed loading of normalized image-based environments +- Fixed `DictRolloutBuffer.add` with multidimensional action space (@younik) Deprecations: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index f14d2f9..95a9c1e 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -246,7 +246,7 @@ class ReplayBuffer(BaseBuffer): obs = obs.reshape((self.n_envs, *self.obs_shape)) next_obs = next_obs.reshape((self.n_envs, *self.obs_shape)) - # Same, for actions + # Reshape to handle multi-dim and discrete action spaces, see GH #970 #1392 action = action.reshape((self.n_envs, self.action_dim)) # Copy to avoid modification by reference @@ -430,7 +430,7 @@ class RolloutBuffer(BaseBuffer): if isinstance(self.observation_space, spaces.Discrete): obs = obs.reshape((self.n_envs, *self.obs_shape)) - # Same reshape, for actions + # Reshape to handle multi-dim and discrete action spaces, see GH #970 #1392 action = action.reshape((self.n_envs, self.action_dim)) self.observations[self.pos] = np.array(obs).copy() @@ -588,7 +588,7 @@ class DictReplayBuffer(ReplayBuffer): next_obs[key] = next_obs[key].reshape((self.n_envs,) + self.obs_shape[key]) self.next_observations[key][self.pos] = np.array(next_obs[key]).copy() - # Same reshape, for actions + # Reshape to handle multi-dim and discrete action spaces, see GH #970 #1392 action = action.reshape((self.n_envs, self.action_dim)) self.actions[self.pos] = np.array(action).copy() @@ -741,6 +741,9 @@ class DictRolloutBuffer(RolloutBuffer): obs_ = obs_.reshape((self.n_envs,) + self.obs_shape[key]) self.observations[key][self.pos] = obs_ + # Reshape to handle multi-dim and discrete action spaces, see GH #970 #1392 + action = action.reshape((self.n_envs, self.action_dim)) + self.actions[self.pos] = np.array(action).copy() self.rewards[self.pos] = np.array(reward).copy() self.episode_starts[self.pos] = np.array(episode_start).copy() diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 01d49ad..15159be 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a11 +1.8.0a12 diff --git a/tests/test_buffers.py b/tests/test_buffers.py index 0e028e6..9dc294c 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -44,7 +44,8 @@ class DummyDictEnv(gym.Env): """ def __init__(self): - self.action_space = spaces.Box(1, 5, (1,)) + # Test for multi-dim action space + self.action_space = spaces.Box(1, 5, shape=(10, 7)) space = spaces.Box(1, 5, (1,)) self.observation_space = spaces.Dict({"observation": space, "achieved_goal": space, "desired_goal": space}) self._observations = [1, 2, 3, 4, 5]