small nits

Author: matteobettini

tests/test_wrappers/__init__.py

@@ -0,0 +1,3 @@
+#  Copyright (c) 2024.
+#  ProrokLab (https://www.proroklab.org/)
+#  All rights reserved.

tests/test_wrappers/test_gym_wrapper.py

@@ -1,3 +1,7 @@
+#  Copyright (c) 2024.
+#  ProrokLab (https://www.proroklab.org/)
+#  All rights reserved.
+
 import gym
 import numpy as np
 import pytest
@@ -10,16 +14,13 @@
 TEST_SCENARIOS = [
     "balance",
     "discovery",
-    "dispersion",
-    "football",
     "give_way",
     "joint_passage",
     "navigation",
     "passage",
-    "reverse_transport",
-    "road_traffic",
     "transport",
     "waterfall",
+    "simple_world_comm",
 ]
 
 
@@ -93,7 +94,10 @@ def test_gym_wrapper(
     _check_obs_type(obss, obs_shapes, dict_space, return_numpy=return_numpy)
 
     for _ in range(max_steps):
-        actions = env.unwrapped.get_random_actions()
+        actions = [
+            env.unwrapped.get_random_action(agent).numpy()
+            for agent in env.unwrapped.agents
+        ]
         obss, rews, done, info = env.step(actions)
         _check_obs_type(obss, obs_shapes, dict_space, return_numpy=return_numpy)
 

tests/test_wrappers/test_gymnasium_vec_wrapper.py

@@ -68,7 +68,10 @@ def test_gymnasium_wrapper(
     ), f"Expected info to be a dictionary but got {type(info)}"
 
     for _ in range(max_steps):
-        actions = env.unwrapped.get_random_actions()
+        actions = [
+            env.unwrapped.get_random_action(agent).numpy()
+            for agent in env.unwrapped.agents
+        ]
         obss, rews, terminated, truncated, info = env.step(actions)
         _check_obs_type(obss, obs_shapes, dict_space, return_numpy=return_numpy)
 

tests/test_wrappers/test_gymnasium_wrapper.py

@@ -65,7 +65,10 @@ def test_gymnasium_wrapper(
     ), f"Expected info to be a dictionary but got {type(info)}"
 
     for _ in range(max_steps):
-        actions = env.unwrapped.get_random_actions()
+        actions = [
+            env.unwrapped.get_random_action(agent).numpy()
+            for agent in env.unwrapped.agents
+        ]
         obss, rews, terminated, truncated, info = env.step(actions)
         _check_obs_type(obss, obs_shapes, dict_space, return_numpy=return_numpy)