Skip to content

Commit afdbecb

Browse files
committed
Fix
1 parent 2d98fc8 commit afdbecb

File tree

4 files changed

+28
-5
lines changed

4 files changed

+28
-5
lines changed

gymnasium/wrappers/vector/vectorize_action.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def __init__(
8585
# TODO: We could compute single_action_space from the action_space if only the latter is provided and avoid the warning below.
8686
if self.action_space != batch_space(self.single_action_space, self.num_envs):
8787
warn(
88-
"The action space and the batched single action space don't match as expected."
88+
f"For {env}, the action space and the batched single action space don't match as expected, action_space={env.action_space}, batched single_action_space={batch_space(self.single_action_space, self.num_envs)}"
8989
)
9090

9191
self.func = func

gymnasium/wrappers/vector/vectorize_observation.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ def __init__(
8484
self.single_observation_space, self.num_envs
8585
):
8686
warn(
87-
"The observation space and the batched single observation space don't match as expected."
87+
f"For {env}, the observation space and the batched single observation space don't match as expected, observation_space={env.observation_space}, batched single_observation_space={batch_space(self.single_observation_space, self.num_envs)}"
8888
)
8989

9090
self.func = func

tests/wrappers/vector/test_transform_action.py

+21-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Test suite for vector TransformAction wrapper."""
22

33
import numpy as np
4+
import pytest
45

56
from gymnasium import spaces, wrappers
67
from gymnasium.vector import SyncVectorEnv
@@ -16,7 +17,7 @@ def create_env():
1617
)
1718

1819

19-
def test_observation_space_from_single_observation_space(
20+
def test_action_space_from_single_action_space(
2021
n_envs: int = 5,
2122
):
2223
vec_env = SyncVectorEnv([create_env for _ in range(n_envs)])
@@ -51,3 +52,22 @@ def test_observation_space_from_single_observation_space(
5152
assert (
5253
vec_env.single_action_space.high == np.array([110, 95, 110], dtype=np.float32)
5354
).all()
55+
56+
57+
def test_warning_on_mismatched_single_action_space(
58+
n_envs: int = 2,
59+
):
60+
vec_env = SyncVectorEnv([create_env for _ in range(n_envs)])
61+
# We only specify action_space without single_action_space, so single_action_space inherits its value from the wrapped env which would not match. This mismatch should give us a warning.
62+
with pytest.warns(
63+
Warning,
64+
match=r"the action space and the batched single action space don't match as expected",
65+
):
66+
vec_env = wrappers.vector.TransformAction(
67+
vec_env,
68+
func=lambda x: x + 100,
69+
action_space=spaces.Box(
70+
low=np.array([[0, -10, -5]] * n_envs, dtype=np.float32) + 100,
71+
high=np.array([[10, -5, 10]] * n_envs, dtype=np.float32) + 100,
72+
),
73+
)

tests/wrappers/vector/test_transform_observation.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,14 @@ def test_observation_space_from_single_observation_space(
8282

8383

8484
def test_warning_on_mismatched_single_observation_space(
85-
n_envs: int = 5,
85+
n_envs: int = 2,
8686
):
8787
vec_env = SyncVectorEnv([create_env for _ in range(n_envs)])
8888
# We only specify observation_space without single_observation_space, so single_observation_space inherits its value from the wrapped env which would not match. This mismatch should give us a warning.
89-
with pytest.warns(Warning):
89+
with pytest.warns(
90+
Warning,
91+
match=r"the observation space and the batched single observation space don't match as expected",
92+
):
9093
vec_env = wrappers.vector.TransformObservation(
9194
vec_env,
9295
func=lambda x: x + 100,

0 commit comments

Comments
 (0)