saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix returned type in predict (#964)

Browse source 
* `arr[0]` to `arr.squeeze(0)`

* `squeeze(axis=0)` to `squeeze(0)`

* Type testing

* Add type test for unvectorized observation

* `squeeze(0)` to `squeeze(axis=0)`

* Treatment of the laziness symptoms

* Update changelog

* Udate changelog

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
This commit is contained in:
Quentin Gallouédec 2022-07-18 11:22:19 +02:00 committed by GitHub
parent a18b91e01a
commit fda3d4d748
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 7 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
docs/misc/changelog.rst
View file

@ -17,6 +17,7 @@ SB3-Contrib
Bug Fixes: Bug Fixes:
^^^^^^^^^^ ^^^^^^^^^^
- Fixed the issue that ``predict`` does not always return action as ``np.ndarray`` (@qgallouedec)
Deprecations: Deprecations:
^^^^^^^^^^^^^ ^^^^^^^^^^^^^
@ -1011,4 +1012,4 @@ And all the contributors:
@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 @buoyancy99 @caburu @xy9485 @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485
@Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede
@Melanol @Melanol @qgallouedec

4
stable_baselines3/common/distributions.py
View file

@ -578,10 +578,10 @@ class StateDependentNoiseDistribution(Distribution):
return th.mm(latent_sde, self.exploration_mat) return th.mm(latent_sde, self.exploration_mat)
# Use batch matrix multiplication for efficient computation # Use batch matrix multiplication for efficient computation
# (batch_size, n_features) -> (batch_size, 1, n_features) # (batch_size, n_features) -> (batch_size, 1, n_features)
latent_sde = latent_sde.unsqueeze(1) latent_sde = latent_sde.unsqueeze(dim=1)
# (batch_size, 1, n_actions) # (batch_size, 1, n_actions)
noise = th.bmm(latent_sde, self.exploration_matrices) noise = th.bmm(latent_sde, self.exploration_matrices)
return noise.squeeze(1) return noise.squeeze(dim=1)
def actions_from_params( def actions_from_params(
self, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor, deterministic: bool = False self, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor, deterministic: bool = False

2
stable_baselines3/common/policies.py
View file

@ -350,7 +350,7 @@ class BasePolicy(BaseModel):
# Remove batch dimension if needed # Remove batch dimension if needed
if not vectorized_env: if not vectorized_env:
actions = actions[0] actions = actions.squeeze(axis=0)
return actions, state return actions, state

2
tests/test_predict.py
View file

@ -73,11 +73,13 @@ def test_predict(model_class, env_id, device):
obs = env.reset() obs = env.reset()
action, _ = model.predict(obs) action, _ = model.predict(obs)
assert isinstance(action, np.ndarray)
assert action.shape == env.action_space.shape assert action.shape == env.action_space.shape
assert env.action_space.contains(action) assert env.action_space.contains(action)
vec_env_obs = vec_env.reset() vec_env_obs = vec_env.reset()
action, _ = model.predict(vec_env_obs) action, _ = model.predict(vec_env_obs)
assert isinstance(action, np.ndarray)
assert action.shape[0] == vec_env_obs.shape[0] assert action.shape[0] == vec_env_obs.shape[0]
# Special case for DQN to check the epsilon greedy exploration # Special case for DQN to check the epsilon greedy exploration