Skip to content

how to use the trainded model? #33

Description

@suoer95

In the following code (CybORG/Agents/Wrappers)/RLLibWrapper.py), we can get a trainded model, but I don't know how to use the model in a new env. Could you give me some suggestion? Thanks for the info!
`import inspect
import numpy as np
from ray.rllib.agents import ppo
from ray.rllib.env import ParallelPettingZooEnv
from ray.tune import register_env
from CybORG import CybORG
from CybORG.Agents import B_lineAgent, GreenAgent
from CybORG.Agents.Wrappers import ChallengeWrapper

from CybORG.Agents.Wrappers.PettingZooParallelWrapper import PettingZooParallelWrapper
from CybORG.Simulator.Scenarios import FileReaderScenarioGenerator, DroneSwarmScenarioGenerator

class RLLibWrapper(ChallengeWrapper):
def init(self, agent_name, env, reward_threshold=None, max_steps=None):
super().init(agent_name, env, reward_threshold, max_steps)

def step(self, action=None):
    obs, reward, done, info = self.env.step(action=action)
    self.step_counter += 1
    if self.max_steps is not None and self.step_counter >= self.max_steps:
        done = True
    return np.float32(obs), reward, done, info

def reset(self):
    self.step_counter = 0
    obs = self.env.reset()
    return np.float32(obs)

def env_creator_CC1(env_config: dict):
path = str(inspect.getfile(CybORG))
path = path[:-7] + f'/Simulator/Scenarios/scenario_files/Scenario1b.yaml'
sg = FileReaderScenarioGenerator(path)
agents = {"Red": B_lineAgent(), "Green": GreenAgent()}
cyborg = CybORG(scenario_generator=sg, environment='sim', agents=agents)
env = RLLibWrapper(env=cyborg, agent_name="Blue", max_steps=100)
return env

def env_creator_CC2(env_config: dict):
path = str(inspect.getfile(CybORG))
path = path[:-7] + f'/Simulator/Scenarios/scenario_files/Scenario2.yaml'
sg = FileReaderScenarioGenerator(path)
agents = {"Red": B_lineAgent(), "Green": GreenAgent()}
cyborg = CybORG(scenario_generator=sg, environment='sim', agents=agents)
env = RLLibWrapper(env=cyborg, agent_name="Blue", max_steps=100)
return env

def env_creator_CC3(env_config: dict):
sg = DroneSwarmScenarioGenerator()
cyborg = CybORG(scenario_generator=sg, environment='sim')
env = ParallelPettingZooEnv(PettingZooParallelWrapper(env=cyborg))
return env

def print_results(results_dict):
train_iter = results_dict["training_iteration"]
r_mean = results_dict["episode_reward_mean"]
r_max = results_dict["episode_reward_max"]
r_min = results_dict["episode_reward_min"]
print(f"{train_iter:4d} \tr_mean: {r_mean:.1f} \tr_max: {r_max:.1f} \tr_min: {r_min: .1f}")

if name == "main":
register_env(name="CC1", env_creator=env_creator_CC1)
register_env(name="CC2", env_creator=env_creator_CC2)
register_env(name="CC3", env_creator=env_creator_CC3)
config = ppo.DEFAULT_CONFIG.copy()
for env in ['CC1', 'CC2', 'CC3']:
agent = ppo.PPOTrainer(config=config, env=env)

    train_steps = 1e2
    total_steps = 0
    while total_steps < train_steps:
        results = agent.train()
        print_results(results)
        total_steps = results["timesteps_total"]`

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions