Merge branch 'main' into football

Author: matteobettini

.github/unittest/install_dependencies.sh

@@ -7,7 +7,7 @@
 
 python -m pip install --upgrade pip
 
-pip install -e .
+pip install -e ".[gymnasium]"
 
 python -m pip install flake8 pytest pytest-cov tqdm matplotlib==3.8
 python -m pip install cvxpylayers # Navigation heuristic

README.md

@@ -39,7 +39,7 @@
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3/", None),
     "sphinx": ("https://www.sphinx-doc.org/en/master/", None),
-    "torch": ("https://pytorch.org/docs/master", None),
+    "torch": ("https://pytorch.org/docs/stable/", None),
     "torchrl": ("https://pytorch.org/rl/stable/", None),
     "tensordict": ("https://pytorch.org/tensordict/stable", None),
     "benchmarl": ("https://benchmarl.readthedocs.io/en/latest/", None),

docs/source/conf.py

@@ -2,6 +2,9 @@
 VMAS
 ====
 
+.. discord_button::
+   https://discord.gg/dg8txxDW5t
+
 .. figure:: https://raw.githubusercontent.com/matteobettini/vmas-media/master/media/vmas_scenarios_more.gif
    :align: center
 

docs/source/index.rst

@@ -29,6 +29,21 @@ Install optional requirements
 By default, vmas has only the core requirements.
 Here are some optional packages you may want to install.
 
+Wrappers
+^^^^^^^^
+
+If you want to use VMAS environment wrappers, you may want to install VMAS
+with the following options:
+
+.. code-block:: console
+
+   # install gymnasium for gymnasium wrapper
+   pip install vmas[gymnasium]
+
+   # install rllib for rllib wrapper
+   pip install vmas[rllib]
+
+
 Training
 ^^^^^^^^
 
@@ -40,12 +55,14 @@ You may want to install one of the following training libraries
    pip install torchrl
    pip install "ray[rllib]"==2.1.0 # We support versions "ray[rllib]<=2.2,>=1.13"
 
-Logging
-^^^^^^^
+Utils
+^^^^^
 
-You may want to install the following rendering and logging tools
+You may want to install the following additional tools
 
 .. code-block:: console
 
-   pip install wandb
-   pip install opencv-python moviepy matplotlib
+   # install rendering dependencies
+   pip install vmas[render]
+   # install testing dependencies
+   pip install vmas[test]

docs/source/usage/installation.rst

@@ -30,5 +30,11 @@ def get_version():
     author_email="[email protected]",
     packages=find_packages(),
     install_requires=["numpy", "torch", "pyglet<=1.5.27", "gym", "six"],
+    extras_require={
+        "gymnasium": ["gymnasium", "shimmy"],
+        "rllib": ["ray[rllib]<=2.2"],
+        "render": ["opencv-python", "moviepy", "matplotlib", "opencv-python"],
+        "test": ["pytest", "pytest-instafail", "pyyaml", "tqdm"],
+    },
     include_package_data=True,
 )

setup.py

@@ -0,0 +1,28 @@
+#  Copyright (c) 2024.
+#  ProrokLab (https://www.proroklab.org/)
+#  All rights reserved.
+
+import torch
+
+from vmas import make_env
+
+
+def test_vectorized_lidar(n_envs=12, n_steps=15):
+    def get_obs(env):
+        rollout_obs = []
+        for _ in range(n_steps):
+            obs, _, _, _ = env.step(env.get_random_actions())
+            obs = torch.stack(obs, dim=-1)
+            rollout_obs.append(obs)
+        return torch.stack(rollout_obs, dim=-1)
+
+    env_vec_lidar = make_env(
+        scenario="pollock", num_envs=n_envs, seed=0, lidar=True, vectorized_lidar=True
+    )
+    obs_vec_lidar = get_obs(env_vec_lidar)
+    env_non_vec_lidar = make_env(
+        scenario="pollock", num_envs=n_envs, seed=0, lidar=True, vectorized_lidar=False
+    )
+    obs_non_vec_lidar = get_obs(env_non_vec_lidar)
+
+    assert torch.allclose(obs_vec_lidar, obs_non_vec_lidar)

tests/test_lidar.py

@@ -24,8 +24,9 @@ def setup_env(
         self.env.seed(0)
 
     @pytest.mark.parametrize("n_agents", [1, 4])
-    def test_heuristic(self, n_agents, n_steps=50, n_envs=4):
-        self.setup_env(n_agents=n_agents, n_envs=n_envs)
+    @pytest.mark.parametrize("agent_lidar", [True, False])
+    def test_heuristic(self, n_agents, agent_lidar, n_steps=50, n_envs=4):
+        self.setup_env(n_agents=n_agents, n_envs=n_envs, use_agent_lidar=agent_lidar)
         policy = discovery.HeuristicPolicy(True)
 
         obs = self.env.reset()

tests/test_scenarios/test_discovery.py

@@ -1,7 +1,8 @@
 #  Copyright (c) 2022-2024.
 #  ProrokLab (https://www.proroklab.org/)
 #  All rights reserved.
-import os
+import math
+import random
 import sys
 from pathlib import Path
 
@@ -16,16 +17,20 @@
 def scenario_names():
     scenarios = []
     scenarios_folder = Path(__file__).parent.parent / "vmas" / "scenarios"
-    for _, _, filenames in os.walk(scenarios_folder):
-        scenarios += filenames
-    scenarios = [
-        scenario.split(".")[0]
-        for scenario in scenarios
-        if scenario.endswith(".py") and not scenario.startswith("__")
-    ]
+    for path in scenarios_folder.glob("**/*.py"):
+        if path.is_file() and not path.name.startswith("__"):
+            scenarios.append(path.stem)
     return scenarios
 
 
+def random_nvecs(count, l_min=2, l_max=6, n_min=2, n_max=6, seed=0):
+    random.seed(seed)
+    return [
+        [random.randint(n_min, n_max) for _ in range(random.randint(l_min, l_max))]
+        for _ in range(count)
+    ]
+
+
 def test_all_scenarios_included():
     from vmas import debug_scenarios, mpe_scenarios, scenarios
 
@@ -70,6 +75,163 @@ def test_multi_discrete_actions(scenario, num_envs=10, n_steps=10):
         env.step(env.get_random_actions())
 
 
+@pytest.mark.parametrize("scenario", scenario_names())
+@pytest.mark.parametrize("multidiscrete_actions", [True, False])
+def test_discrete_action_nvec(scenario, multidiscrete_actions, num_envs=10, n_steps=5):
+    env = make_env(
+        scenario=scenario,
+        num_envs=num_envs,
+        seed=0,
+        multidiscrete_actions=multidiscrete_actions,
+        continuous_actions=False,
+    )
+    if (
+        type(env.scenario).process_action
+        is not vmas.simulator.scenario.BaseScenario.process_action
+    ):
+        pytest.skip("Scenario uses a custom process_action method.")
+
+    random.seed(0)
+    for agent in env.world.agents:
+        agent.discrete_action_nvec = [
+            random.randint(2, 6) for _ in range(agent.action_size)
+        ]
+    env.action_space = env.get_action_space()
+
+    def to_multidiscrete(action, nvec):
+        action_multi = []
+        for i in range(len(nvec)):
+            n = math.prod(nvec[i + 1 :])
+            action_multi.append(action // n)
+            action = action % n
+        return torch.stack(action_multi, dim=-1)
+
+    def full_nvec(agent, world):
+        return list(agent.discrete_action_nvec) + (
+            [world.dim_c] if not agent.silent and world.dim_c != 0 else []
+        )
+
+    for _ in range(n_steps):
+        actions = env.get_random_actions()
+
+        # Check that generated actions are in the action space
+        for a_batch, s in zip(actions, env.action_space.spaces):
+            for a in a_batch:
+                assert a.numpy() in s
+
+        env.step(actions)
+
+        if not multidiscrete_actions:
+            actions = [
+                to_multidiscrete(a.squeeze(-1), full_nvec(agent, env.world))
+                for a, agent in zip(actions, env.world.policy_agents)
+            ]
+
+        # Check that discrete action to continuous control mapping is correct.
+        for i_a, agent in enumerate(env.world.policy_agents):
+            for i, n in enumerate(agent.discrete_action_nvec):
+                a = actions[i_a][:, i]
+                u = agent.action.u[:, i]
+                U = agent.action.u_range_tensor[i]
+                k = agent.action.u_multiplier_tensor[i]
+                for aj, uj in zip(a, u):
+                    assert aj in range(
+                        n
+                    ), f"discrete action {aj} not in [0,{n-1}] (n={n}, U={U}, k={k})"
+                    if n % 2 != 0:
+                        assert (
+                            aj != 0 or uj == 0
+                        ), f"discrete action {aj} maps to control {uj} (n={n}), U={U}, k={k})"
+                        assert (aj < 1 or aj > n // 2) or torch.isclose(
+                            uj / k, (2 * U * (aj - 1)) / (n - 1) - U
+                        ), f"discrete action {aj} maps to control {uj} (n={n}, U={U}, k={k})"
+                        assert (aj <= n // 2) or torch.isclose(
+                            uj / k, 2 * U * (aj / (n - 1)) - U
+                        ), f"discrete action {aj} maps to control {uj} (n={n}), U={U}, k={k})"
+                    else:
+                        assert torch.isclose(
+                            uj / k, 2 * U * (aj / (n - 1)) - U
+                        ), f"discrete action {aj} maps to control {uj} (n={n}), U={U}, k={k})"
+
+
+@pytest.mark.parametrize(
+    "nvecs", list(zip(random_nvecs(10, seed=0), random_nvecs(10, seed=42)))
+)
+def test_discrete_action_nvec_discrete_to_multi(
+    nvecs, scenario="transport", num_envs=10, n_steps=5
+):
+    kwargs = {
+        "scenario": scenario,
+        "num_envs": num_envs,
+        "seed": 0,
+        "continuous_actions": False,
+    }
+    env = make_env(**kwargs, multidiscrete_actions=False)
+    env_multi = make_env(**kwargs, multidiscrete_actions=True)
+    if (
+        type(env.scenario).process_action
+        is not vmas.simulator.scenario.BaseScenario.process_action
+    ):
+        pytest.skip("Scenario uses a custom process_action method.")
+
+    def set_nvec(agent, nvec):
+        agent.action_size = len(nvec)
+        agent.discrete_action_nvec = nvec
+        agent.action.action_size = agent.action_size
+
+    random.seed(0)
+    for agent, agent_multi, nvec in zip(
+        env.world.policy_agents, env_multi.world.policy_agents, nvecs
+    ):
+        set_nvec(agent, nvec)
+        set_nvec(agent_multi, nvec)
+    env.action_space = env.get_action_space()
+    env_multi.action_space = env.get_action_space()
+
+    def full_nvec(agent, world):
+        return list(agent.discrete_action_nvec) + (
+            [world.dim_c] if not agent.silent and world.dim_c != 0 else []
+        )
+
+    def full_action_size(agent, world):
+        return len(full_nvec(agent, world))
+
+    for _ in range(n_steps):
+        actions_multi = env_multi.get_random_actions()
+        prodss = [
+            [
+                math.prod(full_nvec(agent, env.world)[i + 1 :])
+                for i in range(full_action_size(agent, env.world))
+            ]
+            for agent in env.world.policy_agents
+        ]
+        # Compute the expected mapping from multi-discrete to discrete
+        actions = [
+            (a_multi * torch.tensor(prods)).sum(dim=1)
+            for a_multi, prods in zip(actions_multi, prodss)
+        ]
+
+        env_multi.step(actions_multi)
+        env.step(actions)
+
+        # Check that both discrete and multi-discrete actions result in the
+        # same control value
+        for agent, agent_multi, action, action_multi in zip(
+            env.world.policy_agents,
+            env_multi.world.policy_agents,
+            actions,
+            actions_multi,
+        ):
+            U = agent.action.u_range_tensor
+            k = agent.action.u_multiplier_tensor
+            for u, u_multi, a, a_multi in zip(
+                agent.action.u, agent_multi.action.u, action, action_multi
+            ):
+                assert torch.allclose(
+                    u, u_multi
+                ), f"{u} != {u_multi} (nvec={agent.discrete_action_nvec}, a={a}, a_multi={a_multi}, U={U}, k={k})"
+
+
 @pytest.mark.parametrize("scenario", scenario_names())
 def test_non_dict_spaces_actions(scenario, num_envs=10, n_steps=10):
     env = make_env(

tests/test_vmas.py

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