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

#917 - Fixed the RLlib integration

Browse Source 
- Dropped support for overriding the num_episodes and num_steps at the agent level. It's just not needed and will add complexity when overriding and writing output files.
This commit is contained in:
Chris McCarthy
2023-06-30 16:52:57 +01:00
parent 00185d3dad
commit e11fd2ced4
43 changed files with 284 additions and 896 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,9 +3,8 @@ from __future__ import annotations
import json
from datetime import datetime
from pathlib import Path
from typing import Optional, Union
from typing import Union
import tensorflow as tf
from ray.rllib.algorithms import Algorithm
from ray.rllib.algorithms.a2c import A2CConfig
from ray.rllib.algorithms.ppo import PPOConfig
@@ -14,11 +13,7 @@ from ray.tune.registry import register_env
from primaite import getLogger
from primaite.agents.agent import AgentSessionABC
from primaite.common.enums import (
AgentFramework,
AgentIdentifier,
DeepLearningFramework,
)
from primaite.common.enums import AgentFramework, AgentIdentifier
from primaite.environment.primaite_env import Primaite
_LOGGER = getLogger(__name__)
@@ -49,10 +44,7 @@ class RLlibAgent(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.RLLIB:
msg = (
f"Expected RLLIB agent_framework, "
f"got {self._training_config.agent_framework}"
)
msg = f"Expected RLLIB agent_framework, " f"got {self._training_config.agent_framework}"
_LOGGER.error(msg)
raise ValueError(msg)
if self._training_config.agent_identifier == AgentIdentifier.PPO:
@@ -60,10 +52,7 @@ class RLlibAgent(AgentSessionABC):
elif self._training_config.agent_identifier == AgentIdentifier.A2C:
self._agent_config_class = A2CConfig
else:
msg = (
"Expected PPO or A2C agent_identifier, "
f"got {self._training_config.agent_identifier.value}"
)
msg = "Expected PPO or A2C agent_identifier, " f"got {self._training_config.agent_identifier.value}"
_LOGGER.error(msg)
raise ValueError(msg)
self._agent_config: PPOConfig
@@ -94,12 +83,8 @@ class RLlibAgent(AgentSessionABC):
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"
]
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}")
@@ -122,9 +107,7 @@ class RLlibAgent(AgentSessionABC):
),
)
self._agent_config.training(
train_batch_size=self._training_config.num_steps
)
self._agent_config.training(train_batch_size=self._training_config.num_steps)
self._agent_config.framework(framework="tf")
self._agent_config.rollouts(
@@ -132,72 +115,41 @@ class RLlibAgent(AgentSessionABC):
num_envs_per_worker=1,
horizon=self._training_config.num_steps,
)
self._agent: Algorithm = self._agent_config.build(
logger_creator=_custom_log_creator(self.session_path)
)
self._agent: Algorithm = self._agent_config.build(logger_creator=_custom_log_creator(self.learning_path))
def _save_checkpoint(self):
checkpoint_n = self._training_config.checkpoint_every_n_episodes
episode_count = self._current_result["episodes_total"]
if checkpoint_n > 0 and episode_count > 0:
if (episode_count % checkpoint_n == 0) or (
episode_count == self._training_config.num_episodes
):
if (episode_count % checkpoint_n == 0) or (episode_count == self._training_config.num_episodes):
self._agent.save(str(self.checkpoints_path))
def learn(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
# Temporarily override train_batch_size and horizon
if time_steps:
self._agent_config.train_batch_size = time_steps
self._agent_config.horizon = time_steps
time_steps = self._training_config.num_steps
episodes = self._training_config.num_episodes
if not episodes:
episodes = self._training_config.num_episodes
_LOGGER.info(
f"Beginning learning for {episodes} episodes @"
f" {time_steps} time steps..."
)
_LOGGER.info(f"Beginning learning for {episodes} episodes @" f" {time_steps} time steps...")
for i in range(episodes):
self._current_result = self._agent.train()
self._save_checkpoint()
if (
self._training_config.deep_learning_framework
!= DeepLearningFramework.TORCH
):
policy = self._agent.get_policy()
tf.compat.v1.summary.FileWriter(
self.session_path / "ray_results", policy.get_session().graph
)
super().learn()
self._agent.stop()
super().learn()
def evaluate(
self,
time_steps: Optional[int] = None,
episodes: Optional[int] = None,
**kwargs,
):
"""
Evaluate the agent.
:param time_steps: The number of steps per episode. Optional. If not
passed, the value from the training config will be used.
:param episodes: The number of episodes. Optional. If not
passed, the value from the training config will be used.
:param kwargs: Any agent-specific key-word args to be passed.
"""
raise NotImplementedError