def __init__(self, dim, name='wjs_policy'):
super().__init__(dim, name)
self.state_dim = dim * dim * 3
self.gamma = 0.99
self.lr = 0.0005
self.data = ExperienceReplay(batch_size = 100, capacity=10000)
with self.graph.as_default():
self.pl_s = tf.placeholder(tf.float32, [None, self.state_dim], 'state')
self.pl_r = tf.placeholder(tf.float32, [None, 1], 'reward')
self.pl_s_ = tf.placeholder(tf.float32, [None, self.state_dim], 'next_state')
self.pl_done = tf.placeholder(tf.float32, [None, 1], 'done')
self.v = self.v_net('v', self.pl_s)
self.action = tf.argmax(self.v)
self.v_ = self.v_net('v', self.pl_s_)
self.predict = tf.stop_gradient(self.pl_r + self.gamma * self.v_ * (1 - self.pl_done))
self.v_loss = tf.reduce_mean(tf.squared_difference(self.v, self.predict))
self.v_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='v')
optimizer = tf.train.AdamOptimizer(self.lr)
self.train_v = optimizer.minimize(self.v_loss, var_list=self.v_vars, global_step=self.global_step)
tf.summary.scalar('LOSS/v_loss', self.v_loss)
self.summaries = tf.summary.merge_all()
self.sess.run(tf.global_variables_initializer())
def __init__(self,
s_dim,
visual_sources,
visual_resolution,
a_dim_or_list,
action_type,
gamma,
max_episode,
base_dir,
policy_mode=None,
batch_size=1,
buffer_size=1,
use_priority=False,
n_step=False):
super().__init__(a_dim_or_list, action_type, base_dir)
self.s_dim = s_dim
self.visual_sources = visual_sources
self.visual_dim = [visual_sources, *visual_resolution] if visual_sources else [0]
self.a_dim_or_list = a_dim_or_list
self.gamma = gamma
self.max_episode = max_episode
self.policy_mode = policy_mode
self.batch_size = batch_size
self.buffer_size = buffer_size
'''
the biggest diffenernce between policy_modes(ON and OFF) is 'OFF' mode need raise the dimension
of 'r' and 'done'.
'ON' mode means program will call on_store function and use pandas dataframe to store data.
'OFF' mode will call off_store function and use replay buffer to store data.
'''
if self.policy_mode == 'ON':
self.data = pd.DataFrame(columns=['s', 'a', 'r', 's_', 'done'])
elif self.policy_mode == 'OFF':
if use_priority:
if n_step:
print('N-Step PER')
self.data = NStepPrioritizedExperienceReplay(self.batch_size, self.buffer_size, max_episode=self.max_episode,
gamma=self.gamma, alpha=0.6, beta=0.2, epsilon=0.01, agents_num=20, n=4)
else:
print('PER')
self.data = PrioritizedExperienceReplay(self.batch_size, self.buffer_size, max_episode=self.max_episode, alpha=0.6, beta=0.2, epsilon=0.01)
else:
if n_step:
print('N-Step ER')
self.data = NStepExperienceReplay(self.batch_size, self.buffer_size, gamma=self.gamma, agents_num=20, n=4)
else:
print('ER')
self.data = ExperienceReplay(self.batch_size, self.buffer_size)
else:
raise Exception('Please specific a mode of policy!')
with self.graph.as_default():
self.pl_s = tf.placeholder(tf.float32, [None, self.s_dim], 'vector_observation')
self.pl_a = tf.placeholder(tf.float32, [None, self.a_counts], 'pl_action')
self.pl_r = tf.placeholder(tf.float32, [None, 1], 'reward')
self.pl_s_ = tf.placeholder(tf.float32, [None, self.s_dim], 'next_state')
self.pl_done = tf.placeholder(tf.float32, [None, 1], 'done')
self.pl_visual_s = tf.placeholder(tf.float32, [None] + self.visual_dim, 'visual_observation_')
self.pl_visual_s_ = tf.placeholder(tf.float32, [None] + self.visual_dim, 'next_visual_observation')
def __init__(self, dim, name='wjs_policy'):
super().__init__(dim, name)
self.state_dim = dim * dim * 3
self.gamma = 0.99
self.lr = 0.0005
self.data = ExperienceReplay(batch_size=100, capacity=10000)
self.v_net = V(vector_dim=self.state_dim,
name='v_net',
hidden_units=[128, 64, 32])
self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr)
def init_data_memory(self):
'''
the biggest diffenernce between policy_modes(ON and OFF) is 'OFF' mode need raise the dimension
of 'r' and 'done'.
'ON' mode means program will call on_store function and use pandas dataframe to store data.
'OFF' mode will call off_store function and use replay buffer to store data.
'''
if self.policy_mode == 'ON':
self.data = pd.DataFrame(columns=['s', 'a', 'r', 'done'])
elif self.policy_mode == 'OFF':
if self.use_priority:
if self.n_step:
print('N-Step PER')
self.data = NStepPrioritizedExperienceReplay(
self.batch_size,
self.buffer_size,
max_episode=self.max_episode,
gamma=self.gamma,
alpha=er_config['nper_config']['alpha'],
beta=er_config['nper_config']['beta'],
epsilon=er_config['nper_config']['epsilon'],
agents_num=er_config['nper_config']['max_agents'],
n=er_config['nper_config']['n'],
global_v=er_config['nper_config']['global_v'])
else:
print('PER')
self.data = PrioritizedExperienceReplay(
self.batch_size,
self.buffer_size,
max_episode=self.max_episode,
alpha=er_config['per_config']['alpha'],
beta=er_config['per_config']['beta'],
epsilon=er_config['per_config']['epsilon'],
global_v=er_config['nper_config']['global_v'])
else:
if self.n_step:
print('N-Step ER')
self.data = NStepExperienceReplay(
self.batch_size,
self.buffer_size,
gamma=self.gamma,
agents_num=er_config['ner_config']['max_agents'],
n=er_config['ner_config']['n'])
else:
print('ER')
self.data = ExperienceReplay(self.batch_size,
self.buffer_size)
else:
raise Exception('Please specific a mode of policy!')
def __init__(self,
state_dim,
learning_rate=5.0e-4,
buffer_size=10000,
batch_size=128,
epochs=2,
name='wjs_policy',
cp_dir='./models'):
super().__init__(cp_dir=cp_dir)
self.lr = learning_rate
self.epochs = epochs
self.data = ExperienceReplay(batch_size=batch_size,
capacity=buffer_size)
self.net = PV(state_dim=state_dim, name='pv_net')
self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr)
def init_data_memory(self):
if self.use_priority:
if self.n_step:
print('N-Step PER')
self.data = NStepPrioritizedExperienceReplay(
self.batch_size,
self.buffer_size,
max_episode=self.max_episode,
gamma=self.gamma,
alpha=er_config['nper_config']['alpha'],
beta=er_config['nper_config']['beta'],
epsilon=er_config['nper_config']['epsilon'],
agents_num=er_config['nper_config']['max_agents'],
n=er_config['nper_config']['n'],
global_v=er_config['nper_config']['global_v'])
else:
print('PER')
self.data = PrioritizedExperienceReplay(
self.batch_size,
self.buffer_size,
max_episode=self.max_episode,
alpha=er_config['per_config']['alpha'],
beta=er_config['per_config']['beta'],
epsilon=er_config['per_config']['epsilon'],
global_v=er_config['nper_config']['global_v'])
else:
if self.n_step:
print('N-Step ER')
self.data = NStepExperienceReplay(
self.batch_size,
self.buffer_size,
gamma=self.gamma,
agents_num=er_config['ner_config']['max_agents'],
n=er_config['ner_config']['n'])
else:
print('ER')
self.data = ExperienceReplay(self.batch_size, self.buffer_size)
def __init__(self, env_args: Config, model_args: Config,
buffer_args: Config, train_args: Config):
# print("89898989")
self.env_args = env_args
self.model_args = model_args
self.buffer_args = buffer_args
self.train_args = train_args
self.use_GCN = False
self.model_index = str(self.train_args.get('index'))
self.all_learner_print = bool(
self.train_args.get('all_learner_print', False))
if '-' not in self.train_args['name']:
self.train_args['name'] += f'-{self.model_index}'
if self.model_args['load'] is None:
self.train_args['load_model_path'] = os.path.join(
self.train_args['base_dir'], self.train_args['name'])
else:
if '/' in self.model_args['load'] or '\\' in self.model_args[
'load']: # 所有训练进程都以该模型路径初始化，绝对路径
self.train_args['load_model_path'] = self.model_args['load']
elif '-' in self.model_args['load']:
self.train_args['load_model_path'] = os.path.join(
self.train_args['base_dir'],
self.model_args['load']) # 指定了名称和序号，所有训练进程都以该模型路径初始化，相对路径
else: # 只写load的训练名称，不用带进程序号，会自动补
self.train_args['load_model_path'] = os.path.join(
self.train_args['base_dir'],
self.model_args['load'] + f'-{self.model_index}')
# ENV
self.env = make_env(self.env_args.to_dict, self.use_GCN)
# ALGORITHM CONFIG
Model, algorithm_config, _policy_mode = get_model_info(
self.model_args['algo'])
self.model_args['policy_mode'] = _policy_mode
if self.model_args['algo_config'] is not None:
algorithm_config = UpdateConfig(algorithm_config,
self.model_args['algo_config'],
'algo')
ShowConfig(algorithm_config)
# BUFFER
if _policy_mode == 'off-policy':
self.buffer_args['batch_size'] = algorithm_config['batch_size']
self.buffer_args['buffer_size'] = algorithm_config['buffer_size']
if self.model_args['algo'] in ['drqn', 'drdqn']:
self.buffer_args['type'] = 'EpisodeER'
else:
_use_priority = algorithm_config.get('use_priority', False)
_n_step = algorithm_config.get('n_step', False)
if _use_priority and _n_step:
self.buffer_args['type'] = 'NstepPER'
self.buffer_args['NstepPER'][
'max_episode'] = self.train_args['max_episode']
self.buffer_args['NstepPER']['gamma'] = algorithm_config[
'gamma']
algorithm_config['gamma'] = pow(
algorithm_config['gamma'], self.buffer_args['NstepPER']
['n']) # update gamma for n-step training.
elif _use_priority:
self.buffer_args['type'] = 'PER'
self.buffer_args['PER']['max_episode'] = self.train_args[
'max_episode']
elif _n_step:
self.buffer_args['type'] = 'NstepER'
self.buffer_args['NstepER']['gamma'] = algorithm_config[
'gamma']
algorithm_config['gamma'] = pow(
algorithm_config['gamma'],
self.buffer_args['NstepER']['n'])
else:
self.buffer_args['type'] = 'ER'
else:
self.buffer_args['type'] = 'Pandas'
# MODEL
base_dir = os.path.join(
self.train_args['base_dir'], self.train_args['name']
) # train_args['base_dir'] DIR/ENV_NAME/ALGORITHM_NAME
if 'batch_size' in algorithm_config.keys() and train_args['fill_in']:
self.train_args['pre_fill_steps'] = algorithm_config['batch_size']
if self.env_args['type'] == 'gym':
self.eval_env_args = deepcopy(self.env_args)
self.eval_env_args.env_num = 1
self.eval_env = make_env(self.eval_env_args.to_dict)
# buffer ------------------------------
if 'Nstep' in self.buffer_args[
'type'] or 'Episode' in self.buffer_args['type']:
self.buffer_args[self.buffer_args['type']][
'agents_num'] = self.env_args['env_num']
self.buffer = get_buffer(self.buffer_args)
# buffer ------------------------------
# model -------------------------------
model_params = {
's_dim': self.env.s_dim,
'visual_sources': self.env.visual_sources,
'visual_resolution': self.env.visual_resolution,
'a_dim_or_list': self.env.a_dim_or_list,
'is_continuous': self.env.is_continuous,
'max_episode': self.train_args.max_episode,
'base_dir': base_dir,
'logger2file': self.model_args.logger2file,
'seed': self.model_args.seed
}
self.model = Model(**model_params, **algorithm_config)
self.model.set_buffer(self.buffer)
self.model.init_or_restore(self.train_args['load_model_path'])
# model -------------------------------
self.train_args['begin_episode'] = self.model.get_init_episode()
if not self.train_args['inference']:
records_dict = {
'env': self.env_args.to_dict,
'model': self.model_args.to_dict,
'buffer': self.buffer_args.to_dict,
'train': self.train_args.to_dict,
'algo': algorithm_config
}
save_config(os.path.join(base_dir, 'config'), records_dict)
else:
# buffer -----------------------------------
self.buffer_args_s = []
for i in range(self.env.brain_num):
_bargs = deepcopy(self.buffer_args)
if 'Nstep' in _bargs['type'] or 'Episode' in _bargs['type']:
_bargs[_bargs['type']][
'agents_num'] = self.env.brain_agents[i]
self.buffer_args_s.append(_bargs)
buffers = [
get_buffer(self.buffer_args_s[i])
for i in range(self.env.brain_num)
]
# buffer -----------------------------------
# model ------------------------------------
self.model_args_s = []
for i in range(self.env.brain_num):
_margs = deepcopy(self.model_args)
_margs['seed'] = self.model_args['seed'] + i * 10
self.model_args_s.append(_margs)
model_params = [
{
's_dim': self.env.s_dim[i],
'a_dim_or_list': self.env.a_dim_or_list[i],
'visual_sources': self.env.visual_sources[i],
'visual_resolution': self.env.visual_resolutions[i],
'is_continuous': self.env.is_continuous[i],
'max_episode': self.train_args.max_episode,
'base_dir': os.path.join(base_dir, b),
'logger2file': self.model_args_s[i].logger2file,
'seed': self.model_args_s[i].
seed, # 0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100
} for i, b in enumerate(self.env.brain_names)
]
# multi agent training------------------------------------
if self.model_args['algo'][:3] == 'ma_':
self.ma = True
assert self.env.brain_num > 1, 'if using ma* algorithms, number of brains must larger than 1'
self.ma_data = ExperienceReplay(batch_size=10, capacity=1000)
[
mp.update({
'n': self.env.brain_num,
'i': i
}) for i, mp in enumerate(model_params)
]
else:
self.ma = False
# multi agent training------------------------------------
self.models = [
Model(**model_params[i], **algorithm_config)
for i in range(self.env.brain_num)
]
[
model.set_buffer(buffer)
for model, buffer in zip(self.models, buffers)
]
[
self.models[i].init_or_restore(
os.path.join(self.train_args['load_model_path'], b))
for i, b in enumerate(self.env.brain_names)
]
# model ------------------------------------
self.train_args['begin_episode'] = self.models[0].get_init_episode(
)
if not self.train_args['inference']:
for i, b in enumerate(self.env.brain_names):
records_dict = {
'env': self.env_args.to_dict,
'model': self.model_args_s[i].to_dict,
'buffer': self.buffer_args_s[i].to_dict,
'train': self.train_args.to_dict,
'algo': algorithm_config
}
save_config(os.path.join(base_dir, b, 'config'),
records_dict)
def __init__(
self,
s_dim,
visual_sources,
visual_resolution,
a_dim,
is_continuous,
ployak=0.995,
high_scale=1.0,
reward_scale=1.0,
sample_g_nums=100,
sub_goal_steps=10,
fn_goal_dim=0,
intrinsic_reward_mode='os',
high_batch_size=256,
high_buffer_size=100000,
low_batch_size=8,
low_buffer_size=10000,
high_actor_lr=1.0e-4,
high_critic_lr=1.0e-3,
low_actor_lr=1.0e-4,
low_critic_lr=1.0e-3,
hidden_units={
'high_actor': [64, 64],
'high_critic': [64, 64],
'low_actor': [64, 64],
'low_critic': [64, 64]
},
**kwargs):
assert visual_sources == 0, 'HIRO doesn\'t support visual inputs.'
super().__init__(s_dim=s_dim,
visual_sources=visual_sources,
visual_resolution=visual_resolution,
a_dim=a_dim,
is_continuous=is_continuous,
**kwargs)
self.data_high = ExperienceReplay(high_batch_size, high_buffer_size)
self.data_low = ExperienceReplay(low_batch_size, low_buffer_size)
self.ployak = ployak
self.high_scale = np.array(
high_scale if isinstance(high_scale, list) else [high_scale] *
self.s_dim,
dtype=np.float32)
self.reward_scale = reward_scale
self.fn_goal_dim = fn_goal_dim
self.sample_g_nums = sample_g_nums
self.sub_goal_steps = sub_goal_steps
self.sub_goal_dim = self.s_dim - self.fn_goal_dim
self.high_noise = rls.ClippedNormalActionNoise(
mu=np.zeros(self.sub_goal_dim),
sigma=self.high_scale * np.ones(self.sub_goal_dim),
bound=self.high_scale / 2)
self.low_noise = rls.ClippedNormalActionNoise(mu=np.zeros(self.a_dim),
sigma=1.0 *
np.ones(self.a_dim),
bound=0.5)
_high_actor_net = lambda: rls.actor_dpg(self.s_dim, self.sub_goal_dim,
hidden_units['high_actor'])
if self.is_continuous:
_low_actor_net = lambda: rls.actor_dpg(
self.s_dim + self.sub_goal_dim, self.a_dim, hidden_units[
'low_actor'])
else:
_low_actor_net = lambda: rls.actor_discrete(
self.s_dim + self.sub_goal_dim, self.a_dim, hidden_units[
'low_actor'])
self.gumbel_dist = tfd.Gumbel(0, 1)
self.high_actor = _high_actor_net()
self.high_actor_target = _high_actor_net()
self.low_actor = _low_actor_net()
self.low_actor_target = _low_actor_net()
_high_critic_net = lambda: rls.critic_q_one(
self.s_dim, self.sub_goal_dim, hidden_units['high_critic'])
_low_critic_net = lambda: rls.critic_q_one(
self.s_dim + self.sub_goal_dim, self.a_dim, hidden_units[
'low_critic'])
self.high_critic = DoubleQ(_high_critic_net)
self.high_critic_target = DoubleQ(_high_critic_net)
self.low_critic = DoubleQ(_low_critic_net)
self.low_critic_target = DoubleQ(_low_critic_net)
self.update_target_net_weights(
self.low_actor_target.weights + self.low_critic_target.weights +
self.high_actor_target.weights + self.high_critic_target.weights,
self.low_actor.weights + self.low_critic.weights +
self.high_actor.weights + self.high_critic.weights)
self.low_actor_lr, self.low_critic_lr = map(
self.init_lr, [low_actor_lr, low_critic_lr])
self.high_actor_lr, self.high_critic_lr = map(
self.init_lr, [high_actor_lr, high_critic_lr])
self.low_actor_optimizer, self.low_critic_optimizer = map(
self.init_optimizer, [self.low_actor_lr, self.low_critic_lr])
self.high_actor_optimizer, self.high_critic_optimizer = map(
self.init_optimizer, [self.high_actor_lr, self.high_critic_lr])
self.model_recorder(
dict(high_actor=self.high_actor,
high_critic=self.high_critic,
low_actor=self.low_actor,
low_critic=self.low_critic,
low_actor_optimizer=self.low_actor_optimizer,
low_critic_optimizer=self.low_critic_optimizer,
high_actor_optimizer=self.high_actor_optimizer,
high_critic_optimizer=self.high_critic_optimizer))
self.counts = 0
self._high_s = [[] for _ in range(self.n_agents)]
self._noop_subgoal = np.random.uniform(-self.high_scale,
self.high_scale,
size=(self.n_agents,
self.sub_goal_dim))
self.get_ir = self.generate_ir_func(mode=intrinsic_reward_mode)
def unity_run(default_args, share_args, options, max_step, max_episode,
save_frequency, name):
from mlagents.envs import UnityEnvironment
from utils.sampler import create_sampler_manager
try:
tf_version, (model, policy_mode,
_) = get_model_info(options['--algorithm'])
algorithm_config = sth.load_config(
f'./Algorithms/{tf_version}/config.yaml')[options['--algorithm']]
ma = options['--algorithm'][:3] == 'ma_'
except KeyError:
raise NotImplementedError
reset_config = default_args['reset_config']
if options['--unity']:
env = UnityEnvironment()
env_name = 'unity'
else:
file_name = default_args['exe_file'] if options[
'--env'] == 'None' else options['--env']
if os.path.exists(file_name):
env = UnityEnvironment(file_name=file_name,
base_port=int(options['--port']),
no_graphics=False if options['--inference']
else not options['--graphic'])
env_dir = os.path.split(file_name)[0]
env_name = os.path.join(*env_dir.replace('\\', '/').replace(
r'//', r'/').split('/')[-2:])
sys.path.append(env_dir)
if os.path.exists(env_dir + '/env_config.py'):
import env_config
reset_config = env_config.reset_config
max_step = env_config.max_step
if os.path.exists(env_dir + '/env_loop.py'):
from env_loop import Loop
else:
raise Exception('can not find this file.')
sampler_manager, resampling_interval = create_sampler_manager(
options['--sampler'], env.reset_parameters)
if 'Loop' not in locals().keys():
if ma:
from ma_loop import Loop
else:
from loop import Loop
if options['--config-file'] != 'None':
algorithm_config = update_config(algorithm_config,
options['--config-file'])
_base_dir = os.path.join(share_args['base_dir'], env_name,
options['--algorithm'])
base_dir = os.path.join(_base_dir, name)
show_config(algorithm_config)
brain_names = env.external_brain_names
brains = env.brains
brain_num = len(brain_names)
visual_resolutions = {}
for i in brain_names:
if brains[i].number_visual_observations:
visual_resolutions[f'{i}'] = [
brains[i].camera_resolutions[0]['height'],
brains[i].camera_resolutions[0]['width'],
1 if brains[i].camera_resolutions[0]['blackAndWhite'] else 3
]
else:
visual_resolutions[f'{i}'] = []
model_params = [{
's_dim':
brains[i].vector_observation_space_size *
brains[i].num_stacked_vector_observations,
'a_dim_or_list':
brains[i].vector_action_space_size,
'action_type':
brains[i].vector_action_space_type,
'max_episode':
max_episode,
'base_dir':
os.path.join(base_dir, i),
'logger2file':
share_args['logger2file'],
'out_graph':
share_args['out_graph'],
} for i in brain_names]
if ma:
assert brain_num > 1, 'if using ma* algorithms, number of brains must larger than 1'
data = ExperienceReplay(share_args['ma']['batch_size'],
share_args['ma']['capacity'])
extra_params = {'data': data}
models = [
model(n=brain_num, i=i, **model_params[i], **algorithm_config)
for i in range(brain_num)
]
else:
extra_params = {}
models = [
model(visual_sources=brains[i].number_visual_observations,
visual_resolution=visual_resolutions[f'{i}'],
**model_params[index],
**algorithm_config) for index, i in enumerate(brain_names)
]
[
models[index].init_or_restore(
os.path.join(
_base_dir,
name if options['--load'] == 'None' else options['--load'], i))
for index, i in enumerate(brain_names)
]
begin_episode = models[0].get_init_episode()
params = {
'env': env,
'brain_names': brain_names,
'models': models,
'begin_episode': begin_episode,
'save_frequency': save_frequency,
'reset_config': reset_config,
'max_step': max_step,
'max_episode': max_episode,
'sampler_manager': sampler_manager,
'resampling_interval': resampling_interval,
'policy_mode': policy_mode
}
if 'batch_size' in algorithm_config.keys() and options['--fill-in']:
steps = algorithm_config['batch_size']
else:
steps = default_args['no_op_steps']
no_op_params = {
'env': env,
'brain_names': brain_names,
'models': models,
'brains': brains,
'steps': steps,
'choose': options['--noop-choose']
}
params.update(extra_params)
no_op_params.update(extra_params)
if options['--inference']:
Loop.inference(env,
brain_names,
models,
reset_config=reset_config,
sampler_manager=sampler_manager,
resampling_interval=resampling_interval)
else:
try:
[
sth.save_config(os.path.join(base_dir, i, 'config'),
algorithm_config) for i in brain_names
]
Loop.no_op(**no_op_params)
Loop.train(**params)
except Exception as e:
print(e)
finally:
try:
[models[i].close() for i in range(len(models))]
except Exception as e:
print(e)
finally:
env.close()
sys.exit()
def __init__(self, env_args: Config, model_args: Config,
buffer_args: Config, train_args: Config):
self.env_args = env_args
self.model_args = model_args
self.buffer_args = buffer_args
self.train_args = train_args
# training control: max_train_step > max_frame_step > max_train_episode
if self.train_args['max_train_step'] > 0:
self.train_args['max_frame_step'] = sys.maxsize
self.train_args['max_train_episode'] = sys.maxsize
elif self.train_args['max_frame_step'] > 0:
self.train_args['max_train_episode'] = sys.maxsize
elif self.train_args['max_train_episode'] <= 0:
raise ValueError(
'max_train_step/max_frame_step/max_train_episode must be specified at least one with value larger than 0.'
)
self.train_args['inference_episode'] = self.train_args[
'inference_episode'] if self.train_args[
'inference_episode'] > 0 else sys.maxsize
self.model_index = str(self.train_args.get('index'))
self.start_time = time.time()
self.all_learner_print = bool(
self.train_args.get('all_learner_print', False))
if '-' not in self.train_args['name']:
self.train_args['name'] += f'-{self.model_index}'
if self.model_args['load'] is None:
self.train_args['load_model_path'] = os.path.join(
self.train_args['base_dir'], self.train_args['name'])
else:
if '/' in self.model_args['load'] or '\\' in self.model_args[
'load']: # 所有训练进程都以该模型路径初始化，绝对路径
self.train_args['load_model_path'] = self.model_args['load']
elif '-' in self.model_args['load']:
self.train_args['load_model_path'] = os.path.join(
self.train_args['base_dir'],
self.model_args['load']) # 指定了名称和序号，所有训练进程都以该模型路径初始化，相对路径
else: # 只写load的训练名称，不用带进程序号，会自动补
self.train_args['load_model_path'] = os.path.join(
self.train_args['base_dir'],
self.model_args['load'] + f'-{self.model_index}')
# ENV
logger.info('Initialize environment begin...')
self.env = make_env(self.env_args.to_dict)
logger.info('Initialize environment successful.')
# ALGORITHM CONFIG
Model, algorithm_config, _policy_mode = get_model_info(
self.model_args['algo'])
self.model_args['policy_mode'] = _policy_mode
if self.model_args['algo_config'] is not None:
algorithm_config = UpdateConfig(algorithm_config,
self.model_args['algo_config'],
'algo')
algorithm_config['use_rnn'] = self.model_args['use_rnn']
ShowConfig(algorithm_config)
# BUFFER
if _policy_mode == 'off-policy':
if algorithm_config['use_rnn'] == True:
self.buffer_args['type'] = 'EpisodeER'
self.buffer_args['batch_size'] = algorithm_config.get(
'episode_batch_size', 0)
self.buffer_args['buffer_size'] = algorithm_config.get(
'episode_buffer_size', 0)
self.buffer_args['EpisodeER'][
'burn_in_time_step'] = algorithm_config.get(
'burn_in_time_step', 0)
self.buffer_args['EpisodeER'][
'train_time_step'] = algorithm_config.get(
'train_time_step', 0)
else:
self.buffer_args['batch_size'] = algorithm_config.get(
'batch_size', 0)
self.buffer_args['buffer_size'] = algorithm_config.get(
'buffer_size', 0)
_use_priority = algorithm_config.get('use_priority', False)
_n_step = algorithm_config.get('n_step', False)
if _use_priority and _n_step:
self.buffer_args['type'] = 'NstepPER'
self.buffer_args['NstepPER'][
'max_train_step'] = self.train_args['max_train_step']
self.buffer_args['NstepPER']['gamma'] = algorithm_config[
'gamma']
algorithm_config['gamma'] = pow(
algorithm_config['gamma'], self.buffer_args['NstepPER']
['n']) # update gamma for n-step training.
elif _use_priority:
self.buffer_args['type'] = 'PER'
self.buffer_args['PER'][
'max_train_step'] = self.train_args['max_train_step']
elif _n_step:
self.buffer_args['type'] = 'NstepER'
self.buffer_args['NstepER']['gamma'] = algorithm_config[
'gamma']
algorithm_config['gamma'] = pow(
algorithm_config['gamma'],
self.buffer_args['NstepER']['n'])
else:
self.buffer_args['type'] = 'ER'
else:
self.buffer_args['type'] = 'None'
self.train_args[
'pre_fill_steps'] = 0 # if on-policy, prefill experience replay is no longer needed.
# MODEL
base_dir = os.path.join(
self.train_args['base_dir'], self.train_args['name']
) # train_args['base_dir'] DIR/ENV_NAME/ALGORITHM_NAME
if self.env_args['type'] == 'gym':
if self.train_args['use_wandb']:
import wandb
wandb_path = os.path.join(base_dir, 'wandb')
if not os.path.exists(wandb_path):
os.makedirs(wandb_path)
wandb.init(sync_tensorboard=True,
name=self.train_args['name'],
dir=base_dir,
project=self.train_args['wandb_project'])
# buffer ------------------------------
if 'Nstep' in self.buffer_args[
'type'] or 'Episode' in self.buffer_args['type']:
self.buffer_args[self.buffer_args['type']][
'agents_num'] = self.env_args['env_num']
self.buffer = get_buffer(self.buffer_args)
# buffer ------------------------------
# model -------------------------------
model_params = {
's_dim': self.env.s_dim,
'visual_sources': self.env.visual_sources,
'visual_resolution': self.env.visual_resolution,
'a_dim': self.env.a_dim,
'is_continuous': self.env.is_continuous,
'max_train_step': self.train_args.max_train_step,
'base_dir': base_dir,
'logger2file': self.model_args.logger2file,
'seed': self.model_args.seed,
'n_agents': self.env.n
}
self.model = Model(**model_params, **algorithm_config)
self.model.set_buffer(self.buffer)
self.model.init_or_restore(self.train_args['load_model_path'])
# model -------------------------------
_train_info = self.model.get_init_training_info()
self.train_args['begin_train_step'] = _train_info['train_step']
self.train_args['begin_frame_step'] = _train_info['frame_step']
self.train_args['begin_episode'] = _train_info['episode']
if not self.train_args['inference']:
records_dict = {
'env': self.env_args.to_dict,
'model': self.model_args.to_dict,
'buffer': self.buffer_args.to_dict,
'train': self.train_args.to_dict,
'algo': algorithm_config
}
save_config(os.path.join(base_dir, 'config'), records_dict)
if self.train_args['use_wandb']:
wandb.config.update(records_dict)
else:
# buffer -----------------------------------
self.buffer_args_s = []
for i in range(self.env.brain_num):
_bargs = deepcopy(self.buffer_args)
if 'Nstep' in _bargs['type'] or 'Episode' in _bargs['type']:
_bargs[_bargs['type']][
'agents_num'] = self.env.brain_agents[i]
self.buffer_args_s.append(_bargs)
buffers = [
get_buffer(self.buffer_args_s[i])
for i in range(self.env.brain_num)
]
# buffer -----------------------------------
# model ------------------------------------
self.model_args_s = []
for i in range(self.env.brain_num):
_margs = deepcopy(self.model_args)
_margs['seed'] = self.model_args['seed'] + i * 10
self.model_args_s.append(_margs)
model_params = [
{
's_dim': self.env.s_dim[i],
'a_dim': self.env.a_dim[i],
'visual_sources': self.env.visual_sources[i],
'visual_resolution': self.env.visual_resolutions[i],
'is_continuous': self.env.is_continuous[i],
'max_train_step': self.train_args.max_train_step,
'base_dir': os.path.join(base_dir, b),
'logger2file': self.model_args_s[i].logger2file,
'seed': self.model_args_s[i].
seed, # 0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100
'n_agents': self.env.brain_agents[i],
} for i, b in enumerate(self.env.fixed_brain_names)
]
# multi agent training------------------------------------
if self.model_args['algo'][:3] == 'ma_':
self.ma = True
assert self.env.brain_num > 1, 'if using ma* algorithms, number of brains must larger than 1'
self.ma_data = ExperienceReplay(batch_size=10, capacity=1000)
[
mp.update({
'n': self.env.brain_num,
'i': i
}) for i, mp in enumerate(model_params)
]
else:
self.ma = False
# multi agent training------------------------------------
self.models = [
Model(**model_params[i], **algorithm_config)
for i in range(self.env.brain_num)
]
[
model.set_buffer(buffer)
for model, buffer in zip(self.models, buffers)
]
[
self.models[i].init_or_restore(
os.path.join(self.train_args['load_model_path'], b))
for i, b in enumerate(self.env.fixed_brain_names)
]
# model ------------------------------------
_train_info = self.models[0].get_init_training_info()
self.train_args['begin_train_step'] = _train_info['train_step']
self.train_args['begin_frame_step'] = _train_info['frame_step']
self.train_args['begin_episode'] = _train_info['episode']
if not self.train_args['inference']:
for i, b in enumerate(self.env.fixed_brain_names):
records_dict = {
'env': self.env_args.to_dict,
'model': self.model_args_s[i].to_dict,
'buffer': self.buffer_args_s[i].to_dict,
'train': self.train_args.to_dict,
'algo': algorithm_config
}
save_config(os.path.join(base_dir, b, 'config'),
records_dict)
pass