Gymnasium/tests/vector/test_step_compatibility_vector.py

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()