[Tests] Improve tests (#118)

Author: matteobettini

tests/test_vmas.py

@@ -36,7 +36,9 @@ def test_all_scenarios_included():
 
 @pytest.mark.parametrize("scenario", scenario_names())
 @pytest.mark.parametrize("continuous_actions", [True, False])
-def test_use_vmas_env(scenario, continuous_actions, num_envs=10, n_steps=10):
+def test_use_vmas_env(
+    scenario, continuous_actions, dict_spaces=True, num_envs=10, n_steps=10
+):
     render = True
     if sys.platform.startswith("win32"):
         # Windows on github servers has issues with pyglet
@@ -51,9 +53,36 @@ def test_use_vmas_env(scenario, continuous_actions, num_envs=10, n_steps=10):
         continuous_actions=continuous_actions,
         num_envs=num_envs,
         n_steps=n_steps,
+        dict_spaces=dict_spaces,
     )
 
 
+@pytest.mark.parametrize("scenario", scenario_names())
+def test_multi_discrete_actions(scenario, num_envs=10, n_steps=10):
+    env = make_env(
+        scenario=scenario,
+        num_envs=num_envs,
+        seed=0,
+        multidiscrete_actions=True,
+        continuous_actions=False,
+    )
+    for _ in range(n_steps):
+        env.step(env.get_random_actions())
+
+
+@pytest.mark.parametrize("scenario", scenario_names())
+def test_non_dict_spaces_actions(scenario, num_envs=10, n_steps=10):
+    env = make_env(
+        scenario=scenario,
+        num_envs=num_envs,
+        seed=0,
+        continuous_actions=True,
+        dict_spaces=False,
+    )
+    for _ in range(n_steps):
+        env.step(env.get_random_actions())
+
+
 @pytest.mark.parametrize("scenario", scenario_names())
 def test_partial_reset(scenario, num_envs=10, n_steps=10):
     env = make_env(
@@ -70,6 +99,19 @@ def test_partial_reset(scenario, num_envs=10, n_steps=10):
             env_index = 0
 
 
+@pytest.mark.parametrize("scenario", scenario_names())
+def test_global_reset(scenario, num_envs=10, n_steps=10):
+    env = make_env(
+        scenario=scenario,
+        num_envs=num_envs,
+        seed=0,
+    )
+    for step in range(n_steps):
+        env.step(env.get_random_actions())
+        if step == n_steps // 2:
+            env.reset()
+
+
 @pytest.mark.parametrize("scenario", vmas.scenarios + vmas.mpe_scenarios)
 def test_vmas_differentiable(scenario, n_steps=10, n_envs=10):
     if scenario == "football" or scenario == "simple_crypto":

vmas/simulator/environment/environment.py

@@ -406,12 +406,25 @@ def get_random_action(self, agent: Agent) -> torch.Tensor:
                     )
             action = torch.stack(actions, dim=-1)
         else:
-            action = torch.randint(
-                low=0,
-                high=self.get_agent_action_space(agent).n,
-                size=(agent.batch_dim,),
-                device=agent.device,
-            )
+            action_space = self.get_agent_action_space(agent)
+            if self.multidiscrete_actions:
+                actions = [
+                    torch.randint(
+                        low=0,
+                        high=action_space.nvec[action_index],
+                        size=(agent.batch_dim,),
+                        device=agent.device,
+                    )
+                    for action_index in range(action_space.shape[0])
+                ]
+                action = torch.stack(actions, dim=-1)
+            else:
+                action = torch.randint(
+                    low=0,
+                    high=action_space.n,
+                    size=(agent.batch_dim,),
+                    device=agent.device,
+                )
         return action
 
     def get_random_actions(self) -> Sequence[torch.Tensor]: