Adding state function to multiwalker

pettingzoo/sisl/multiwalker/multiwalker.py

@@ -120,6 +120,7 @@
 """
 
 import numpy as np
+from gymnasium import spaces
 from gymnasium.utils import EzPickle
 
 from pettingzoo import AECEnv
@@ -160,6 +161,14 @@ def __init__(self, *args, **kwargs):
         # spaces
         self.action_spaces = dict(zip(self.agents, self.env.action_space))
         self.observation_spaces = dict(zip(self.agents, self.env.observation_space))
+        self.state_space = spaces.Box(
+            low=-np.float32(np.inf),
+            high=+np.float32(np.inf),
+            shape=(
+                self.env.n_walkers * 24 + 3,
+            ),  # 24 is the observation space of each walker, 3 is the package observation space
+            dtype=np.float32,
+        )
         self.steps = 0
 
     def observation_space(self, agent):
@@ -239,3 +248,17 @@ def step(self, action):
 
         if self.render_mode == "human":
             self.render()
+
+    def state(self):
+        all_walker_obs = self.env.get_last_obs()
+        all_walker_obs = np.array(list(all_walker_obs.values())).flatten()
+        package_obs = np.array(
+            [
+                self.env.package.position.x,
+                self.env.package.position.y,
+                self.env.package.angle,
+            ]
+        )
+        global_state = np.concatenate((all_walker_obs, package_obs)).astype(np.float32)
+
+        return global_state