-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
executable file
·174 lines (137 loc) · 5.42 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
from config_parser import ConfigParser
from stable_baselines3 import DQN, A2C, PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecMonitor
from stable_baselines3.common.utils import set_random_seed
import numpy as np
from bio_world import BioGymWorld
from agents.no_action import NoAction
from agents.random_action import RandomAction
from agents.user_action import UserAction
np.set_printoptions(suppress=True, formatter={'float': "{0:0.3f}".format})
config_parser = ConfigParser("bio_env_configs/2x2_10x.ini")
bio_env, renderer, max_steps, reduced_actions = config_parser.create_bio_gym_world(
)
env = BioGymWorld(bio_env, renderer, max_steps, reduced_actions)
def make_env(rank, seed=0):
"""
Utility function for multiprocessed env.
:param env_id: (str) the environment ID
:param seed: (int) the inital seed for RNG
:param rank: (int) index of the subprocess
"""
def _init():
env_instance = BioGymWorld(bio_env, renderer, max_steps,
reduced_actions)
env_instance.renderer.render_mode = "off"
# use a seed for reproducibility
# Important: use a different seed for each environment
# otherwise they would generate the same experiences
# env_instance.reset(seed=seed + rank)
env_instance.reset()
return env_instance
# set_random_seed(seed)
return _init
def make_train_env():
train_env = SubprocVecEnv(
[make_env(i) for i in range(8)],
start_method="fork",
)
train_env = VecMonitor(train_env)
return train_env
model_type = DQN
def train_model(model_name, timesteps):
train_env = make_train_env()
print("train_env.n_envs = " + str(train_env.num_envs))
# train_env = env
# train_env.renderer.render_mode = "off"
model = model_type("MultiInputPolicy",
train_env,
verbose=1,
tensorboard_log="./logs/")
model.learn(total_timesteps=timesteps,
progress_bar=True,
tb_log_name=model_name)
model.save("trained_models/" + model_name)
def get_agent(name):
if name == "no action":
return NoAction(env)
elif name == "random":
return RandomAction(env)
elif name == "user":
return UserAction(env)
else:
raise NotImplementedError("Agent not implemented.")
def run(episodes,
render_mode,
show_episode_history,
agent_name,
model_name=None):
"""
Run episodes of environment with given RL agent.
"""
score_history = []
species_abundance_history = []
shannon_index_history = []
num_add_pop_actions = 0
num_remove_pop_actions = 0
num_no_actions = 0
env.renderer.render_mode = render_mode
if agent_name == "model":
agent = model_type.load("trained_models/" + model_name)
else:
agent = get_agent(agent_name)
for i in range(1, episodes + 1):
obs, info = env.reset()
done = False
score = 0
while not done:
if agent_name == "model":
action, _state = agent.predict(obs, deterministic=True)
else:
action = agent.predict(obs, env.current_step)
if action < env.bio_env.get_no_action():
num_remove_pop_actions += 1
elif action > env.bio_env.get_no_action():
num_add_pop_actions += 1
else:
num_no_actions += 1
obs, reward, terminated, truncated, info = env.step(action)
score += reward
done = terminated or truncated
score_history.append(score)
species_abundance_history.append(
env.bio_env.get_average_species_abundance())
shannon_index_history.append(env.bio_env.get_average_shannon_index())
if show_episode_history:
env.show_episode_history()
print("Average episode score: " +
str(sum(score_history) / len(score_history)))
print(
"Average episode species abundance (relative to critical threshold): "
+
str(sum(species_abundance_history) / len(species_abundance_history)) +
", std = " + str(np.std(np.array(species_abundance_history))))
print("Average episode Shannon index: " +
str(sum(shannon_index_history) / len(shannon_index_history)) +
", std = " + str(np.std(np.array(shannon_index_history))))
print("----------------------------------")
sum_actions = num_add_pop_actions + num_remove_pop_actions + num_no_actions
print("Num. add actions = " + str(num_add_pop_actions) + " (" +
str(num_add_pop_actions * 100 / sum_actions) + "%)")
print("Num. remove actions = " + str(num_remove_pop_actions) + " (" +
str(num_remove_pop_actions * 100 / sum_actions) + "%)")
print("Num. no actions = " + str(num_no_actions) + " (" +
str(num_no_actions * 100 / sum_actions) + "%)")
env.show_run_history(score_history, species_abundance_history,
shannon_index_history)
env.close()
if __name__ == '__main__':
# for i in range(1, 21):
# model_name = "8env_PPO_2x2_10x_800k_" + str(i)
# train_model(model_name, 800_000)
run(episodes=50,
render_mode="off",
show_episode_history=False,
agent_name="model",
model_name="DQN/2x2_10x/DQN_2x2_10x_200k_8")