2623 Ray single agent action masking

src/primaite/session/environment.py

@@ -1,9 +1,10 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 import json
 from os import PathLike
-from typing import Any, Dict, List, Optional, SupportsFloat, Tuple, Union
+from typing import Any, Dict, Optional, SupportsFloat, Tuple, Union
 
 import gymnasium
+import numpy as np
 from gymnasium.core import ActType, ObsType
 
 from primaite import getLogger
@@ -40,10 +41,8 @@ class PrimaiteGymEnv(gymnasium.Env):
         """Current episode number."""
         self.total_reward_per_episode: Dict[int, float] = {}
         """Average rewards of agents per episode."""
-        self.action_masking: bool = False
-        """Whether to use action masking."""
 
-    def action_masks(self) -> List[bool]:
+    def action_masks(self) -> np.ndarray:
         """
         Return the action mask for the agent.
 
@@ -54,13 +53,13 @@ class PrimaiteGymEnv(gymnasium.Env):
         :rtype: List[bool]
         """
         mask = [True] * len(self.agent.action_manager.action_map)
-        if not self.action_masking:
+        if not self.agent.action_masking:
             return mask
 
         for i, action in self.agent.action_manager.action_map.items():
             request = self.agent.action_manager.form_request(action_identifier=action[0], action_options=action[1])
             mask[i] = self.game.simulation._request_manager.check_valid(request, {})
-        return mask
+        return np.asarray(mask)
 
     @property
     def agent(self) -> ProxyAgent:

src/primaite/session/ray_envs.py

@@ -3,6 +3,7 @@ import json
 from typing import Dict, SupportsFloat, Tuple
 
 import gymnasium
+from gymnasium import spaces
 from gymnasium.core import ActType, ObsType
 from ray.rllib.env.multi_agent_env import MultiAgentEnv
 
@@ -38,15 +39,10 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
 
         self.terminateds = set()
         self.truncateds = set()
-        self.observation_space = gymnasium.spaces.Dict(
-            {
-                name: gymnasium.spaces.flatten_space(agent.observation_manager.space)
-                for name, agent in self.agents.items()
-            }
-        )
-        self.action_space = gymnasium.spaces.Dict(
-            {name: agent.action_manager.space for name, agent in self.agents.items()}
+        self.observation_space = spaces.Dict(
+            {name: spaces.flatten_space(agent.observation_manager.space) for name, agent in self.agents.items()}
         )
+        self.action_space = spaces.Dict({name: agent.action_manager.space for name, agent in self.agents.items()})
         self._obs_space_in_preferred_format = True
         self._action_space_in_preferred_format = True
         super().__init__()
@@ -158,15 +154,30 @@ class PrimaiteRayEnv(gymnasium.Env):
         self.env = PrimaiteGymEnv(env_config=env_config)
         # self.env.episode_counter -= 1
         self.action_space = self.env.action_space
-        self.observation_space = self.env.observation_space
+        if self.env.agent.action_masking:
+            self.observation_space = spaces.Dict(
+                {"action_mask": spaces.MultiBinary(self.env.action_space.n), "observations": self.env.observation_space}
+            )
+        else:
+            self.observation_space = self.env.observation_space
 
     def reset(self, *, seed: int = None, options: dict = None) -> Tuple[ObsType, Dict]:
         """Reset the environment."""
+        if self.env.agent.action_masking:
+            obs, *_ = self.env.reset(seed=seed)
+            new_obs = {"action_mask": self.env.action_masks(), "observations": obs}
+            return new_obs, *_
         return self.env.reset(seed=seed)
 
     def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict]:
         """Perform a step in the environment."""
-        return self.env.step(action)
+        # if action masking is enabled, intercept the step method and add action mask to observation
+        if self.env.agent.action_masking:
+            obs, *_ = self.env.step(action)
+            new_obs = {"action_mask": self.env.action_masks(), "observations": obs}
+            return new_obs, *_
+        else:
+            return self.env.step(action)
 
     def close(self):
         """Close the simulation."""