Move policy module into session

src/primaite/game/policy/__init__.py

@@ -1,4 +0,0 @@
-from primaite.game.policy.rllib import RaySingleAgentPolicy
-from primaite.game.policy.sb3 import SB3Policy
-
-__all__ = ["SB3Policy", "RaySingleAgentPolicy"]

src/primaite/session/policy/__init__.py

@@ -0,0 +1,4 @@
+from primaite.session.policy.rllib import RaySingleAgentPolicy
+from primaite.session.policy.sb3 import SB3Policy
+
+__all__ = ["SB3Policy", "RaySingleAgentPolicy"]

src/primaite/session/policy/rllib.py

@@ -1,8 +1,8 @@
 from pathlib import Path
 from typing import Literal, Optional, TYPE_CHECKING
 
-from primaite.game.policy.policy import PolicyABC
 from primaite.session.environment import PrimaiteRayEnv, PrimaiteRayMARLEnv
+from primaite.session.policy.policy import PolicyABC
 
 if TYPE_CHECKING:
     from primaite.session.session import PrimaiteSession, TrainingOptions

src/primaite/session/policy/sb3.py

@@ -8,7 +8,7 @@ from stable_baselines3.common.callbacks import CheckpointCallback
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.ppo import MlpPolicy as PPO_MLP
 
-from primaite.game.policy.policy import PolicyABC
+from primaite.session.policy.policy import PolicyABC
 
 if TYPE_CHECKING:
     from primaite.session.session import PrimaiteSession, TrainingOptions

src/primaite/session/session.py

@@ -5,12 +5,12 @@ from typing import Dict, List, Literal, Optional, Union
 from pydantic import BaseModel, ConfigDict
 
 from primaite.game.game import PrimaiteGame
-
-# from primaite.game.game import PrimaiteGame
-from primaite.game.policy.policy import PolicyABC
 from primaite.session.environment import PrimaiteGymEnv, PrimaiteRayEnv, PrimaiteRayMARLEnv
 from primaite.session.io import SessionIO, SessionIOSettings
 
+# from primaite.game.game import PrimaiteGame
+from primaite.session.policy.policy import PolicyABC
+
 
 class TrainingOptions(BaseModel):
     """Options for training the RL agent."""