Fix bug in gym to gymnasium upgrade (#40)

- Explicitly use gymnasium instead of gym everywhere (docs and imports)
- Fix a bug in multi-agent handler, trying to access env.done, which was removed
- Added required infos key for RLlib

docs/components.md

@@ -47,7 +47,7 @@ config['movement'] = LazyUEMovement
 config['movement_params'].update({'lazy_ue': 3})
 
 # create environment with lazy UE!
-env = gym.make('mobile-small-central-v0', config=config)
+env = gymnasium.make('mobile-small-central-v0', config=config)
 ```
 
 This example demonstrates that each core component can be replaced **without changing source code**. It also shows that we can pass parameters to the components' initialization via the ``config`` dictionary. Adapting the channel model, etc., works similarly. For example, replacing the default [Okumura–Hata](https://en.wikipedia.org/wiki/Hata_model) channel model by a (simplified) path loss model can be as easy as this:
@@ -79,7 +79,7 @@ config['channel'] = PathLoss
 config['channel_params'].update({'gamma': 2.0})
 
 # create environment with custom channel model
-env = gym.make('mobile-small-central-v0', config=config)
+env = gymnasium.make('mobile-small-central-v0', config=config)
 ...
 ```
 
@@ -110,7 +110,7 @@ class CustomHandler(Handler):
         ...
 
 config = {'handler': CustomHandler}
-env = gym.make('mobile-small-central-v0', config=config)
+env = gymnasium.make('mobile-small-central-v0', config=config)
 ```
 
 So far, [mobile-env](https://mobile-env.readthedocs.io/en/latest/index.html) implements handlers for a multi-agent and centralized control setting.
@@ -129,7 +129,7 @@ Example usage of our **RLlibMAWrapper** class:
 ```python
 from mobile_env.wrappers.multi_agent import RLlibMAWrapper
 
-env = gym.make('mobile-small-ma-v0')
+env = gymnasium.make('mobile-small-ma-v0')
 # wrap multi-agent env for RLlib compatibility
 env = RLlibMAWrapper(env)
 # use RLlib to train on the environment ...

examples/test.ipynb

@@ -69,13 +69,13 @@
     }
    ],
    "source": [
-    "import gymnasium as gym\n",
+    "import gymnasium\n",
     "import matplotlib.pyplot as plt\n",
-    "# importing mobile_env automatically registers the predefined scenarios in Gym\n",
+    "# importing mobile_env automatically registers the predefined scenarios in Gymnasium\n",
     "import mobile_env\n",
     "\n",
     "# create a small mobile environment for a single, centralized control agent\n",
-    "env = gym.make(\"mobile-small-central-v0\", render_mode='rgb_array')\n",
+    "env = gymnasium.make(\"mobile-small-central-v0\", render_mode='rgb_array')\n",
     "env.reset()\n",
     "\n",
     "# select a random action: control cell connections for each UE\n",

mobile_env/core/base.py

@@ -2,7 +2,7 @@
 from collections import Counter, defaultdict
 from typing import Dict, List, Set, Tuple
 
-import gymnasium as gym
+import gymnasium
 import matplotlib.patheffects as pe
 import matplotlib.pyplot as plt
 import numpy as np
@@ -23,7 +23,7 @@
 from mobile_env.handlers.central import MComCentralHandler
 
 
-class MComCore(gym.Env):
+class MComCore(gymnasium.Env):
     NOOP_ACTION = 0
     metadata = {"render_modes": ["rgb_array", "human"]}
 

mobile_env/handlers/multi_agent.py

@@ -1,6 +1,6 @@
 from typing import Dict
 
-import gymnasium as gym
+import gymnasium
 import numpy as np
 
 from mobile_env.handlers.handler import Handler
@@ -20,23 +20,23 @@ def ue_obs_size(cls, env) -> int:
         return sum(env.feature_sizes[ftr] for ftr in cls.features)
 
     @classmethod
-    def action_space(cls, env) -> gym.spaces.Dict:
-        return gym.spaces.Dict(
+    def action_space(cls, env) -> gymnasium.spaces.Dict:
+        return gymnasium.spaces.Dict(
             {
-                ue.ue_id: gym.spaces.Discrete(env.NUM_STATIONS + 1)
+                ue.ue_id: gymnasium.spaces.Discrete(env.NUM_STATIONS + 1)
                 for ue in env.users.values()
             }
         )
 
     @classmethod
-    def observation_space(cls, env) -> gym.spaces.Dict:
+    def observation_space(cls, env) -> gymnasium.spaces.Dict:
         size = cls.ue_obs_size(env)
         space = {
-            ue_id: gym.spaces.Box(low=-1, high=1, shape=(size,), dtype=np.float32)
+            ue_id: gymnasium.spaces.Box(low=-1, high=1, shape=(size,), dtype=np.float32)
             for ue_id in env.users
         }
 
-        return gym.spaces.Dict(space)
+        return gymnasium.spaces.Dict(space)
 
     @classmethod
     def reward(cls, env):
@@ -68,7 +68,7 @@ def observation(cls, env) -> Dict[int, np.ndarray]:
         """Select features for MA setting & flatten each UE's features."""
 
         # get features for currently active UEs
-        active = set([ue.ue_id for ue in env.active if not env.done])
+        active = set([ue.ue_id for ue in env.active if not env.time_is_up])
         features = env.features()
         features = {ue_id: obs for ue_id, obs in features.items() if ue_id in active}
 

mobile_env/scenarios/registry.py

@@ -1,6 +1,6 @@
 import itertools
 
-import gymnasium as gym
+import gymnasium
 
 from mobile_env.handlers.central import MComCentralHandler
 from mobile_env.handlers.multi_agent import MComMAHandler
@@ -14,7 +14,7 @@
 for scenario, handler in itertools.product(scenarios, handlers):
     env_name = scenarios[scenario].__name__
     config = {"handler": handlers[handler]}
-    gym.envs.register(
+    gymnasium.envs.register(
         id=f"mobile-{scenario}-{handler}-v0",
         entry_point=f"mobile_env.scenarios.{scenario}:{env_name}",
         kwargs={"config": config},

mobile_env/wrappers/multi_agent.py

@@ -1,6 +1,6 @@
 from typing import Tuple
 
-import gymnasium as gym
+import gymnasium
 import numpy as np
 from ray.rllib.env.multi_agent_env import MultiAgentEnv
 from ray.rllib.utils.typing import MultiAgentDict
@@ -19,9 +19,9 @@ def __init__(self, env):
         # override action and observation space defined for wrapped environment
         # RLlib expects the action and observation space
         # to be defined per actor, i.e, per UE
-        self.action_space = gym.spaces.Discrete(self.env.NUM_STATIONS + 1)
+        self.action_space = gymnasium.spaces.Discrete(self.env.NUM_STATIONS + 1)
         size = self.env.handler.ue_obs_size(self.env)
-        self.observation_space = gym.spaces.Box(
+        self.observation_space = gymnasium.spaces.Box(
             low=-1, high=1, shape=(size,), dtype=np.float32
         )
 
@@ -57,6 +57,10 @@ def step(
         # update keys of previous observation dictionary
         self.prev_step_ues = set(obs.keys())
 
+        # RLlib expects the keys of infos to be a subset of obs + __common__
+        # Put all infos under __common__
+        infos = {"__common__": infos}
+
         return obs, rews, terminateds, truncateds, infos
 
     def render(self) -> None:

setup.py

@@ -20,7 +20,7 @@
 
 setup(
     name="mobile-env",
-    version="2.0.0",
+    version="2.0.1",
     author="Stefan Schneider, Stefan Werner",
     description="mobile-env: An Open Environment for Autonomous Coordination in "
     "Wireless Mobile Networks",

tests/test_central_envs.py

@@ -1,20 +1,20 @@
 import unittest
 
-import gymnasium as gym
+import gymnasium
 from stable_baselines3.common.env_checker import check_env
 
 import mobile_env  # noqa: F401
 
 
 class TestCentralEnvs(unittest.TestCase):
     def test_central_small(self):
-        check_env(gym.make("mobile-small-central-v0"))
+        check_env(gymnasium.make("mobile-small-central-v0"))
 
     def test_central_medium(self):
-        check_env(gym.make("mobile-medium-central-v0"))
+        check_env(gymnasium.make("mobile-medium-central-v0"))
 
     def test_central_large(self):
-        check_env(gym.make("mobile-large-central-v0"))
+        check_env(gymnasium.make("mobile-large-central-v0"))
 
 
 if __name__ == "__main__":

tests/test_env_stepping.py

@@ -1,9 +1,9 @@
 """Simple test of small env similar to test notebook."""
-import gymnasium as gym
+import gymnasium
 import pytest
 
 # importing mobile_env automatically registers the predefined scenarios in Gym
-import mobile_env
+import mobile_env  # noqa: F401
 
 
 @pytest.mark.parametrize(
@@ -16,7 +16,7 @@ def test_env_stepping(env_name):
     Just to ensure that it does not crash.
     """
     # create a small mobile environment for a single, centralized control agent
-    env = gym.make(env_name)
+    env = gymnasium.make(env_name)
     obs, info = env.reset()
     done = False