def run_sweep_cmds(cmds):
output_dir = Path.cwd().joinpath('sp_outputs')
output_dir.mkdir(parents=True, exist_ok=True)
processes = []
nbsrs = []
for idx, cmd in enumerate(cmds):
logger.info(f'CMD_{idx}:{cmd}')
p = subprocess.Popen(cmd,
shell=True,
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE)
# p = subprocess.Popen(shlex.split(cmd),
# stderr=subprocess.STDOUT,
# stdout=subprocess.PIPE)
processes.append(p)
nbsrs.append(NBSR(p.stdout))
try:
all_done = [False for i in range(len(processes))]
while True:
for idx, p in enumerate(processes):
stime = time.time()
proc_print = False
while True:
lines = nbsrs[idx].readline(0.2)
if lines:
if not proc_print:
logger.info(
f'====================================')
logger.info(f'Process {idx}:')
proc_print = True
logger.info(lines.decode('utf-8'))
if time.time() - stime > 10:
break
else:
break
if p.poll() is not None:
all_done[idx] = True
if all(all_done):
break
time.sleep(2)
logger.info('All processes are completed.')
except KeyboardInterrupt:
logger.warning('Keyboard interruption.')
finally:
print('Exiting...')
for p in processes:
p.terminate()
sys.exit()
def make_vec_env(env_id=None,
num_envs=1,
seed=1,
env_func=None,
no_timeout=False,
env_kwargs=None,
distributed=False,
extra_wrapper=None,
wrapper_kwargs=None):
logger.info(f'Creating {num_envs} environments.')
if env_kwargs is None:
env_kwargs = {}
if wrapper_kwargs is None:
wrapper_kwargs = {}
if distributed:
import horovod.torch as hvd
seed_offset = hvd.rank() * 100000
seed += seed_offset
def make_env(env_id, rank, seed, no_timeout, env_kwargs, extra_wrapper,
wrapper_kwargs):
def _thunk():
if env_func is not None:
env = env_func(**env_kwargs)
else:
env = gym.make(env_id, **env_kwargs)
if no_timeout:
env = NoTimeOutEnv(env)
if extra_wrapper is not None:
env = extra_wrapper(env, **wrapper_kwargs)
env.seed(seed + rank)
return env
return _thunk
envs = [
make_env(env_id, idx, seed, no_timeout, env_kwargs, extra_wrapper,
wrapper_kwargs) for idx in range(num_envs)
]
if num_envs > 1:
envs = ShmemVecEnv(envs, context='spawn')
else:
envs = DummyVecEnv(envs)
return envs
def check_if_run_distributed(cfg):
from easyrl import HOROVOD_AVAILABLE
if HOROVOD_AVAILABLE:
import horovod.torch as hvd
hvd.init()
if hvd.size() > 1:
cfg.distributed = True
if cfg.distributed and not HOROVOD_AVAILABLE:
logger.error('Horovod is not installed! Will not run in distributed training')
distributed = HOROVOD_AVAILABLE and cfg.distributed
cfg.distributed = distributed
if distributed:
gpu_id = hvd.local_rank() + cfg.gpu_shift
if cfg.gpus is not None:
gpu_id = cfg.gpus[gpu_id]
logger.info(f'Rank {hvd.local_rank()} GPU ID: {gpu_id}')
torch.cuda.set_device(gpu_id)
logger.info(f'Using Horovod for distributed training, number of processes:{hvd.size()}')
return distributed
def save_model(self, is_best=False, step=None):
self.save_env(cfg.alg.model_dir)
data_to_save = {
'step': step,
'actor_state_dict': self.actor.state_dict(),
'q1_state_dict': self.q1.state_dict(),
'q1_tgt_state_dict': self.q1_tgt.state_dict(),
'q2_state_dict': self.q2.state_dict(),
'q2_tgt_state_dict': self.q2_tgt.state_dict(),
'pi_optim_state_dict': self.pi_optimizer.state_dict(),
'q_optim_state_dict': self.q_optimizer.state_dict(),
}
if cfg.alg.alpha is None:
data_to_save['log_alpha'] = self.log_alpha
data_to_save[
'alpha_optim_state_dict'] = self.alpha_optimizer.state_dict()
save_model(data_to_save, cfg.alg, is_best=is_best, step=step)
logger.info(f'Saving the replay buffer to: {self.mem_file}.')
save_to_pickle(self.memory, self.mem_file)
logger.info('The replay buffer is saved.')
def save_model(data, cfg, is_best=False, step=None):
if not cfg.save_best_only and step is not None:
ckpt_file = cfg.model_dir \
.joinpath('ckpt_{:012d}.pt'.format(step))
else:
ckpt_file = None
if is_best:
best_model_file = cfg.model_dir.joinpath('model_best.pt')
else:
best_model_file = None
if not cfg.save_best_only:
saved_model_files = sorted(cfg.model_dir.glob('*.pt'))
if len(saved_model_files) > cfg.max_saved_models:
saved_model_files[0].unlink()
logger.info(f'Exploration steps: {step}')
for fl in [ckpt_file, best_model_file]:
if fl is not None:
logger.info(f'Saving checkpoint: {fl}.')
torch.save(data, fl)
def restore_cfg(self, skip_params=None, path=None):
if path is None:
path = self.data_dir
hp_file = path.joinpath('hp.json')
with hp_file.open() as f:
cfg_stored = json.load(f)
if skip_params is None:
skip_params = []
skip_params.extend(['resume',
'resume_step',
'render',
'test',
'test_num',
'eval_num_envs',
'save_test_traj',
'save_dir',
'diff_cfg'])
for key, val in cfg_stored.items():
if hasattr(self, key) and key not in skip_params:
setattr(self, key, val)
logger.info(f'Restoring {key} to {val}.')
def print_param_grad_status(self):
logger.info('Requires Grad?')
logger.info('================== Actor ================== ')
for name, param in self.actor.named_parameters():
print(f'{name}: {param.requires_grad}')
logger.info('================== Critic ================== ')
for name, param in self.critic.named_parameters():
print(f'{name}: {param.requires_grad}')
def make_vec_env(env_id, num_envs, seed=1, no_timeout=False, env_kwargs=None):
logger.info(f'Creating {num_envs} environments.')
if env_kwargs is None:
env_kwargs = {}
def make_env(env_id, rank, seed, no_timeout, env_kwargs):
def _thunk():
from gym import envs
#print(envs.registry.all())
try:
env = gym.make(env_id, **env_kwargs)
except Exception:
if env_id == "CheetahMPCEnv-v0":
env = CheetahMPCEnv(**env_kwargs)
elif env_id == "CheetahRSSPMTGEnv-v0":
env = CheetahRSSPMTGEnv - v0(**env_kwargs)
elif env_id == "CheetahRSSFlatEnv-v0":
env = CheetahRSSFlatEnv - v0(**env_kwargs)
else:
print(f"ENVIRONMENT {env_id} NOT REGISTERED")
raise Exception
if no_timeout:
env = NoTimeOutEnv(env)
env.seed(seed + rank)
return env
return _thunk
envs = [
make_env(env_id, idx, seed, no_timeout, env_kwargs)
for idx in range(num_envs)
]
if num_envs > 1:
envs = ShmemVecEnv(envs, context='spawn')
else:
envs = DummyVecEnv(envs)
return envs