fix test problems and slowness

tests/e2e_integration_tests/environments/test_rllib_multi_agent_environment.py

@@ -1,7 +1,5 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
-import ray
 import yaml
-from ray import air, tune
 from ray.rllib.algorithms.ppo import PPOConfig
 
 from primaite.session.ray_envs import PrimaiteRayMARLEnv
@@ -12,7 +10,6 @@ MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
 
 def test_rllib_multi_agent_compatibility():
     """Test that the PrimaiteRayEnv class can be used with a multi agent RLLIB system."""
-
     with open(MULTI_AGENT_PATH, "r") as f:
         cfg = yaml.safe_load(f)
 
@@ -26,14 +23,5 @@ def test_rllib_multi_agent_compatibility():
         )
         .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()
+    algo = config.build()
+    algo.train()

tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py

@@ -3,7 +3,6 @@ import tempfile
 from pathlib import Path
 
 import pytest
-import ray
 import yaml
 from ray.rllib.algorithms import ppo
 

tests/e2e_integration_tests/environments/test_sb3_environment.py

@@ -20,7 +20,7 @@ def test_sb3_compatibility():
     gym = PrimaiteGymEnv(env_config=cfg)
     model = PPO("MlpPolicy", gym)
 
-    model.learn(total_timesteps=1000)
+    model.learn(total_timesteps=256)
 
     save_path = Path(tempfile.gettempdir()) / "model.zip"
     model.save(save_path)