mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-27 08:47:08 +00:00
89 lines
2.3 KiB
Python
89 lines
2.3 KiB
Python
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
import gym
|
||
|
|
from gym.spaces import Discrete
|
||
|
|
from gym.vector import AsyncVectorEnv, SyncVectorEnv
|
||
|
|
|
||
|
|
|
||
|
|
class OldStepEnv(gym.Env):
|
||
|
|
def __init__(self):
|
||
|
|
self.action_space = Discrete(2)
|
||
|
|
self.observation_space = Discrete(2)
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def step(self, action):
|
||
|
|
obs = self.observation_space.sample()
|
||
|
|
rew = 0
|
||
|
|
done = False
|
||
|
|
info = {}
|
||
|
|
return obs, rew, done, info
|
||
|
|
|
||
|
|
|
||
|
|
class NewStepEnv(gym.Env):
|
||
|
|
def __init__(self):
|
||
|
|
self.action_space = Discrete(2)
|
||
|
|
self.observation_space = Discrete(2)
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def step(self, action):
|
||
|
|
obs = self.observation_space.sample()
|
||
|
|
rew = 0
|
||
|
|
terminated = False
|
||
|
|
truncated = False
|
||
|
|
info = {}
|
||
|
|
return obs, rew, terminated, truncated, info
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("VecEnv", [AsyncVectorEnv, SyncVectorEnv])
|
||
|
|
def test_vector_step_compatibility_new_env(VecEnv):
|
||
|
|
|
||
|
|
envs = [
|
||
|
|
OldStepEnv(),
|
||
|
|
NewStepEnv(),
|
||
|
|
]
|
||
|
|
|
||
|
|
vec_env = VecEnv([lambda: env for env in envs])
|
||
|
|
vec_env.reset()
|
||
|
|
step_returns = vec_env.step([0, 0])
|
||
|
|
assert len(step_returns) == 4
|
||
|
|
_, _, dones, _ = step_returns
|
||
|
|
assert dones.dtype == np.bool_
|
||
|
|
vec_env.close()
|
||
|
|
|
||
|
|
vec_env = VecEnv([lambda: env for env in envs], new_step_api=True)
|
||
|
|
vec_env.reset()
|
||
|
|
step_returns = vec_env.step([0, 0])
|
||
|
|
assert len(step_returns) == 5
|
||
|
|
_, _, terminateds, truncateds, _ = step_returns
|
||
|
|
assert terminateds.dtype == np.bool_
|
||
|
|
assert truncateds.dtype == np.bool_
|
||
|
|
vec_env.close()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("async_bool", [True, False])
|
||
|
|
def test_vector_step_compatibility_existing(async_bool):
|
||
|
|
|
||
|
|
env = gym.vector.make("CartPole-v1", num_envs=3, asynchronous=async_bool)
|
||
|
|
env.reset()
|
||
|
|
step_returns = env.step(env.action_space.sample())
|
||
|
|
assert len(step_returns) == 4
|
||
|
|
_, _, dones, _ = step_returns
|
||
|
|
assert dones.dtype == np.bool_
|
||
|
|
env.close()
|
||
|
|
|
||
|
|
env = gym.vector.make(
|
||
|
|
"CartPole-v1", num_envs=3, asynchronous=async_bool, new_step_api=True
|
||
|
|
)
|
||
|
|
env.reset()
|
||
|
|
step_returns = env.step(env.action_space.sample())
|
||
|
|
assert len(step_returns) == 5
|
||
|
|
_, _, terminateds, truncateds, _ = step_returns
|
||
|
|
assert terminateds.dtype == np.bool_
|
||
|
|
assert truncateds.dtype == np.bool_
|
||
|
|
env.close()
|