Resolve merge conflicts

tests/conftest.py

@@ -6,7 +6,8 @@ import pytest
 import yaml
 
 from primaite import getLogger
-from primaite.game.session import PrimaiteSession
+from primaite.game.game import PrimaiteGame
+from primaite.session.session import PrimaiteSession
 
 # from primaite.environment.primaite_env import Primaite
 # from primaite.primaite_session import PrimaiteSession

tests/e2e_integration_tests/environments/test_rllib_multi_agent_environment.py

@@ -0,0 +1,45 @@
+import pytest
+import ray
+import yaml
+from ray import air, tune
+from ray.rllib.algorithms.ppo import PPOConfig
+
+from primaite.config.load import example_config_path
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteRayMARLEnv
+
+
+@pytest.mark.skip(reason="Slow, reenable later")
+def test_rllib_multi_agent_compatibility():
+    """Test that the PrimaiteRayEnv class can be used with a multi agent RLLIB system."""
+
+    with open(example_config_path(), "r") as f:
+        cfg = yaml.safe_load(f)
+
+    game = PrimaiteGame.from_config(cfg)
+
+    ray.shutdown()
+    ray.init()
+
+    env_config = {"game": game}
+    config = (
+        PPOConfig()
+        .environment(env=PrimaiteRayMARLEnv, env_config={"game": game})
+        .rollouts(num_rollout_workers=0)
+        .multi_agent(
+            policies={agent.agent_name for agent in game.rl_agents},
+            policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,
+        )
+        .training(train_batch_size=128)
+    )
+
+    tune.Tuner(
+        "PPO",
+        run_config=air.RunConfig(
+            stop={"training_iteration": 128},
+            checkpoint_config=air.CheckpointConfig(
+                checkpoint_frequency=10,
+            ),
+        ),
+        param_space=config,
+    ).fit()

tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py

@@ -0,0 +1,40 @@
+import tempfile
+from pathlib import Path
+
+import pytest
+import ray
+import yaml
+from ray.rllib.algorithms import ppo
+
+from primaite.config.load import example_config_path
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteRayEnv
+
+
+@pytest.mark.skip(reason="Slow, reenable later")
+def test_rllib_single_agent_compatibility():
+    """Test that the PrimaiteRayEnv class can be used with a single agent RLLIB system."""
+    with open(example_config_path(), "r") as f:
+        cfg = yaml.safe_load(f)
+
+    game = PrimaiteGame.from_config(cfg)
+
+    ray.shutdown()
+    ray.init()
+
+    env_config = {"game": game}
+    config = {
+        "env": PrimaiteRayEnv,
+        "env_config": env_config,
+        "disable_env_checking": True,
+        "num_rollout_workers": 0,
+    }
+
+    algo = ppo.PPO(config=config)
+
+    for i in range(5):
+        result = algo.train()
+
+    save_file = Path(tempfile.gettempdir()) / "ray/"
+    algo.save(save_file)
+    assert save_file.exists()

tests/e2e_integration_tests/environments/test_sb3_environment.py

@@ -0,0 +1,27 @@
+"""Test that we can create a primaite environment and train sb3 agent with no crash."""
+import tempfile
+from pathlib import Path
+
+import yaml
+from stable_baselines3 import PPO
+
+from primaite.config.load import example_config_path
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteGymEnv
+
+
+def test_sb3_compatibility():
+    """Test that the Gymnasium environment can be used with an SB3 agent."""
+    with open(example_config_path(), "r") as f:
+        cfg = yaml.safe_load(f)
+
+    game = PrimaiteGame.from_config(cfg)
+    gym = PrimaiteGymEnv(game=game)
+    model = PPO("MlpPolicy", gym)
+
+    model.learn(total_timesteps=1000)
+
+    save_path = Path(tempfile.gettempdir()) / "model.zip"
+    model.save(save_path)
+
+    assert (save_path).exists()

tests/e2e_integration_tests/test_primaite_session.py

@@ -8,6 +8,7 @@ CFG_PATH = TEST_ASSETS_ROOT / "configs/test_primaite_session.yaml"
 TRAINING_ONLY_PATH = TEST_ASSETS_ROOT / "configs/train_only_primaite_session.yaml"
 EVAL_ONLY_PATH = TEST_ASSETS_ROOT / "configs/eval_only_primaite_session.yaml"
 MISCONFIGURED_PATH = TEST_ASSETS_ROOT / "configs/bad_primaite_session.yaml"
+MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
 
 
 class TestPrimaiteSession:
@@ -19,15 +20,15 @@ class TestPrimaiteSession:
                 raise AssertionError
 
             assert session is not None
-            assert session.simulation
-            assert len(session.agents) == 3
-            assert len(session.rl_agents) == 1
+            assert session.game.simulation
+            assert len(session.game.agents) == 3
+            assert len(session.game.rl_agents) == 1
 
             assert session.policy
             assert session.env
 
-            assert session.simulation.network
-            assert len(session.simulation.network.nodes) == 10
+            assert session.game.simulation.network
+            assert len(session.game.simulation.network.nodes) == 10
 
     @pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
     def test_start_session(self, temp_primaite_session):
@@ -64,6 +65,13 @@ class TestPrimaiteSession:
             session.start_session()
             # TODO: include checks that the model was loaded and that the eval-only session ran
 
+    @pytest.mark.skip(reason="Slow, reenable later")
+    @pytest.mark.parametrize("temp_primaite_session", [[MULTI_AGENT_PATH]], indirect=True)
+    def test_multi_agent_session(self, temp_primaite_session):
+        """Check that we can run a training session with a multi agent system."""
+        with temp_primaite_session as session:
+            session.start_session()
+
     def test_error_thrown_on_bad_configuration(self):
         with pytest.raises(pydantic.ValidationError):
             session = TempPrimaiteSession.from_config(MISCONFIGURED_PATH)
@@ -72,12 +80,12 @@ class TestPrimaiteSession:
     def test_session_sim_reset(self, temp_primaite_session):
         with temp_primaite_session as session:
             session: TempPrimaiteSession
-            client_1 = session.simulation.network.get_node_by_hostname("client_1")
+            client_1 = session.game.simulation.network.get_node_by_hostname("client_1")
             client_1.software_manager.uninstall("DataManipulationBot")
 
             assert "DataManipulationBot" not in client_1.software_manager.software
 
-            session.reset()
-            client_1 = session.simulation.network.get_node_by_hostname("client_1")
+            session.game.reset()
+            client_1 = session.game.simulation.network.get_node_by_hostname("client_1")
 
             assert "DataManipulationBot" in client_1.software_manager.software