oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix typehints

Browse Source
This commit is contained in:
Marek Wolan
2023-11-16 15:28:38 +00:00
parent 13c49bf3ea
commit 0b9bdedebd
5 changed files with 8 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/primaite/game/agent/rewards.py
View File

@@ -26,7 +26,7 @@ the structure:
``` ```
""" """
from abc import abstractmethod from abc import abstractmethod
from typing import Dict, List, Tuple, TYPE_CHECKING from typing import Dict, List, Tuple, Type, TYPE_CHECKING
from primaite import getLogger from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
@@ -228,7 +228,7 @@ class WebServer404Penalty(AbstractReward):
class RewardFunction: class RewardFunction:
"""Manages the reward function for the agent.""" """Manages the reward function for the agent."""
__rew_class_identifiers: Dict[str, type[AbstractReward]] = { __rew_class_identifiers: Dict[str, Type[AbstractReward]] = {
"DUMMY": DummyReward, "DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity, "DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty, "WEB_SERVER_404_PENALTY": WebServer404Penalty,

4
src/primaite/game/policy/policy.py
View File

@@ -1,7 +1,7 @@
"""Base class and common logic for RL policies.""" """Base class and common logic for RL policies."""
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pathlib import Path from pathlib import Path
from typing import Any, Dict, TYPE_CHECKING from typing import Any, Dict, Type, TYPE_CHECKING
if TYPE_CHECKING: if TYPE_CHECKING:
from primaite.game.session import PrimaiteSession, TrainingOptions from primaite.game.session import PrimaiteSession, TrainingOptions
@@ -10,7 +10,7 @@ if TYPE_CHECKING:
class PolicyABC(ABC): class PolicyABC(ABC):
"""Base class for reinforcement learning agents.""" """Base class for reinforcement learning agents."""
_registry: Dict[str, type["PolicyABC"]] = {} _registry: Dict[str, Type["PolicyABC"]] = {}
""" """
Registry of policy types, keyed by name. Registry of policy types, keyed by name.

4
src/primaite/game/policy/sb3.py
View File

@@ -1,6 +1,6 @@
"""Stable baselines 3 policy.""" """Stable baselines 3 policy."""
from pathlib import Path from pathlib import Path
from typing import Literal, Optional, TYPE_CHECKING, Union from typing import Literal, Optional, Type, TYPE_CHECKING, Union
from stable_baselines3 import A2C, PPO from stable_baselines3 import A2C, PPO
from stable_baselines3.a2c import MlpPolicy as A2C_MLP from stable_baselines3.a2c import MlpPolicy as A2C_MLP
@@ -21,7 +21,7 @@ class SB3Policy(PolicyABC, identifier="SB3"):
"""Initialize a stable baselines 3 policy.""" """Initialize a stable baselines 3 policy."""
super().__init__(session=session) super().__init__(session=session)
self._agent_class: type[Union[PPO, A2C]] self._agent_class: Type[Union[PPO, A2C]]
if algorithm == "PPO": if algorithm == "PPO":
self._agent_class = PPO self._agent_class = PPO
policy = PPO_MLP policy = PPO_MLP

4
src/primaite/game/session.py
View File

@@ -52,7 +52,7 @@ class PrimaiteGymEnv(gymnasium.Env):
self.session: "PrimaiteSession" = session self.session: "PrimaiteSession" = session
self.agent: ProxyAgent = agents[0] self.agent: ProxyAgent = agents[0]
def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, dict[str, Any]]: def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
"""Perform a step in the environment.""" """Perform a step in the environment."""
# make ProxyAgent store the action chosen my the RL policy # make ProxyAgent store the action chosen my the RL policy
self.agent.store_action(action) self.agent.store_action(action)
@@ -70,7 +70,7 @@ class PrimaiteGymEnv(gymnasium.Env):
return next_obs, reward, terminated, truncated, info return next_obs, reward, terminated, truncated, info
def reset(self, seed: Optional[int] = None) -> tuple[ObsType, dict[str, Any]]: def reset(self, seed: Optional[int] = None) -> Tuple[ObsType, Dict[str, Any]]:
"""Reset the environment.""" """Reset the environment."""
self.session.reset() self.session.reset()
state = self.session.get_sim_state() state = self.session.get_sim_state()

2
tests/e2e_integration_tests/test_primaite_session.py
View File

@@ -40,8 +40,6 @@ class TestPrimaiteSession:
with temp_primaite_session as session: with temp_primaite_session as session:
session: TempPrimaiteSession session: TempPrimaiteSession
session.start_session() session.start_session()
for i in range(100):
print(session.io_manager.generate_session_path())
# TODO: include checks that the model was trained, e.g. that the loss changed and checkpoints were saved? # TODO: include checks that the model was trained, e.g. that the loss changed and checkpoints were saved?
@pytest.mark.parametrize("temp_primaite_session", [[EVAL_ONLY_PATH]], indirect=True) @pytest.mark.parametrize("temp_primaite_session", [[EVAL_ONLY_PATH]], indirect=True)