def __init__(self, config, buffer_remote, param_server, save_dir): self.device = torch.device(config['device']) self.num_train_steps = config['num_train_steps'] # environment_info network_type = config['network_type'] obs_dim = config['obs_dim'] action_dim = config['action_dim'] # buffer info self.use_per = config['use_per'] self.priority_alpha = config['priority_alpha'] self.priority_beta_start = config['priority_beta_start'] self.priority_beta_end = config['priority_beta_end'] self.priority_beta = self.priority_beta_start self.priority_beta_increment = ( (self.priority_beta_end - self.priority_beta_start) / self.num_train_steps) self.batch_size = config['batch_size'] # hyperparameters self.gamma = config['gamma'] self.unroll_steps = config['unroll_steps'] q_lr = config['q_lr'] policy_lr = config['policy_lr'] self.tau = config['tau'] num_quantiles = config['num_quantiles'] # initialize networks self.q_network = create_critic(network_type, obs_dim, action_dim) self.target_q_network = create_critic(network_type, obs_dim, action_dim) self.policy = create_policy(obs_dim, action_dim) self.target_policy = create_policy(obs_dim, action_dim) self.param_list = list(self.policy.state_dict()) self.param_server = param_server for target_param, param in zip(self.target_q_network.parameters(), self.q_network.parameters()): target_param.data.copy_(param) for target_param, param in zip(self.target_policy.parameters(), self.policy.parameters()): target_param.data.copy_(param) # initialize optimizers self.q_optimizer = optim.Adam(self.q_network.parameters(), lr=q_lr) self.policy_optimizer = optim.Adam(self.policy.parameters(), lr=policy_lr) # parallelization self.buffer_remote = buffer_remote self.update_step = 0 # logging self.save_dir = save_dir self.log = list()
def __init__(self, worker_id, config, save_dir, buffer_remote, param_server): self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.worker_id = worker_id env_id = config['env'] random_seed = config['random_seed'] self.unroll_cnt = config['unroll_steps'] self.num_train_step = config['num_train_steps'] self.max_ep_length = config['max_ep_length'] self.env = gym.make(env_id) random.seed(worker_id + random_seed) self.env.seed(random.randint(1, random_seed)) self.action_range = [ config['action_bound_low'], config['action_bound_high'] ] self.policy = create_policy(config['obs_dim'], config['action_dim']) # parallelization self.buffer_remote = buffer_remote self.param_server = param_server # logging self.save_dir = save_dir self.log = list()
def main(): torch.set_num_threads(1) device = torch.device("cuda:0" if args.cuda else "cpu") if args.vis: from torch.utils.tensorboard import SummaryWriter save_path = os.path.join(args.save_dir, args.algo) try: os.makedirs(save_path) except OSError: pass viz = SummaryWriter( os.path.join( save_path, args.env_name + time.strftime("_%d_%b_%H_%M", time.localtime()))) train_envs = make_vec_envs(args.env_name, args.seed, args.num_processes, args.gamma, args.no_norm, args.num_stack, args.log_dir, args.add_timestep, device, allow_early_resets=False) if args.eval_interval: eval_envs = make_vec_envs(args.env_name, args.seed + args.num_processes, args.num_processes, args.gamma, args.no_norm, args.num_stack, eval_log_dir, args.add_timestep, device, allow_early_resets=True, eval=True) if eval_envs.venv.__class__.__name__ == "VecNormalize": eval_envs.venv.ob_rms = train_envs.venv.ob_rms else: eval_envs = None # FIXME this is very specific to Pommerman env right now actor_critic = create_policy(train_envs.observation_space, train_envs.action_space, name='pomm', nn_kwargs={ 'batch_norm': False if args.algo == 'acktr' else True, 'recurrent': args.recurrent_policy, 'hidden_size': 512, }, train=True) if args.load_path != '': actor_critic.load_state_dict(state_dict) actor_critic.to(device) if args.algo.startswith('a2c'): agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef, args.entropy_coef, lr=args.lr, lr_schedule=lr_update_schedule, eps=args.eps, alpha=args.alpha, max_grad_norm=args.max_grad_norm) elif args.algo.startswith('ppo'): agent = algo.PPO(actor_critic, args.clip_param, args.ppo_epoch, args.num_mini_batch, args.value_loss_coef, args.entropy_coef, lr=args.lr, lr_schedule=lr_update_schedule, eps=args.eps, max_grad_norm=args.max_grad_norm) elif args.algo == 'acktr': agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef, args.entropy_coef, acktr=True) if args.algo.endswith('sil'): agent = algo.SIL(agent, update_ratio=args.sil_update_ratio, epochs=args.sil_epochs, batch_size=args.sil_batch_size, value_loss_coef=args.sil_value_loss_coef or args.value_loss_coef, entropy_coef=args.sil_entropy_coef or args.entropy_coef) replay = ReplayStorage(5e5, args.num_processes, args.gamma, 0.1, train_envs.observation_space.shape, train_envs.action_space, actor_critic.recurrent_hidden_state_size, device=device) else: replay = None rollouts = RolloutStorage(args.num_steps, args.num_processes, train_envs.observation_space.shape, train_envs.action_space, actor_critic.recurrent_hidden_state_size) obs = train_envs.reset() rollouts.obs[0].copy_(obs) rollouts.to(device) episode_rewards = deque(maxlen=10) start = time.time() for j in range(num_updates): for step in range(args.num_steps): # Sample actions with torch.no_grad(): value, action, action_log_prob, recurrent_hidden_states = actor_critic.act( rollouts.obs[step], rollouts.recurrent_hidden_states[step], rollouts.masks[step]) # Obser reward and next obs obs, reward, done, infos = train_envs.step(action) for info in infos: if 'episode' in info.keys(): episode_rewards.append(info['episode']['r']) # If done then clean the history of observations. masks = torch.tensor([[0.0] if done_ else [1.0] for done_ in done], device=device) rollouts.insert(obs, recurrent_hidden_states, action, action_log_prob, value, reward, masks) if replay is not None: replay.insert(rollouts.obs[step], rollouts.recurrent_hidden_states[step], action, reward, done) with torch.no_grad(): next_value = actor_critic.get_value( rollouts.obs[-1], rollouts.recurrent_hidden_states[-1], rollouts.masks[-1]).detach() rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau) value_loss, action_loss, dist_entropy, other_metrics = agent.update( rollouts, j, replay) rollouts.after_update() if j % args.save_interval == 0 and args.save_dir != "": save_path = os.path.join(args.save_dir, args.algo) try: os.makedirs(save_path) except OSError: pass # A really ugly way to save a model to CPU save_model = actor_critic if args.cuda: save_model = copy.deepcopy(actor_critic).cpu() save_model = [ save_model.state_dict(), hasattr(train_envs.venv, 'ob_rms') and train_envs.venv.ob_rms or None ] torch.save(save_model, os.path.join(save_path, args.env_name + ".pt")) total_num_steps = (j + 1) * update_factor if j % args.log_interval == 0 and len(episode_rewards) > 1: end = time.time() print( "Updates {}, num timesteps {}, FPS {}, last {} mean/median reward {:.1f}/{:.1f}, " "min / max reward {:.1f}/{:.1f}, dist_entropy{:.5f} value/action loss {:.5f}/{:.5f}" .format(j, total_num_steps, int(total_num_steps / (end - start)), len(episode_rewards), np.mean(episode_rewards), np.median(episode_rewards), np.min(episode_rewards), np.max(episode_rewards), dist_entropy, value_loss, action_loss), end=', ' if other_metrics else '\n') if 'sil_value_loss' in other_metrics: print("SIL value/action loss {:.1f}/{:.1f}.".format( other_metrics['sil_value_loss'], other_metrics['sil_action_loss'])) if args.vis: viz.add_scalar('episode_rewards/mean', np.mean(episode_rewards), total_num_steps) viz.add_scalar('episode_rewards/median', np.median(episode_rewards), total_num_steps) viz.add_scalar('episode_rewards/min', np.min(episode_rewards), total_num_steps) viz.add_scalar('episode_rewards/max', np.max(episode_rewards), total_num_steps) viz.add_scalar('train/value_loss', value_loss, total_num_steps) viz.add_scalar('train/action_loss', action_loss, total_num_steps) viz.add_scalar('train/dist_entropy', dist_entropy, total_num_steps) if args.eval_interval and len( episode_rewards) > 1 and j > 0 and j % args.eval_interval == 0: eval_episode_rewards = [] obs = eval_envs.reset() eval_recurrent_hidden_states = torch.zeros( args.num_processes, actor_critic.recurrent_hidden_state_size, device=device) eval_masks = torch.zeros(args.num_processes, 1, device=device) while len(eval_episode_rewards) < 50: with torch.no_grad(): _, action, _, eval_recurrent_hidden_states = actor_critic.act( obs, eval_recurrent_hidden_states, eval_masks, deterministic=True) # Obser reward and next obs obs, reward, done, infos = eval_envs.step(action) eval_masks = torch.tensor([[0.0] if done_ else [1.0] for done_ in done], device=device) for info in infos: if 'episode' in info.keys(): eval_episode_rewards.append(info['episode']['r']) print(" Evaluation using {} episodes: mean reward {:.5f}\n".format( len(eval_episode_rewards), np.mean(eval_episode_rewards)))
def main(): torch.set_num_threads(1) device = torch.device("cuda:0" if args.cuda else "cpu") if args.vis: from visdom import Visdom viz = Visdom(port=args.port) win = None train_envs = make_vec_envs(args.env_name, args.seed, args.num_processes, args.gamma, args.no_norm, args.num_stack, args.log_dir, args.add_timestep, device, allow_early_resets=False) if args.eval_interval: eval_seed = args.seed if args.seed is None else args.seed + args.num_processes eval_envs = make_vec_envs(args.env_name, eval_seed, args.num_processes // 4, args.gamma, args.no_norm, args.num_stack, eval_log_dir, args.add_timestep, device=device, allow_early_resets=True, eval=True, rank_offsest=args.num_processes) if eval_envs.venv.__class__.__name__ == "VecNormalize": eval_envs.venv.ob_rms = train_envs.venv.ob_rms else: eval_envs = None print(train_envs.observation_space.shape) noisy_net = True actor_critic = create_policy( train_envs.observation_space, train_envs.action_space, name='basic', nn_kwargs={ #'batch_norm': False if args.algo == 'acktr' else True, 'recurrent': 'lstm' if args.recurrent_policy else '', 'hidden_size': 512, }, noisy_net=noisy_net, train=True) if args.resume and os.path.isfile(args.resume): print('Resuming from checkpoint (%s)' % args.resume) state_dict, ob_rms = torch.load(args.resume, map_location='cpu') actor_critic.load_state_dict(state_dict) actor_critic.to(device) if args.algo.startswith('a2c'): agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef, args.entropy_coef, lr=args.lr, lr_schedule=lr_update_schedule, eps=args.eps, alpha=args.alpha, max_grad_norm=args.max_grad_norm) elif args.algo.startswith('ppo'): agent = algo.PPO(actor_critic, args.clip_param, args.ppo_epoch, args.num_mini_batch, args.value_loss_coef, args.entropy_coef, lr=args.lr, lr_schedule=lr_update_schedule, eps=args.eps, max_grad_norm=args.max_grad_norm) elif args.algo == 'acktr': agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef, args.entropy_coef, acktr=True) if args.algo.endswith('sil'): agent = algo.SIL(agent, update_ratio=args.sil_update_ratio, epochs=args.sil_epochs, batch_size=args.sil_batch_size, value_loss_coef=args.sil_value_loss_coef or args.value_loss_coef, entropy_coef=args.sil_entropy_coef or args.entropy_coef) replay = ReplayStorage(1e5, args.num_processes, args.gamma, 0.1, train_envs.observation_space.shape, train_envs.action_space, actor_critic.recurrent_hidden_state_size, device=device) else: replay = None rollouts = RolloutStorage(args.num_steps, args.num_processes, train_envs.observation_space.shape, train_envs.action_space, actor_critic.recurrent_hidden_state_size) obs = train_envs.reset() rollouts.obs[0].copy_(obs) rollouts.to(device) episode_rewards = deque(maxlen=10) start = time.time() for j in range(num_updates): if noisy_net: actor_critic.reset_noise() for step in range(args.num_steps): # Sample actions with torch.no_grad(): value, action, action_log_prob, recurrent_hidden_states = actor_critic.act( rollouts.obs[step], rollouts.recurrent_hidden_states[step], rollouts.masks[step]) # Obser reward and next obs obs, reward, done, infos = train_envs.step(action) for info in infos: if 'episode' in info.keys(): episode_rewards.append(info['episode']['r']) # If done then clean the history of observations. masks = torch.tensor([[0.0] if done_ else [1.0] for done_ in done], device=device) rollouts.insert(obs, recurrent_hidden_states, action, action_log_prob, value, reward, masks) if replay is not None: replay.insert(rollouts.obs[step], rollouts.recurrent_hidden_states[step], action, reward, done) with torch.no_grad(): next_value = actor_critic.get_value( rollouts.obs[-1], rollouts.recurrent_hidden_states[-1], rollouts.masks[-1]).detach() rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau) value_loss, action_loss, dist_entropy, other_metrics = agent.update( rollouts, j, replay) rollouts.after_update() if j % args.save_interval == 0 and args.save_dir != "": save_path = os.path.join(args.save_dir, args.algo) try: os.makedirs(save_path) except OSError: pass # A really ugly way to save a model to CPU save_model = actor_critic if args.cuda: save_model = copy.deepcopy(actor_critic).cpu() save_model = [ save_model.state_dict(), hasattr(train_envs.venv, 'ob_rms') and train_envs.venv.ob_rms or None ] torch.save(save_model, os.path.join(save_path, args.env_name + ".pt")) total_num_steps = (j + 1) * update_factor if j % args.log_interval == 0 and len(episode_rewards) > 1: end = time.time() print( "Updates {}, num timesteps {}, FPS {}, last {} mean/median reward {:.1f}/{:.1f}, " "min / max reward {:.1f}/{:.1f}, value/action loss {:.5f}/{:.5f}" .format(j, total_num_steps, int(total_num_steps / (end - start)), len(episode_rewards), np.mean(episode_rewards), np.median(episode_rewards), np.min(episode_rewards), np.max(episode_rewards), dist_entropy, value_loss, action_loss), end=', ' if other_metrics else '\n') if 'sil_value_loss' in other_metrics: print("SIL value/action loss {:.1f}/{:.1f}.".format( other_metrics['sil_value_loss'], other_metrics['sil_action_loss'])) if args.eval_interval and len( episode_rewards) > 1 and j > 0 and j % args.eval_interval == 0: actor_critic.eval() eval_episode_rewards = [] num_eval_processes = args.num_processes // 4 obs = eval_envs.reset() eval_recurrent_hidden_states = torch.zeros( 2, num_eval_processes, actor_critic.recurrent_hidden_state_size, device=device) eval_masks = torch.zeros(num_eval_processes, 1, device=device) while len(eval_episode_rewards) < 50: with torch.no_grad(): _, action, _, eval_recurrent_hidden_states = actor_critic.act( obs, eval_recurrent_hidden_states, eval_masks, deterministic=True) # Obser reward and next obs obs, reward, done, infos = eval_envs.step(action) eval_masks = torch.tensor([[0.0] if done_ else [1.0] for done_ in done], device=device) for info in infos: if 'episode' in info.keys(): eval_episode_rewards.append(info['episode']['r']) print(" Evaluation using {} episodes: mean reward {:.5f}\n".format( len(eval_episode_rewards), np.mean(eval_episode_rewards))) actor_critic.train() if args.vis and j % args.vis_interval == 0: try: # Sometimes monitor doesn't properly flush the outputs win = visdom_plot(viz, win, args.log_dir, args.env_name, args.algo, args.num_frames) except IOError: pass