def make_robotics_env(env_id, seed, rank=0):
"""
Create a wrapped, monitored gym.Env for MuJoCo.
"""
set_global_seeds(seed)
env = gym.make(env_id)
env = FlattenDictWrapper(env, ['observation', 'desired_goal'])
env = Monitor(
env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)),
info_keywords=('is_success',))
env.seed(seed)
return env
def make_mujoco_env(env_id, seed):
"""
Create a wrapped, monitored gym.Env for MuJoCo.
"""
set_global_seeds(seed)
env = gym.make(env_id)
env = Monitor(env, logger.get_dir())
env.seed(seed)
return env
def launch(env_name,
logdir,
n_epochs,
num_cpu,
seed,
replay_strategy,
policy_save_interval,
clip_return,
override_params={},
save_policies=True):
# Fork for multi-CPU MPI implementation.
if num_cpu > 1:
whoami = mpi_fork(num_cpu)
if whoami == 'parent':
sys.exit(0)
import baselines.common.tf_util as U
U.single_threaded_session().__enter__()
rank = MPI.COMM_WORLD.Get_rank()
# Configure logging
if rank == 0:
if logdir or logger.get_dir() is None:
logger.configure(dir=logdir)
else:
logger.configure()
logdir = logger.get_dir()
assert logdir is not None
os.makedirs(logdir, exist_ok=True)
# Seed everything.
rank_seed = seed + 1000000 * rank
set_global_seeds(rank_seed)
# Prepare params.
params = config.DEFAULT_PARAMS
params['env_name'] = env_name
params['replay_strategy'] = replay_strategy
if env_name in config.DEFAULT_ENV_PARAMS:
params.update(config.DEFAULT_ENV_PARAMS[env_name]
) # merge env-specific parameters in
params.update(
**override_params) # makes it possible to override any parameter
with open(os.path.join(logger.get_dir(), 'params.json'), 'w') as f:
json.dump(params, f)
params = config.prepare_params(params)
config.log_params(params, logger=logger)
dims = config.configure_dims(params)
policy = config.configure_ddpg(dims=dims,
params=params,
clip_return=clip_return)
rollout_params = {
'exploit': False,
'use_target_net': False,
'use_demo_states': True,
'compute_Q': False,
'T': params['T'],
}
eval_params = {
'exploit': True,
'use_target_net': params['test_with_polyak'],
'use_demo_states': False,
'compute_Q': True,
'T': params['T'],
}
for name in [
'T', 'rollout_batch_size', 'gamma', 'noise_eps', 'random_eps'
]:
rollout_params[name] = params[name]
eval_params[name] = params[name]
rollout_worker = RolloutWorker(params['make_env'], policy, dims, logger,
**rollout_params)
rollout_worker.seed(rank_seed)
evaluator = RolloutWorker(params['make_env'], policy, dims, logger,
**eval_params)
evaluator.seed(rank_seed)
train(logdir=logdir,
policy=policy,
rollout_worker=rollout_worker,
evaluator=evaluator,
n_epochs=n_epochs,
n_test_rollouts=params['n_test_rollouts'],
n_cycles=params['n_cycles'],
n_batches=params['n_batches'],
policy_save_interval=policy_save_interval,
save_policies=save_policies)
def train(policy, rollout_worker, evaluator, n_epochs, n_test_rollouts,
n_cycles, n_batches, policy_save_interval, save_policies, **kwargs):
rank = MPI.COMM_WORLD.Get_rank()
latest_policy_path = os.path.join(logger.get_dir(), 'policy_latest.pkl')
best_policy_path = os.path.join(logger.get_dir(), 'policy_best.pkl')
periodic_policy_path = os.path.join(logger.get_dir(), 'policy_{}.pkl')
logger.info("Training...")
# logger.info("Epoch -1 | Finish Time :{}".format(datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")))
starttime = datetime.datetime.now()
best_success_rate = -1
for epoch in range(n_epochs):
# train
rollout_worker.clear_history()
for _ in range(n_cycles):
episode = rollout_worker.generate_rollouts()
# print(episode['info_is_success'])
policy.store_episode(episode)
for _ in range(n_batches):
policy.train()
policy.update_target_net()
# test
evaluator.clear_history()
for _ in range(n_test_rollouts):
evaluator.generate_rollouts()
# logger.info("Epoch: {} | Finish Time :{}".format(epoch, datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")))
# record logs
logger.record_tabular('epoch', epoch)
for key, val in evaluator.logs('test'):
logger.record_tabular(key, mpi_average(val))
for key, val in rollout_worker.logs('train'):
logger.record_tabular(key, mpi_average(val))
for key, val in policy.logs():
logger.record_tabular(key, mpi_average(val))
endtime = datetime.datetime.now()
logger.record_tabular('time',
str(endtime - starttime).replace(',', '-'))
if rank == 0:
logger.dump_tabular()
# save the policy if it's better than the previous ones
success_rate = mpi_average(evaluator.current_success_rate())
if rank == 0 and success_rate >= best_success_rate and save_policies:
best_success_rate = success_rate
logger.info(
'New best success rate: {}. Saving policy to {} ...'.format(
best_success_rate, best_policy_path))
evaluator.save_policy(best_policy_path)
evaluator.save_policy(latest_policy_path)
if rank == 0 and policy_save_interval > 0 and epoch % policy_save_interval == 0 and save_policies:
policy_path = periodic_policy_path.format(epoch)
logger.info('Saving periodic policy to {} ...'.format(policy_path))
evaluator.save_policy(policy_path)
# make sure that different threads have different seeds
local_uniform = np.random.uniform(size=(1, ))
root_uniform = local_uniform.copy()
MPI.COMM_WORLD.Bcast(root_uniform, root=0)
if rank != 0:
assert local_uniform[0] != root_uniform[0]
def _thunk():
env = make_atari(env_id)
env.seed(seed + rank)
env = Monitor(env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)))
return wrap_deepmind(env, **wrapper_kwargs)
def train(policy, rollout_worker, evaluator, n_epochs, n_test_rollouts,
n_cycles, n_batches, policy_save_interval, save_policies, **kwargs):
rank = MPI.COMM_WORLD.Get_rank()
latest_policy_path = os.path.join(logger.get_dir(), 'policy_latest.pkl')
best_policy_path = os.path.join(logger.get_dir(), 'policy_best.pkl')
periodic_policy_path = os.path.join(logger.get_dir(), 'policy_{}.pkl')
r_mean_logdir = os.path.join(logger.get_dir(), 'total_rbias_mean.npy')
r_std_logdir = os.path.join(logger.get_dir(), 'total_rbias_std.npy')
logger.info("Training...")
starttime = datetime.datetime.now()
best_success_rate = -1
for epoch in range(n_epochs):
policy.epcoch_num = epoch
# train
rollout_worker.clear_history()
for _ in range(n_cycles):
episode = rollout_worker.generate_rollouts()
policy.store_episode(episode)
for _ in range(n_batches):
policy.train()
policy.update_target_net()
# lky
if rank == 0:
policy.isPlot = False
policy.picdir = os.path.join(logger.get_dir(),
'rew_epoch_' + str(epoch) + '.pdf')
policy.rewdir = os.path.join(logger.get_dir(),
'rew_epoch_' + str(epoch) + '.npy')
# test
evaluator.clear_history()
for _ in range(n_test_rollouts):
evaluator.generate_rollouts()
# record logs
logger.record_tabular('epoch', epoch)
for key, val in evaluator.logs('test'):
logger.record_tabular(key, mpi_average(val))
for key, val in rollout_worker.logs('train'):
logger.record_tabular(key, mpi_average(val))
for key, val in policy.logs():
logger.record_tabular(key, mpi_average(val))
endtime = datetime.datetime.now()
logger.record_tabular('time',
str(endtime - starttime).replace(',', '-'))
if rank == 0:
logger.dump_tabular()
# save reward
if rank == 0:
with open(r_mean_logdir, "wb") as fp:
pickle.dump(policy.total_epoch_r_mean_bias, fp)
with open(r_std_logdir, "wb") as fp:
pickle.dump(policy.total_epoch_r_std_bias, fp)
# save the policy if it's better than the previous ones
success_rate = mpi_average(evaluator.current_success_rate())
if rank == 0 and success_rate >= best_success_rate and save_policies:
best_success_rate = success_rate
logger.info(
'New best success rate: {}. Saving policy to {} ...'.format(
best_success_rate, best_policy_path))
evaluator.save_policy(best_policy_path)
evaluator.save_policy(latest_policy_path)
if rank == 0 and policy_save_interval > 0 and epoch % policy_save_interval == 0 and save_policies:
policy_path = periodic_policy_path.format(epoch)
logger.info('Saving periodic policy to {} ...'.format(policy_path))
evaluator.save_policy(policy_path)
# make sure that different threads have different seeds
local_uniform = np.random.uniform(size=(1, ))
root_uniform = local_uniform.copy()
MPI.COMM_WORLD.Bcast(root_uniform, root=0)
if rank != 0:
assert local_uniform[0] != root_uniform[0]
def train(policy, rollout_worker, evaluator, n_epochs, n_test_rollouts,
n_cycles, n_batches, policy_save_interval, save_policies, num_cpu,
dump_buffer, rank_method, fit_interval, prioritization, **kwargs):
rank = MPI.COMM_WORLD.Get_rank()
latest_policy_path = os.path.join(logger.get_dir(), 'policy_latest.pkl')
best_policy_path = os.path.join(logger.get_dir(), 'policy_best.pkl')
periodic_policy_path = os.path.join(logger.get_dir(), 'policy_{}.pkl')
logger.info("Training...")
best_success_rate = -1
t = 1
starttime = datetime.datetime.now()
for epoch in range(n_epochs):
# train
rollout_worker.clear_history()
for cycle in range(n_cycles):
episode = rollout_worker.generate_rollouts()
if (cycle % fit_interval
== 0) and (not cycle == 0) or (cycle == n_cycles - 1):
if prioritization == 'entropy':
policy.fit_density_model()
policy.store_episode(episode, dump_buffer, rank_method, epoch)
for batch in range(n_batches):
t = ((epoch * n_cycles * n_batches) +
(cycle * n_batches) + batch) * num_cpu
policy.train(t, dump_buffer)
policy.update_target_net()
# test
evaluator.clear_history()
for _ in range(n_test_rollouts):
evaluator.generate_rollouts()
# record logs
logger.record_tabular('epoch', epoch)
for key, val in evaluator.logs('test'):
logger.record_tabular(key, mpi_average(val))
for key, val in rollout_worker.logs('train'):
logger.record_tabular(key, mpi_average(val))
for key, val in policy.logs():
logger.record_tabular(key, mpi_average(val))
endtime = datetime.datetime.now()
logger.record_tabular('time',
str(endtime - starttime).replace(',', '-'))
if rank == 0:
logger.dump_tabular()
if dump_buffer:
policy.dump_buffer(epoch)
# save the policy if it's better than the previous ones
success_rate = mpi_average(evaluator.current_success_rate())
if rank == 0 and success_rate >= best_success_rate and save_policies:
best_success_rate = success_rate
logger.info(
'New best success rate: {}. Saving policy to {} ...'.format(
best_success_rate, best_policy_path))
evaluator.save_policy(best_policy_path)
evaluator.save_policy(latest_policy_path)
if rank == 0 and policy_save_interval > 0 and epoch % policy_save_interval == 0 and save_policies:
policy_path = periodic_policy_path.format(epoch)
logger.info('Saving periodic policy to {} ...'.format(policy_path))
evaluator.save_policy(policy_path)
# make sure that different threads have different seeds
local_uniform = np.random.uniform(size=(1, ))
root_uniform = local_uniform.copy()
MPI.COMM_WORLD.Bcast(root_uniform, root=0)
if rank != 0:
assert local_uniform[0] != root_uniform[0]
def launch(env_name,
logdir,
n_epochs,
num_cpu,
seed,
replay_strategy,
policy_save_interval,
clip_return,
temperature,
prioritization,
binding,
logging,
version,
dump_buffer,
n_cycles,
rank_method,
fit_interval,
override_params={},
save_policies=True):
# Fork for multi-CPU MPI implementation.
if num_cpu > 1:
whoami = mpi_fork(num_cpu, binding)
if whoami == 'parent':
sys.exit(0)
import baselines.common.tf_util as U
U.single_threaded_session().__enter__()
rank = MPI.COMM_WORLD.Get_rank()
# Configure logging
'''
if logging:
logdir = 'logs/'+str(env_name)+'-prioritization'+str(prioritization)+'-replay_strategy'+str(replay_strategy)+\
'-n_epochs'+str(n_epochs)+'-num_cpu'+str(num_cpu)+'-seed'+str(seed)+'-version'+str(version)
else:
logdir = osp.join(tempfile.gettempdir(),
datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"))
'''
# Configure logging
if rank == 0:
if logdir or logger.get_dir() is None:
logger.configure(dir=logdir)
else:
logger.configure()
logdir = logger.get_dir()
assert logdir is not None
os.makedirs(logdir, exist_ok=True)
if rank == 0:
if logdir or logger.get_dir() is None:
logger.configure(dir=logdir)
else:
logger.configure()
logdir = logger.get_dir()
assert logdir is not None
os.makedirs(logdir, exist_ok=True)
# Seed everything.
rank_seed = seed + 1000000 * rank
set_global_seeds(rank_seed)
# Prepare params.
params = config.DEFAULT_PARAMS
params['env_name'] = env_name
params['replay_strategy'] = replay_strategy
params['temperature'] = temperature
params['prioritization'] = prioritization
params['binding'] = binding
params['max_timesteps'] = n_epochs * params['n_cycles'] * params[
'n_batches'] * num_cpu
params['version'] = version
params['dump_buffer'] = dump_buffer
params['n_cycles'] = n_cycles
params['rank_method'] = rank_method
params['fit_interval'] = fit_interval
params['n_epochs'] = n_epochs
params['num_cpu'] = num_cpu
if params['dump_buffer']:
params['alpha'] = 0
if env_name in config.DEFAULT_ENV_PARAMS:
params.update(config.DEFAULT_ENV_PARAMS[env_name]
) # merge env-specific parameters in
params.update(
**override_params) # makes it possible to override any parameter
with open(os.path.join(logger.get_dir(), 'params.json'), 'w') as f:
json.dump(params, f)
params = config.prepare_params(params)
config.log_params(params, logger=logger)
dims = config.configure_dims(params)
policy = config.configure_ddpg(dims=dims,
params=params,
clip_return=clip_return)
rollout_params = {
'exploit': False,
'use_target_net': False,
'use_demo_states': True,
'compute_Q': False,
'T': params['T'],
}
eval_params = {
'exploit': True,
'use_target_net': params['test_with_polyak'],
'use_demo_states': False,
'compute_Q': True,
'T': params['T'],
}
for name in [
'T', 'rollout_batch_size', 'gamma', 'noise_eps', 'random_eps'
]:
rollout_params[name] = params[name]
eval_params[name] = params[name]
rollout_worker = RolloutWorker(params['make_env'], policy, dims, logger,
**rollout_params)
rollout_worker.seed(rank_seed)
evaluator = RolloutWorker(params['make_env'], policy, dims, logger,
**eval_params)
evaluator.seed(rank_seed)
train(logdir=logdir,
policy=policy,
rollout_worker=rollout_worker,
evaluator=evaluator,
n_epochs=n_epochs,
n_test_rollouts=params['n_test_rollouts'],
n_cycles=params['n_cycles'],
n_batches=params['n_batches'],
policy_save_interval=policy_save_interval,
save_policies=save_policies,
num_cpu=num_cpu,
dump_buffer=dump_buffer,
rank_method=rank_method,
fit_interval=fit_interval,
prioritization=prioritization)