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

#917 - Almost there. All output files being writen for SB3/RLLIB PPO & A2C. Just need to bring in the hardcoded agents then update the testa and docs.

Browse Source
This commit is contained in:
Chris McCarthy
2023-06-19 21:53:25 +01:00
parent 23bafde457
commit 03ae4884e0
14 changed files with 321 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/primaite/agents/agent.py
View File

@@ -5,19 +5,18 @@ from pathlib import Path
from typing import Optional, Final, Dict, Union, List
from uuid import uuid4
from primaite import getLogger
from primaite import getLogger, SESSIONS_DIR
from primaite.common.enums import OutputVerboseLevel
from primaite.config import lay_down_config
from primaite.config import training_config
from primaite.config.training_config import TrainingConfig
from primaite.environment.primaite_env import Primaite
from primaite.transactions.transactions_to_file import \
write_transaction_to_file
_LOGGER = getLogger(__name__)
def _get_temp_session_path(session_timestamp: datetime) -> Path:
def _get_session_path(session_timestamp: datetime) -> Path:
"""
Get a temp directory session path the test session will output to.
@@ -26,7 +25,7 @@ def _get_temp_session_path(session_timestamp: datetime) -> Path:
"""
date_dir = session_timestamp.strftime("%Y-%m-%d")
session_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
session_path = Path("./") / date_dir / session_path
session_path = SESSIONS_DIR / date_dir / session_path
session_path.mkdir(exist_ok=True, parents=True)
return session_path
@@ -57,16 +56,16 @@ class AgentSessionABC(ABC):
self._env: Primaite
self._agent = None
self._transaction_list: List[Dict] = []
self._can_learn: bool = False
self._can_evaluate: bool = False
self._uuid = str(uuid4())
self.session_timestamp: datetime = datetime.now()
"The session timestamp"
self.session_path = _get_temp_session_path(self.session_timestamp)
self.session_path = _get_session_path(self.session_timestamp)
"The Session path"
self.checkpoints_path = self.session_path / "checkpoints"
self.checkpoints_path.mkdir(parents=True, exist_ok=True)
"The Session checkpoints path"
self.timestamp_str = self.session_timestamp.strftime(
@@ -167,11 +166,6 @@ class AgentSessionABC(ABC):
):
if self._can_learn:
_LOGGER.debug("Writing transactions")
write_transaction_to_file(
transaction_list=self._transaction_list,
session_path=self.session_path,
timestamp_str=self.timestamp_str,
)
self._update_session_metadata_file()
self._can_evaluate = True

97
src/primaite/agents/rllib.py
View File

@@ -1,18 +1,23 @@
import glob
import time
from enum import Enum
import json
from datetime import datetime
from pathlib import Path
from typing import Union, Optional
from pathlib import Path
from typing import Optional
from ray.rllib.algorithms import Algorithm
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.algorithms.a2c import A2CConfig
from ray.tune.logger import UnifiedLogger
from ray.tune.registry import register_env
from primaite import getLogger
from primaite.agents.agent import AgentSessionABC
from primaite.config import training_config
from primaite.common.enums import AgentFramework, RedAgentIdentifier
from primaite.environment.primaite_env import Primaite
_LOGGER = getLogger(__name__)
def _env_creator(env_config):
return Primaite(
training_config_path=env_config["training_config_path"],
@@ -23,7 +28,17 @@ def _env_creator(env_config):
)
class RLlibPPO(AgentSessionABC):
def _custom_log_creator(session_path: Path):
logdir = session_path / "ray_results"
logdir.mkdir(parents=True, exist_ok=True)
def logger_creator(config):
return UnifiedLogger(config, logdir, loggers=None)
return logger_creator
class RLlibAgent(AgentSessionABC):
def __init__(
self,
@@ -31,17 +46,63 @@ class RLlibPPO(AgentSessionABC):
lay_down_config_path
):
super().__init__(training_config_path, lay_down_config_path)
self._ppo_config: PPOConfig
if not self._training_config.agent_framework == AgentFramework.RLLIB:
msg = (f"Expected RLLIB agent_framework, "
f"got {self._training_config.agent_framework}")
_LOGGER.error(msg)
raise ValueError(msg)
if self._training_config.red_agent_identifier == RedAgentIdentifier.PPO:
self._agent_config_class = PPOConfig
elif self._training_config.red_agent_identifier == RedAgentIdentifier.A2C:
self._agent_config_class = A2CConfig
else:
msg = ("Expected PPO or A2C red_agent_identifier, "
f"got {self._training_config.red_agent_identifier.value}")
_LOGGER.error(msg)
raise ValueError(msg)
self._agent_config: PPOConfig
self._current_result: dict
self._setup()
self._agent.save()
_LOGGER.debug(
f"Created {self.__class__.__name__} using: "
f"agent_framework={self._training_config.agent_framework}, "
f"red_agent_identifier="
f"{self._training_config.red_agent_identifier}, "
f"deep_learning_framework="
f"{self._training_config.deep_learning_framework}"
)
def _update_session_metadata_file(self):
"""
Update the ``session_metadata.json`` file.
Updates the `session_metadata.json`` in the ``session_path`` directory
with the following key/value pairs:
- end_datetime: The date & time the session ended in iso format.
- total_episodes: The total number of training episodes completed.
- total_time_steps: The total number of training time steps completed.
"""
with open(self.session_path / "session_metadata.json", "r") as file:
metadata_dict = json.load(file)
metadata_dict["end_datetime"] = datetime.now().isoformat()
metadata_dict["total_episodes"] = self._current_result["episodes_total"]
metadata_dict["total_time_steps"] = self._current_result["timesteps_total"]
filepath = self.session_path / "session_metadata.json"
_LOGGER.debug(f"Updating Session Metadata file: {filepath}")
with open(filepath, "w") as file:
json.dump(metadata_dict, file)
_LOGGER.debug("Finished updating session metadata file")
def _setup(self):
super()._setup()
register_env("primaite", _env_creator)
self._ppo_config = PPOConfig()
self._agent_config = self._agent_config_class()
self._ppo_config.environment(
self._agent_config.environment(
env="primaite",
env_config=dict(
training_config_path=self._training_config_path,
@@ -52,19 +113,21 @@ class RLlibPPO(AgentSessionABC):
)
)
self._ppo_config.training(
self._agent_config.training(
train_batch_size=self._training_config.num_steps
)
self._ppo_config.framework(
framework=self._training_config.deep_learning_framework.value
self._agent_config.framework(
framework=self._training_config.deep_learning_framework
)
self._ppo_config.rollouts(
self._agent_config.rollouts(
num_rollout_workers=1,
num_envs_per_worker=1,
horizon=self._training_config.num_steps
)
self._agent: Algorithm = self._ppo_config.build()
self._agent: Algorithm = self._agent_config.build(
logger_creator=_custom_log_creator(self.session_path)
)
def _save_checkpoint(self):
checkpoint_n = self._training_config.checkpoint_every_n_episodes
@@ -84,8 +147,8 @@ class RLlibPPO(AgentSessionABC):
):
# Temporarily override train_batch_size and horizon
if time_steps:
self._ppo_config.train_batch_size = time_steps
self._ppo_config.horizon = time_steps
self._agent_config.train_batch_size = time_steps
self._agent_config.horizon = time_steps
if not episodes:
episodes = self._training_config.num_episodes

34
src/primaite/agents/sb3.py
View File

@@ -1,23 +1,48 @@
from typing import Optional
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3 import PPO, A2C
from primaite import getLogger
from primaite.agents.agent import AgentSessionABC
from primaite.common.enums import RedAgentIdentifier, AgentFramework
from primaite.environment.primaite_env import Primaite
from stable_baselines3.ppo import MlpPolicy as PPOMlp
_LOGGER = getLogger(__name__)
class SB3PPO(AgentSessionABC):
class SB3Agent(AgentSessionABC):
def __init__(
self,
training_config_path,
lay_down_config_path
):
super().__init__(training_config_path, lay_down_config_path)
if not self._training_config.agent_framework == AgentFramework.SB3:
msg = (f"Expected SB3 agent_framework, "
f"got {self._training_config.agent_framework}")
_LOGGER.error(msg)
raise ValueError(msg)
if self._training_config.red_agent_identifier == RedAgentIdentifier.PPO:
self._agent_class = PPO
elif self._training_config.red_agent_identifier == RedAgentIdentifier.A2C:
self._agent_class = A2C
else:
msg = ("Expected PPO or A2C red_agent_identifier, "
f"got {self._training_config.red_agent_identifier.value}")
_LOGGER.error(msg)
raise ValueError(msg)
self._tensorboard_log_path = self.session_path / "tensorboard_logs"
self._tensorboard_log_path.mkdir(parents=True, exist_ok=True)
self._setup()
_LOGGER.debug(
f"Created {self.__class__.__name__} using: "
f"agent_framework={self._training_config.agent_framework}, "
f"red_agent_identifier="
f"{self._training_config.red_agent_identifier}"
)
def _setup(self):
super()._setup()
@@ -28,10 +53,10 @@ class SB3PPO(AgentSessionABC):
session_path=self.session_path,
timestamp_str=self.timestamp_str
)
self._agent = PPO(
self._agent = self._agent_class(
PPOMlp,
self._env,
verbose=1,
verbose=self._training_config.output_verbose_level,
n_steps=self._training_config.num_steps,
tensorboard_log=self._tensorboard_log_path
)
@@ -65,6 +90,7 @@ class SB3PPO(AgentSessionABC):
for i in range(episodes):
self._agent.learn(total_timesteps=time_steps)
self._save_checkpoint()
self._env.close()
super().learn()
def evaluate(