Add test for wrapping gymnasium environments in torch

tests/torch/test_wrapper_gymnasium.py

@@ -0,0 +1,80 @@
+import pytest
+
+from collections.abc import Mapping
+import gymnasium as gym
+
+import torch
+
+from skrl.envs.wrappers.torch import GymnasiumWrapper, wrap_env
+
+
+def test_env(capsys: pytest.CaptureFixture):
+    num_envs = 1
+    action = torch.ones((num_envs, 1))
+
+    # load wrap the environment
+    original_env = gym.make("Pendulum-v1")
+    env = wrap_env(original_env, "gymnasium")
+    assert isinstance(env, GymnasiumWrapper)
+
+    # check properties
+    assert env.state_space is None
+    assert isinstance(env.observation_space, gym.Space) and env.observation_space.shape == (3,)
+    assert isinstance(env.action_space, gym.Space) and env.action_space.shape == (1,)
+    assert isinstance(env.num_envs, int) and env.num_envs == num_envs
+    assert isinstance(env.num_agents, int) and env.num_agents == 1
+    assert isinstance(env.device, torch.device)
+    # check internal properties
+    assert env._env is original_env
+    assert env._unwrapped is original_env.unwrapped
+    # check methods
+    for _ in range(2):
+        observation, info = env.reset()
+        assert isinstance(observation, torch.Tensor) and observation.shape == torch.Size([num_envs, 3])
+        assert isinstance(info, Mapping)
+        for _ in range(3):
+            observation, reward, terminated, truncated, info = env.step(action)
+            env.render()
+            assert isinstance(observation, torch.Tensor) and observation.shape == torch.Size([num_envs, 3])
+            assert isinstance(reward, torch.Tensor) and reward.shape == torch.Size([num_envs, 1])
+            assert isinstance(terminated, torch.Tensor) and terminated.shape == torch.Size([num_envs, 1])
+            assert isinstance(truncated, torch.Tensor) and truncated.shape == torch.Size([num_envs, 1])
+            assert isinstance(info, Mapping)
+
+    env.close()
+
+def test_vectorized_env(capsys: pytest.CaptureFixture):
+    num_envs = 10
+    action = torch.ones((num_envs, 1))
+
+    # load wrap the environment
+    original_env = gym.make_vec("Pendulum-v1", num_envs=num_envs)
+    env = wrap_env(original_env, "gymnasium")
+    assert isinstance(env, GymnasiumWrapper)
+
+    # check properties
+    assert env.state_space is None
+    assert isinstance(env.observation_space, gym.Space) and env.observation_space.shape == (3,)
+    assert isinstance(env.action_space, gym.Space) and env.action_space.shape == (1,)
+    assert isinstance(env.num_envs, int) and env.num_envs == num_envs
+    assert isinstance(env.num_agents, int) and env.num_agents == 1
+    assert isinstance(env.device, torch.device)
+    # check internal properties
+    assert env._env is original_env
+    assert env._unwrapped is original_env.unwrapped
+    assert env._vectorized is True
+    # check methods
+    for _ in range(2):
+        observation, info = env.reset()
+        assert isinstance(observation, torch.Tensor) and observation.shape == torch.Size([num_envs, 3])
+        assert isinstance(info, Mapping)
+        for _ in range(3):
+            observation, reward, terminated, truncated, info = env.step(action)
+            env.render()
+            assert isinstance(observation, torch.Tensor) and observation.shape == torch.Size([num_envs, 3])
+            assert isinstance(reward, torch.Tensor) and reward.shape == torch.Size([num_envs, 1])
+            assert isinstance(terminated, torch.Tensor) and terminated.shape == torch.Size([num_envs, 1])
+            assert isinstance(truncated, torch.Tensor) and truncated.shape == torch.Size([num_envs, 1])
+            assert isinstance(info, Mapping)
+
+    env.close()