mirror of
https://github.com/illiumst/marl-factory-grid.git
synced 2025-05-22 14:56:43 +02:00
116 lines
4.8 KiB
Python
116 lines
4.8 KiB
Python
import warnings
|
|
|
|
from pathlib import Path
|
|
import time
|
|
|
|
from stable_baselines3.common.callbacks import CallbackList
|
|
from stable_baselines3.common.vec_env import SubprocVecEnv
|
|
|
|
from environments.factory.factory_dirt_item import DirtItemFactory
|
|
from environments.factory.factory_item import ItemFactory, ItemProperties
|
|
from environments.factory.factory_dirt import DirtProperties, DirtFactory
|
|
from environments.logging.monitor import MonitorCallback
|
|
from environments.logging.recorder import RecorderCallback
|
|
from environments.utility_classes import MovementProperties
|
|
from plotting.compare_runs import compare_seed_runs, compare_model_runs
|
|
|
|
warnings.filterwarnings('ignore', category=FutureWarning)
|
|
warnings.filterwarnings('ignore', category=UserWarning)
|
|
|
|
|
|
def make_env(env_kwargs_dict):
|
|
|
|
def _init():
|
|
with DirtFactory(**env_kwargs_dict) as init_env:
|
|
return init_env
|
|
|
|
return _init
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# combine_runs(Path('debug_out') / 'A2C_1630314192')
|
|
# exit()
|
|
|
|
# compare_runs(Path('debug_out'), 1623052687, ['step_reward'])
|
|
# exit()
|
|
|
|
from stable_baselines3 import PPO, DQN, A2C
|
|
# from algorithms.reg_dqn import RegDQN
|
|
# from sb3_contrib import QRDQN
|
|
|
|
dirt_props = DirtProperties(clean_amount=2, gain_amount=0.1, max_global_amount=20,
|
|
max_local_amount=1, spawn_frequency=16, max_spawn_ratio=0.05,
|
|
dirt_smear_amount=0.0, agent_can_interact=True)
|
|
item_props = ItemProperties(n_items=10, agent_can_interact=True,
|
|
spawn_frequency=30, n_drop_off_locations=2,
|
|
max_agent_inventory_capacity=15)
|
|
move_props = MovementProperties(allow_diagonal_movement=True,
|
|
allow_square_movement=True,
|
|
allow_no_op=False)
|
|
train_steps = 5e6
|
|
time_stamp = int(time.time())
|
|
|
|
out_path = None
|
|
|
|
for modeL_type in [A2C, PPO, DQN]: # ,RegDQN, QRDQN]:
|
|
for seed in range(3):
|
|
env_kwargs = dict(n_agents=1,
|
|
# item_properties=item_props,
|
|
dirt_properties=dirt_props,
|
|
movement_properties=move_props,
|
|
pomdp_r=2, max_steps=1000, parse_doors=False,
|
|
level_name='rooms', frames_to_stack=4,
|
|
omit_agent_in_obs=True, combin_agent_obs=True, record_episodes=False,
|
|
cast_shadows=True, doors_have_area=False, env_seed=seed, verbose=False,
|
|
)
|
|
|
|
if modeL_type.__name__ in ["PPO", "A2C"]:
|
|
kwargs = dict(ent_coef=0.01)
|
|
env = SubprocVecEnv([make_env(env_kwargs) for _ in range(10)], start_method="spawn")
|
|
elif modeL_type.__name__ in ["RegDQN", "DQN", "QRDQN"]:
|
|
env = make_env(env_kwargs)()
|
|
kwargs = dict(buffer_size=50000,
|
|
learning_starts=64,
|
|
batch_size=64,
|
|
target_update_interval=5000,
|
|
exploration_fraction=0.25,
|
|
exploration_final_eps=0.025
|
|
)
|
|
else:
|
|
raise NameError(f'The model "{modeL_type.__name__}" has the wrong name.')
|
|
|
|
model = modeL_type("MlpPolicy", env, verbose=1, seed=seed, device='cpu', **kwargs)
|
|
|
|
out_path = Path('debug_out') / f'{model.__class__.__name__}_{time_stamp}'
|
|
|
|
# identifier = f'{seed}_{model.__class__.__name__}_{time_stamp}'
|
|
identifier = f'{seed}_{model.__class__.__name__}_{time_stamp}'
|
|
out_path /= identifier
|
|
|
|
callbacks = CallbackList(
|
|
[MonitorCallback(filepath=out_path / f'monitor_{identifier}.pick'),
|
|
RecorderCallback(filepath=out_path / f'recorder_{identifier}.json', occupation_map=False,
|
|
trajectory_map=False
|
|
)]
|
|
)
|
|
|
|
model.learn(total_timesteps=int(train_steps), callback=callbacks)
|
|
|
|
save_path = out_path / f'model_{identifier}.zip'
|
|
save_path.parent.mkdir(parents=True, exist_ok=True)
|
|
model.save(save_path)
|
|
param_path = out_path.parent / f'env_{model.__class__.__name__}_{time_stamp}.json'
|
|
try:
|
|
env.env_method('save_params', param_path)
|
|
except AttributeError:
|
|
env.save_params(param_path)
|
|
print("Model Trained and saved")
|
|
print("Model Group Done.. Plotting...")
|
|
|
|
if out_path:
|
|
compare_seed_runs(out_path.parent)
|
|
print("All Models Done... Evaluating")
|
|
if out_path:
|
|
compare_model_runs(Path('debug_out'), time_stamp, 'step_reward')
|