-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patherl_config.py
184 lines (159 loc) · 9.89 KB
/
erl_config.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
175
176
177
178
179
180
181
182
183
184
import os
import torch
import numpy as np
class Config:
def __init__(self, agent_class=None, env_class=None, env_args=None):
self.num_envs = None
self.agent_class = agent_class # agent = agent_class(...)
self.if_off_policy = self.get_if_off_policy() # whether off-policy or on-policy of DRL algorithm
'''Argument of environment'''
self.env_class = env_class # env = env_class(**env_args)
self.env_args = env_args # env = env_class(**env_args)
if env_args is None: # dummy env_args
env_args = {'env_name': None,
'num_envs': 1,
'max_step': 12345,
'state_dim': None,
'action_dim': None,
'if_discrete': None, }
env_args.setdefault('num_envs', 1) # `num_envs=1` in default in single env.
env_args.setdefault('max_step', 12345) # `max_step=12345` in default, which is a large enough value.
self.env_name = env_args['env_name'] # the name of environment. Be used to set 'cwd'.
self.num_envs = env_args['num_envs'] # the number of sub envs in vectorized env. `num_envs=1` in single env.
self.max_step = env_args['max_step'] # the max step number of an episode. 'set as 12345 in default.
self.state_dim = env_args['state_dim'] # vector dimension (feature number) of state
self.action_dim = env_args['action_dim'] # vector dimension (feature number) of action
self.if_discrete = env_args['if_discrete'] # discrete or continuous action space
'''Arguments for reward shaping'''
self.gamma = 0.99 # discount factor of future rewards
self.reward_scale = 2 ** 0 # an approximate target reward usually be closed to 256
'''Arguments for training'''
self.net_dims = (64, 32) # the middle layer dimension of MLP (MultiLayer Perceptron)
self.learning_rate = 6e-5 # the learning rate for network updating
self.clip_grad_norm = 3.0 # 0.1 ~ 4.0, clip the gradient after normalization
self.state_value_tau = 0 # the tau of normalize for value and state `std = (1-std)*std + tau*std`
self.soft_update_tau = 5e-3 # 2 ** -8 ~= 5e-3. the tau of soft target update `net = (1-tau)*net + tau*net1`
if self.if_off_policy: # off-policy
self.batch_size = int(64) # num of transitions sampled from replay buffer.
self.horizon_len = int(512) # collect horizon_len step while exploring, then update networks
self.buffer_size = int(1e6) # ReplayBuffer size. First in first out for off-policy.
self.repeat_times = 1.0 # repeatedly update network using ReplayBuffer to keep critic's loss small
self.if_use_per = False # use PER (Prioritized Experience Replay) for sparse reward
else: # on-policy
self.batch_size = int(128) # num of transitions sampled from replay buffer.
self.horizon_len = int(2048) # collect horizon_len step while exploring, then update network
self.buffer_size = None # ReplayBuffer size. Empty the ReplayBuffer for on-policy.
self.repeat_times = 8.0 # repeatedly update network using ReplayBuffer to keep critic's loss small
self.if_use_vtrace = False # use V-trace + GAE (Generalized Advantage Estimation) for sparse reward
'''Arguments for device'''
self.gpu_id = int(0) # `int` means the ID of single GPU, -1 means CPU
self.num_workers = 2 # rollout workers number pre GPU (adjust it to get high GPU usage)
self.num_threads = 8 # cpu_num for pytorch, `torch.set_num_threads(self.num_threads)`
self.random_seed = 1 # initialize random seed in self.init_before_training()
self.learner_gpus = 0 # `int` means the ID of single GPU, -1 means CPU
'''Arguments for evaluate'''
self.cwd = None # current working directory to save model. None means set automatically
self.if_remove = True # remove the cwd folder? (True, False, None:ask me)
self.break_step = np.inf # break training if 'total_step > break_step'
self.break_score = np.inf # break training if `cumulative_rewards > break_score`
self.if_keep_save = True # keeping save the checkpoint. False means save until stop training.
self.if_over_write = False # overwrite the best policy network. `self.cwd/actor.pth`
self.if_save_buffer = False # if save the replay buffer for continuous training after stop training
self.save_gap = int(8) # save actor f"{cwd}/actor_*.pth" for learning curve.
self.eval_times = int(3) # number of times that get the average episodic cumulative return
self.eval_per_step = int(2e4) # evaluate the agent per training steps
self.eval_env_class = None # eval_env = eval_env_class(*eval_env_args)
self.eval_env_args = None # eval_env = eval_env_class(*eval_env_args)
def init_before_training(self):
np.random.seed(self.random_seed % (2 ** 32))
torch.manual_seed(self.random_seed % (2 ** 32))
torch.set_num_threads(self.num_threads)
torch.set_default_dtype(torch.float32)
'''set cwd (current working directory) for saving model'''
if self.cwd is None: # set cwd (current working directory) for saving model
self.cwd = f'./{self.env_name}_{self.agent_class.__name__[5:]}_{self.random_seed}'
'''remove history'''
if self.if_remove is None:
self.if_remove = bool(input(f"| Arguments PRESS 'y' to REMOVE: {self.cwd}? ") == 'y')
if self.if_remove:
import shutil
shutil.rmtree(self.cwd, ignore_errors=True)
print(f"| Arguments Remove cwd: {self.cwd}")
else:
print(f"| Arguments Keep cwd: {self.cwd}")
os.makedirs(self.cwd, exist_ok=True)
def get_if_off_policy(self) -> bool:
agent_name = self.agent_class.__name__ if self.agent_class else ''
on_policy_names = ('SARSA', 'VPG', 'A2C', 'A3C', 'TRPO', 'PPO', 'MPO')
return all([agent_name.find(s) == -1 for s in on_policy_names])
def print(self):
from pprint import pprint
pprint(vars(self)) # prints out args in a neat, readable format
def build_env(env_class=None, env_args: dict = None, gpu_id: int = -1):
env_args['gpu_id'] = gpu_id # set gpu_id for vectorized env before build it
env = env_class(**kwargs_filter(env_class.__init__, env_args.copy()))
for attr_str in ('env_name', 'num_envs', 'max_step', 'state_dim', 'action_dim', 'if_discrete'):
setattr(env, attr_str, env_args[attr_str])
if env_args.get("eval_sequential", False):
setattr(env, "max_step", (env.full_seq_len - env.seq_len) // env.step_gap)
return env
def kwargs_filter(function, kwargs: dict) -> dict:
import inspect
sign = inspect.signature(function).parameters.values()
sign = {val.name for val in sign}
common_args = sign.intersection(kwargs.keys())
return {key: kwargs[key] for key in common_args} # filtered kwargs
def get_gym_env_args(env, if_print: bool) -> dict:
"""get a dict about a standard OpenAI gym env information.
assert 0.18.0 <= gym.__version__ <= 0.25.3
env: a standard OpenAI gym env
if_print: [bool] print the dict about env information.
return: env_args [dict]
env_args = {
'env_name': env_name, # [str] the environment name, such as XxxXxx-v0
'num_envs': num_envs. # [int] the number of sub envs in vectorized env. `num_envs=1` in single env.
'max_step': max_step, # [int] the max step number of an episode.
'state_dim': state_dim, # [int] the dimension of state
'action_dim': action_dim, # [int] the dimension of action or the number of discrete action
'if_discrete': if_discrete, # [bool] action space is discrete or continuous
}
"""
import gym
if_gym_standard_env = {'unwrapped', 'observation_space', 'action_space', 'spec'}.issubset(dir(env))
if if_gym_standard_env and (not hasattr(env, 'num_envs')): # isinstance(env, gym.Env):
assert '0.18.0' <= gym.__version__ <= '0.25.2' # pip3 install gym==0.24.0
env_name = env.unwrapped.spec.id
num_envs = getattr(env, 'num_envs', 1)
max_step = getattr(env, '_max_episode_steps', 12345)
state_shape = env.observation_space.shape
state_dim = state_shape[0] if len(state_shape) == 1 else state_shape # sometimes state_dim is a list
if_discrete = isinstance(env.action_space, gym.spaces.Discrete)
if if_discrete: # make sure it is discrete action space
action_dim = getattr(env.action_space, 'n')
elif isinstance(env.action_space, gym.spaces.Box): # make sure it is continuous action space
action_dim = env.action_space.shape[0]
if any(env.action_space.high - 1):
print('WARNING: env.action_space.high', env.action_space.high)
if any(env.action_space.low + 1):
print('WARNING: env.action_space.low', env.action_space.low)
else:
raise RuntimeError('\n| Error in get_gym_env_info(). Please set these value manually:'
'\n `state_dim=int; action_dim=int; if_discrete=bool;`'
'\n And keep action_space in range (-1, 1).')
else:
env_name = getattr(env, 'env_name', 'env')
num_envs = getattr(env, 'num_envs', 1)
max_step = getattr(env, 'max_step', 12345)
state_dim = env.state_dim
action_dim = env.action_dim
if_discrete = env.if_discrete
env_args = {'env_name': env_name,
'num_envs': num_envs,
'max_step': max_step,
'state_dim': state_dim,
'action_dim': action_dim,
'if_discrete': if_discrete, }
if if_print:
env_args_str = repr(env_args).replace(',', f",\n{'':11}")
print(f"env_args = {env_args_str}")
return env_args