-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_truly_hierarchical.py
99 lines (92 loc) · 2.79 KB
/
train_truly_hierarchical.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
import os
import ray
from ray import air, tune
from ray.rllib.algorithms.ppo import PPO, PPOConfig
from gymnasium.spaces import Discrete, Box
from ray.rllib.algorithms.ppo import PPOConfig
from envs.truly_heirarchical_env import TrulyHeirarchicalDCRL
from envs.heirarchical_env import HeirarchicalDCRL, DEFAULT_CONFIG
from utils.create_trainable import create_wrapped_trainable
from utils.rllib_callbacks import CustomMetricsCallback
NUM_WORKERS = 1
NAME = "test"
RESULTS_DIR = './results/'
# Dummy env to get obs and action space
hdcrl_env = HeirarchicalDCRL()
CONFIG = (
PPOConfig()
.environment(
env=TrulyHeirarchicalDCRL,
env_config=DEFAULT_CONFIG
)
.framework("torch")
.rollouts(
num_rollout_workers=NUM_WORKERS,
rollout_fragment_length=2,
)
.training(
gamma=0.99,
lr=1e-5,
kl_coeff=0.2,
clip_param=0.1,
entropy_coeff=0.0,
use_gae=True,
train_batch_size=4096,
num_sgd_iter=10,
model={'fcnet_hiddens': [64, 64]},
shuffle_sequences=True
)
.multi_agent(
policies={
"high_level_policy": (
None,
hdcrl_env.observation_space,
hdcrl_env.action_space,
PPOConfig()
),
"DC1_ls_policy": (
None,
Box(-1.0, 1.0, (14,)),
Discrete(3),
PPOConfig()
),
"DC2_ls_policy": (
None,
Box(-1.0, 1.0, (14,)),
Discrete(3),
PPOConfig()
),
"DC3_ls_policy": (
None,
Box(-1.0, 1.0, (14,)),
Discrete(3),
PPOConfig()
),
},
policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: agent_id,
)
.callbacks(CustomMetricsCallback)
.resources(num_gpus=0)
.debugging(seed=0)
)
if __name__ == "__main__":
os.environ["RAY_DEDUP_LOGS"] = "0"
# ray.init(local_mode=True, ignore_reinit_error=True)
ray.init(ignore_reinit_error=True)
tune.Tuner(
create_wrapped_trainable(PPO),
param_space=CONFIG.to_dict(),
run_config=air.RunConfig(
stop={"timesteps_total": 100_000_000},
verbose=0,
local_dir=RESULTS_DIR,
# storage_path=RESULTS_DIR,
name=NAME,
checkpoint_config=ray.air.CheckpointConfig(
checkpoint_frequency=5,
num_to_keep=5,
checkpoint_score_attribute="episode_reward_mean",
checkpoint_score_order="max"
),
)
).fit()