Gymnasium/tests/vector/test_sync_vector_env.py

"""Test the `SyncVectorEnv` implementation."""

import re

import numpy as np
import pytest

from gymnasium.envs.registration import EnvSpec
from gymnasium.spaces import Box, Discrete, MultiDiscrete, Tuple
from gymnasium.vector import SyncVectorEnv
from tests.envs.utils import all_testing_env_specs
from tests.vector.testing_utils import (
    CustomSpace,
    assert_rng_equal,
    make_custom_space_env,
    make_env,
)


def test_create_sync_vector_env():
    """Tests creating the sync vector environment."""
    env_fns = [make_env("FrozenLake-v1", i) for i in range(8)]
    env = SyncVectorEnv(env_fns)
    env.close()

    assert env.num_envs == 8


def test_reset_sync_vector_env():
    """Tests sync vector `reset` function."""
    env_fns = [make_env("CartPole-v1", i) for i in range(8)]
    env = SyncVectorEnv(env_fns)
    observations, infos = env.reset()
    env.close()

    assert isinstance(env.observation_space, Box)
    assert isinstance(observations, np.ndarray)
    assert observations.dtype == env.observation_space.dtype
    assert observations.shape == (8,) + env.single_observation_space.shape
    assert observations.shape == env.observation_space.shape

    del observations


@pytest.mark.parametrize("use_single_action_space", [True, False])
def test_step_sync_vector_env(use_single_action_space):
    """Test sync vector `steps` function."""
    env = SyncVectorEnv([make_env("FrozenLake-v1", i) for i in range(8)])
    env.reset()

    assert isinstance(env.single_action_space, Discrete)
    assert isinstance(env.action_space, MultiDiscrete)

    if use_single_action_space:
        actions = [env.single_action_space.sample() for _ in range(8)]
    else:
        actions = env.action_space.sample()
    observations, rewards, terminations, truncations, _ = env.step(actions)

    env.close()

    assert isinstance(env.observation_space, MultiDiscrete)
    assert isinstance(observations, np.ndarray)
    assert observations.dtype == env.observation_space.dtype
    assert observations.shape == (8,) + env.single_observation_space.shape
    assert observations.shape == env.observation_space.shape

    assert isinstance(rewards, np.ndarray)
    assert isinstance(rewards[0], (float, np.floating))
    assert rewards.ndim == 1
    assert rewards.size == 8

    assert isinstance(terminations, np.ndarray)
    assert terminations.dtype == np.bool_
    assert terminations.ndim == 1
    assert terminations.size == 8

    assert isinstance(truncations, np.ndarray)
    assert truncations.dtype == np.bool_
    assert truncations.ndim == 1
    assert truncations.size == 8


def test_render_sync_vector():
    envs = SyncVectorEnv(
        [make_env("CartPole-v1", i, render_mode="rgb_array") for i in range(3)]
    )
    assert envs.render_mode == "rgb_array"

    envs.reset()
    rendered_frames = envs.render()
    assert isinstance(rendered_frames, tuple)
    assert len(rendered_frames) == envs.num_envs
    assert all(isinstance(frame, np.ndarray) for frame in rendered_frames)

    envs = SyncVectorEnv([make_env("CartPole-v1", i) for i in range(3)])
    assert envs.render_mode is None


def test_call_sync_vector_env():
    """Test sync vector `call` on sub-environments."""
    env_fns = [
        make_env("CartPole-v1", i, render_mode="rgb_array_list") for i in range(4)
    ]

    env = SyncVectorEnv(env_fns)
    _ = env.reset()
    images = env.call("render")
    gravity = env.call("gravity")

    env.close()

    assert isinstance(images, tuple)
    assert len(images) == 4
    for i in range(4):
        assert len(images[i]) == 1
        assert isinstance(images[i][0], np.ndarray)

    assert isinstance(gravity, tuple)
    assert len(gravity) == 4
    for i in range(4):
        assert isinstance(gravity[i], float)
        assert gravity[i] == 9.8


def test_set_attr_sync_vector_env():
    """Test sync vector `set_attr` function."""
    env_fns = [make_env("CartPole-v1", i) for i in range(4)]

    env = SyncVectorEnv(env_fns)
    env.set_attr("gravity", [9.81, 3.72, 8.87, 1.62])
    gravity = env.get_attr("gravity")
    assert gravity == (9.81, 3.72, 8.87, 1.62)

    env.close()


def test_check_spaces_sync_vector_env():
    """Tests the sync vector `check_spaces` function."""
    # CartPole-v1 - observation_space: Box(4,), action_space: Discrete(2)
    env_fns = [make_env("CartPole-v1", i) for i in range(8)]
    # FrozenLake-v1 - Discrete(16), action_space: Discrete(4)
    env_fns[1] = make_env("FrozenLake-v1", 1)
    with pytest.raises(
        AssertionError,
        match=re.escape(
            "SyncVectorEnv(..., observation_mode='same') however the sub-environments observation spaces are not equivalent."
        ),
    ):
        env = SyncVectorEnv(env_fns)
        env.close()


def test_custom_space_sync_vector_env():
    """Test the use of custom spaces with sync vector environment."""
    env_fns = [make_custom_space_env(i) for i in range(4)]

    env = SyncVectorEnv(env_fns)
    reset_observations, infos = env.reset()

    assert isinstance(env.single_action_space, CustomSpace)
    assert isinstance(env.action_space, Tuple)

    actions = ("action-2", "action-3", "action-5", "action-7")
    step_observations, _, _, _, _ = env.step(actions)

    env.close()

    assert isinstance(env.single_observation_space, CustomSpace)
    assert isinstance(env.observation_space, Tuple)

    assert isinstance(reset_observations, tuple)
    assert reset_observations == ("reset", "reset", "reset", "reset")

    assert isinstance(step_observations, tuple)
    assert step_observations == (
        "step(action-2)",
        "step(action-3)",
        "step(action-5)",
        "step(action-7)",
    )


def test_sync_vector_env_seed():
    """Test seeding for sync vector environments."""
    env = make_env("BipedalWalker-v3", seed=123)()
    sync_vector_env = SyncVectorEnv([make_env("BipedalWalker-v3", seed=123)])

    assert_rng_equal(env.action_space.np_random, sync_vector_env.action_space.np_random)
    for _ in range(100):
        env_action = env.action_space.sample()
        vector_action = sync_vector_env.action_space.sample()
        assert np.all(env_action == vector_action)

    env.close()


@pytest.mark.parametrize(
    "spec", all_testing_env_specs, ids=[spec.id for spec in all_testing_env_specs]
)
def test_sync_vector_determinism(spec: EnvSpec, seed: int = 123, n: int = 3):
    """Check that for all environments, the sync vector envs produce the same action samples using the same seeds."""
    env_1 = SyncVectorEnv([make_env(spec.id, seed=seed) for _ in range(n)])
    env_2 = SyncVectorEnv([make_env(spec.id, seed=seed) for _ in range(n)])
    assert_rng_equal(env_1.action_space.np_random, env_2.action_space.np_random)

    for _ in range(100):
        env_1_samples = env_1.action_space.sample()
        env_2_samples = env_2.action_space.sample()
        assert np.all(env_1_samples == env_2_samples)

    env_1.close()
    env_2.close()