Merged PR 246: Make flattening observation spaces optional.

## Summary
Simple change that gives the user an option to turn on/off observation space flattening.

## Test process
Confirmed that if the new setting is set to 'false', SB3 complains about composite spaces, but RL agents notebook works fine.

## Checklist
- [x] PR is linked to a **work item**
- [x] **acceptance criteria** of linked ticket are met
- [x] performed **self-review** of the code
- [ ] written **tests** for any new functionality added with this PR
- [x] updated the **documentation** if this PR changes or adds functionality
- [ ] written/updated **design docs** if this PR implements new functionality
- [x] updated the **change log**
- [x] ran **pre-commit** checks for code style
- [x] attended to any **TO-DOs** left in the code

Make flattening observation spaces optional.

Related work items: #2161

CHANGELOG.md

@@ -6,7 +6,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
-- Made packet capture and system logging optional
+- Made packet capture and system logging optional (off by default). To turn on, change the io_settings.save_pcap_logs and io_settings.save_sys_logs settings in the config.
+- Made observation space flattening optional (on by default). To turn off for an agent, change the agent_settings.flatten_obs setting in the config.
 
 
 ### Added

docs/source/config.rst

@@ -13,7 +13,25 @@ This section allows selecting which training framework and algorithm to use, and
 
 ``io_settings``
 ---------------
-This section configures how the ``PrimaiteSession`` saves data.
+This section configures how PrimAITE saves data during simulation and training.
+
+**save_final_model**: Only used if training with PrimaiteSession, if true, the policy will be saved after the final training iteration.
+
+**save_checkpoints**: Only used if training with PrimaiteSession, if true, the policy will be saved periodically during training.
+
+**checkpoint_interval**: Only used if training with PrimaiteSession and if ``save_checkpoints`` is true. Defines how often to save the policy during training.
+
+**save_logs**: *currently unused*.
+
+**save_transactions**: *currently unused*.
+
+**save_tensorboard_logs**: *currently unused*.
+
+**save_step_metadata**: Whether to save the RL agents' action, environment state, and other data at every single step.
+
+**save_pcap_logs**: Whether to save pcap files of all network traffic during the simulation.
+
+**save_sys_logs**: Whether to save system logs from all nodes during the simulation.
 
 ``game``
 --------
@@ -56,6 +74,10 @@ Description of configurable items:
 **agent_settings**:
 Settings passed to the agent during initialisation. These depend on the agent class.
 
+Reinforcement learning agents use the ``ProxyAgent`` class, they accept these agent settings:
+
+**flatten_obs**: If true, gymnasium flattening will be performed on the observation space before sending to the agent. Set this to true if your agent does not support nested observation spaces.
+
 ``simulation``
 --------------
 In this section the network layout is defined. This part of the config follows a hierarchical structure. Almost every component defines a ``ref`` field which acts as a human-readable unique identifier, used by other parts of the config, such as agents.

src/primaite/config/_package_data/example_config.yaml

@@ -525,7 +525,7 @@ agents:
 
 
     agent_settings:
-      # ...
+      flatten_obs: true
 
 
 

src/primaite/game/agent/interface.py

@@ -44,6 +44,7 @@ class AgentSettings(BaseModel):
 
     start_settings: Optional[AgentStartSettings] = None
     "Configuration for when an agent begins performing it's actions"
+    flatten_obs: bool = True
 
     @classmethod
     def from_config(cls, config: Optional[Dict]) -> "AgentSettings":
@@ -166,6 +167,7 @@ class ProxyAgent(AbstractAgent):
         action_space: Optional[ActionManager],
         observation_space: Optional[ObservationManager],
         reward_function: Optional[RewardFunction],
+        agent_settings: Optional[AgentSettings] = None,
     ) -> None:
         super().__init__(
             agent_name=agent_name,
@@ -174,6 +176,7 @@ class ProxyAgent(AbstractAgent):
             reward_function=reward_function,
         )
         self.most_recent_action: ActType
+        self.flatten_obs: bool = agent_settings.flatten_obs
 
     def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
         """

src/primaite/game/game.py

@@ -427,6 +427,7 @@ class PrimaiteGame:
                     action_space=action_space,
                     observation_space=obs_space,
                     reward_function=rew_function,
+                    agent_settings=agent_settings,
                 )
                 game.agents.append(new_agent)
                 game.rl_agents.append(new_agent)

src/primaite/notebooks/training_example_ray_single_agent.ipynb

@@ -39,6 +39,15 @@
     "#### Create a Ray algorithm and pass it our config."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(cfg['agents'][2]['agent_settings'])"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -76,6 +85,13 @@
     "    param_space=config\n",
     ").fit()\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

src/primaite/session/environment.py

@@ -23,6 +23,7 @@ class PrimaiteGymEnv(gymnasium.Env):
         super().__init__()
         self.game: "PrimaiteGame" = game
         self.agent: ProxyAgent = self.game.rl_agents[0]
+        self.flatten_obs: bool = False
 
     def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
         """Perform a step in the environment."""
@@ -81,13 +82,19 @@ class PrimaiteGymEnv(gymnasium.Env):
     @property
     def observation_space(self) -> gymnasium.Space:
         """Return the observation space of the environment."""
-        return gymnasium.spaces.flatten_space(self.agent.observation_manager.space)
+        if self.agent.flatten_obs:
+            return gymnasium.spaces.flatten_space(self.agent.observation_manager.space)
+        else:
+            return self.agent.observation_manager.space
 
     def _get_obs(self) -> ObsType:
         """Return the current observation."""
-        unflat_space = self.agent.observation_manager.space
+        if not self.agent.flatten_obs:
-        unflat_obs = self.agent.observation_manager.current_observation
+            return self.agent.observation_manager.current_observation
-        return gymnasium.spaces.flatten(unflat_space, unflat_obs)
+        else:
+            unflat_space = self.agent.observation_manager.space
+            unflat_obs = self.agent.observation_manager.current_observation
+            return gymnasium.spaces.flatten(unflat_space, unflat_obs)
 
 
 class PrimaiteRayEnv(gymnasium.Env):

src/primaite/session/session.py

@@ -101,9 +101,9 @@ class PrimaiteSession:
 
         # CREATE ENVIRONMENT
         if sess.training_options.rl_framework == "RLLIB_single_agent":
-            sess.env = PrimaiteRayEnv(env_config={"game": game})
+            sess.env = PrimaiteRayEnv(env_config={"cfg": cfg})
         elif sess.training_options.rl_framework == "RLLIB_multi_agent":
-            sess.env = PrimaiteRayMARLEnv(env_config={"game": game})
+            sess.env = PrimaiteRayMARLEnv(env_config={"cfg": cfg})
         elif sess.training_options.rl_framework == "SB3":
             sess.env = PrimaiteGymEnv(game=game)