def train(rank, args, shared_model, optimizer, env_conf): start_time = time.time() ptitle('Training Agent: {}'.format(rank)) #log = {} #setup_logger('{}_train_log'.format(args.env), r'{0}{1}_train_log'.format( # args.log_dir, args.env)) #log['{}_train_log'.format(args.env)] = logging.getLogger( # '{}_train_log'.format(args.env)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) if 'micropolis' in args.env.lower(): env = micropolis_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id if 'micropolis' in args.env.lower(): modelInit = getattr(model, args.design_head) player.model = modelInit(player.env.observation_space.shape[0], player.env.action_space, player.env.env.env.MAP_X) player.lstm_sizes = player.model.getMemorySizes() else: player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space) lstm_size = 512 if 'micropolis' in args.env.lower(): if 'arcade' not in args.env.lower(): lstm_size = (1, 16, env.env.env.MAP_X, env.env.env.MAP_Y) player.lstm_size = lstm_size player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() player.eps_len += 2 log_counter = 0 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) num_lstm_layers = len(player.lstm_sizes) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = [ Variable(torch.zeros(player.lstm_sizes[i]).cuda()) for i in range(num_lstm_layers) ] player.hx = [ Variable(torch.zeros(player.lstm_sizes[i]).cuda()) for i in range(num_lstm_layers) ] else: player.cx = [ Variable(torch.zeros(lstm_sizes[i])) for i in range(num_lstm_layers) ] player.hx = [ Variable(torch.zeros(lstm_sizes[i])) for i in range(num_lstm_layers) ] else: player.cx = [ Variable(player.cx[i].data) for i in range(num_lstm_layers) ] player.hx = [ Variable(player.hx[i].data) for i in range(num_lstm_layers) ] for step in range(args.num_steps): player.action_train() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if args.randomize_exploration: player.certainty = np.random.uniform(0.5, 1.5) if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: values, logit, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) if values.size()[1] == 1: value = values else: prob = torch.nn.functional.softmax(logit, dim=1) action = prob.multinomial(1).data value = values[0][action] R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = torch.zeros(1, 1).cuda() R = Variable(R).cuda() else: R = Variable(R) player.values.append(R) policy_loss = 0 value_loss = 0 for i in reversed(range(len(player.rewards))): if gpu_id >= 0: with torch.cuda.device(gpu_id): player.rewards[i] = torch.Tensor([player.rewards[i] ]).cuda() R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = Variable(gae.cuda()) else: gae = Variable(gae) policy_loss = policy_loss - \ player.log_probs[i] * Variable(gae) - 0.01 * player.entropies[i] #if log_counter % 10 == 0: # log['{}_train_log'.format(args.env)].info( # "Time {0}, reward {1}, policy loss {2}, value loss {3}, entropy {4}". # format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - start_time)), # '{:9.2e}'.format(float(sum(player.rewards) / len(player.rewards))), # '{:9.2e}'.format(float(policy_loss.data.item())), # '{:9.2e}'.format(float(value_loss.data.item())), # '{:10.8e}'.format(float(sum(player.entropies))))) #log_counter += 1 optimizer.zero_grad() a3c = args.lmbda * (policy_loss + 0.5 * value_loss) a3c.backward() torch.nn.utils.clip_grad_norm_(player.model.parameters(), 40) ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf, datasets): ptitle('Training Agent: {}'.format(rank)) print('Start training agent: ', rank) if rank == 0: logger = Logger(args.log_dir) train_step = 0 gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) raw, gt_lbl = datasets env = EM_env(raw, gt_lbl, env_conf) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) gamma = torch.tensor(1.0) if gpu_id >= 0: with torch.cuda.device(gpu_id): gamma = gamma.cuda() # env.seed (args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id # player.model = A3Clstm (env.observation_space.shape, env_conf["num_action"], args.hidden_feat) player.model = SimpleCNN(env.observation_space.shape, env_conf["num_action"]) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() if rank == 0: eps_reward = 0 pinned_eps_reward = 0 mean_log_prob = 0 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: player.eps_len = 0 if rank == 0: pinned_eps_reward = eps_reward eps_reward = 0 mean_log_prob = 0 for step in range(args.num_steps): player.action_train() # print ('step: ', step, 'reward_len: ', len (player.rewards)) if rank == 0: eps_reward += player.reward # mean_log_prob += player.log_probs [-1] if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _ = player.model(Variable(player.state.unsqueeze(0))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) # print ("updating -------------------") # print ("values:", player.values) # print ("gamma:", args.gamma) # print ("rewards:", player.rewards) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # print ("advatage: ", advantage) # print ("value_loss: ", value_loss) # print ("delta_t: ", player.values[i + 1].data + player.rewards[i]) # Generalized Advantage Estimataion delta_t = player.values[i + 1].data * args.gamma + player.rewards[i] - \ player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] player.model.zero_grad() sum_loss = (policy_loss + value_loss) sum_loss.backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions() if rank == 0: train_step += 1 if train_step % (args.log_period) == 0: log_info = { 'train: sum_loss': sum_loss, 'train: value_loss': value_loss, 'train: policy_loss': policy_loss, 'train: advanage': advantage, # 'entropy': entropy, 'train: eps reward': pinned_eps_reward, # 'mean log prob': mean_log_prob } for tag, value in log_info.items(): logger.scalar_summary(tag, value, train_step)
def train(rank, args, input_model=None, max_iter=100000, step_test=-1, log=False): if rank >= 0: torch.manual_seed(args.seed + rank) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = create_env(args) env.seed(args.seed + rank) if log: log = setup_logger("{0}_{1}_log".format(args.scale_legs, rank), "logs/{0}_{1}_log".format(args.scale_legs, rank)) # player initialization player = Agent(None, env, args, None) player.gpu_id = gpu_id if args.model == 'MLP': player.model = A3C_MLP(player.env.observation_space.shape[0], player.env.action_space, args.stack_frames) if args.model == 'CONV': player.model = A3C_CONV(args.stack_frames, player.env.action_space) # load the input model to the player if input_model != None: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(input_model.state_dict()) else: player.model.load_state_dict(input_model.state_dict()) # initialize the player optimizer optimizer = None if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(player.model.dictForOptimizer(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(player.model.dictForOptimizer(), lr=args.lr) else: optimizer = optim.SGD(player.model.dictForOptimizer(), lr=args.lr) # reset the environment and initialize the player state player.state = player.env.reset(args) player.state = torch.from_numpy(player.state).float() # If on GPU, do as GPU if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() last_iter = 0 mean_buf = Buffer(5) # Start looping over episodes for iteration in range(max_iter): last_iter += iteration # reset cx and hx if the enlvironmnent is over. if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 128).cuda()) player.hx = Variable(torch.zeros(1, 128).cuda()) else: player.cx = Variable(torch.zeros(1, 128)) player.hx = Variable(torch.zeros(1, 128)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) # Roll out actions and collect reward for one episode for step in range(args.num_steps): player.action_train() if player.done: break if player.done: player.eps_len = 0 # reset state state = player.env.reset(args) player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() if gpu_id >= 0: with torch.cuda.device(gpu_id): R = torch.zeros(1, 1).cuda() else: R = torch.zeros(1, 1) if not player.done: state = player.state if args.model == 'CONV': state = state.unsqueeze(0) value, _, _, _ = player.model( (Variable(state), (player.hx, player.cx))) R = value.data player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 R = Variable(R) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = torch.zeros(1, 1).cuda() else: gae = torch.zeros(1, 1) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ (player.log_probs[i].sum() * Variable(gae)) - \ (0.01 * player.entropies[i].sum()) player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() optimizer.step() player.clear_actions() if step_test > 0 and iteration % step_test == 0: tester = Tester(args, player.model) score = tester.test(last_iter) mean_buf.push(score) recent_mean = sum(mean_buf.bf) / mean_buf.current_size text = "Iteration {0}, episode reward {1}, recent reward mean {2}".format( iteration, score, recent_mean) log.info(text) tester = Tester(args, player.model) fitness = tester.test(last_iter) return fitness
def train(rank, args, shared_model, optimizer): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = create_env(args.env, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id if args.model == 'MLP': player.model = A3C_MLP(player.env.observation_space.shape[0], player.env.action_space, args.stack_frames) if args.model == 'CONV': player.model = A3C_CONV(args.stack_frames, player.env.action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 128).cuda()) player.hx = Variable(torch.zeros(1, 128).cuda()) else: player.cx = Variable(torch.zeros(1, 128)) player.hx = Variable(torch.zeros(1, 128)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train() if player.done: break if player.done: player.eps_len = 0 state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() if gpu_id >= 0: with torch.cuda.device(gpu_id): R = torch.zeros(1, 1).cuda() else: R = torch.zeros(1, 1) if not player.done: state = player.state if args.model == 'CONV': state = state.unsqueeze(0) value, _, _, _ = player.model( (Variable(state), (player.hx, player.cx))) R = value.data player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 R = Variable(R) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = torch.zeros(1, 1).cuda() else: gae = torch.zeros(1, 1) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion # print(player.rewards[i]) delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ (player.log_probs[i].sum() * Variable(gae)) - \ (0.01 * player.entropies[i].sum()) player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf): ptitle('Training Agent: {}'.format(rank)) print("prank:", rank, "os.pid:", os.getpid()) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = AllowBacktracking( make_local_env(env_conf['game'], env_conf['level'], stack=False, scale_rew=False)) print("Got a local env; obs space:", env.observation_space) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() print("player.state.shape:", player.state.shape) player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() player.eps_len += 2 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 512).cuda()) player.hx = Variable(torch.zeros(1, 512).cuda()) else: player.cx = Variable(torch.zeros(1, 512)) player.hx = Variable(torch.zeros(1, 512)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train() if player.done: break if player.done: # if player.info['ale.lives'] == 0 or player.max_length: # player.eps_len = 0 state = player.env.reset() player.eps_len += 2 player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() torch.nn.utils.clip_grad_norm(player.model.parameters(), 100.0) ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(args, envs, observation_space, action_space): gpu_id = 0 #每个单独的work,独立的环境和model,在cuda中运行 player = Agent(envs, args) player.model = A3Clstm(observation_space, action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() with torch.cuda.device(gpu_id): player.model = player.model.cuda() player.state = player.state.cuda() player.cx = torch.zeros(args.workers, 512).cuda() player.hx = torch.zeros(args.workers, 512).cuda() optimizer = torch.optim.Adam(player.model.parameters(), lr=args.lr, amsgrad=args.amsgrad) #切换到训练模式 player.model.train() while True: #训练20步或者game over就结束训练 for step in range(args.num_steps): #训练时,保存每一步的相关信息到list player.env.get_images() player.action_train() if player.dones[-1][0]: break if not player.dones[-1][0]: value, _, _ = player.model((player.state, (player.hx, player.cx))) R = value.detach() else: R = torch.zeros(args.workers, 1) with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(R) for j in range(args.num_ppo_train): policy_loss = 0 value_loss = 0 gae = 0 for i in reversed(range(len(player.rewards))): value, logit, _ = player.model( (player.states[i], (player.hxs[i], player.cxs[i]))) prob = F.softmax(logit, dim=1) log_prob = F.log_softmax(logit, dim=1) entropy = -(log_prob * prob).sum(1) log_probs_current = log_prob.gather(1, player.actions[i]) R = args.gamma * R + player.rewards[i] advantage = R - value value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * player.values[ i + 1].detach() - player.values[i].detach() gae = gae * args.gamma * args.tau + delta_t ratio = torch.exp(log_probs_current - player.log_probs[i]) surr1 = ratio surr2 = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) policy_loss = policy_loss - torch.min( surr1, surr2) * gae * -0.01 * entropy optimizer.zero_grad() (policy_loss + 0.5 * value_loss).mean().backward() optimizer.step() #game over时记忆值重置 if player.dones[-1][0]: with torch.cuda.device(gpu_id): player.cx = torch.zeros(args.workers, 512).cuda() player.hx = torch.zeros(args.workers, 512).cuda() else: player.cx = player.cx.detach() player.hx = player.hx.detach() player.clear_actions() # advantage[0:n] # 第0,1,2,...n 到 n+1的估值差 r[0-n],r[1-n],r[2-n]....rn Value(N+1) 取反: # 第n,n-1,n-2,n-3,......3,2,1 # r[n] + Value(N+1) - Value(N) # r[n:n-1] + Value(N+1) - Value(N-1) # ... # r[n:2] + Value(N + 1) - Value(2) # r[n:1] + Value(N + 1) - Value(1) # R = args.gamma * R + player.rewards[i] # advantage = R - player.values[i] # value_loss = value_loss + 0.5 * advantage.pow(2) # value_loss = 0.5 * advantage.pow(2) # advantage = args.gamma * R + player.rewards[i] - player.values[i] #entropy = -(log_prob * prob).sum(1) #self.entropies.append(entropy) #通过prob 采样对应的动作和动作logprob # 计算每次的概率和entropy(entropies)和entropy的sum,sum是每一步所有动作概率的熵值
def train(rank, args, shared_model, optimizer, train_modes, n_iters, env=None): n_steps = 0 n_iter = 0 writer = SummaryWriter(os.path.join(args.log_dir, 'Agent:{}'.format(rank))) ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) training_mode = args.train_mode env_name = args.env train_modes.append(training_mode) n_iters.append(n_iter) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) device = torch.device('cuda:' + str(gpu_id)) else: device = torch.device('cpu') if env == None: env = create_env(env_name) params = shared_model.parameters() if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(params, lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(params, lr=args.lr) env.seed(args.seed + rank) player = Agent(None, env, args, None, device) player.gpu_id = gpu_id player.env.reset() # prepare model player.model = build_model(action_space=player.env.action_space, pose_space=player.reset_cam_pose(), args=args,) player.model = player.model.to(device) player.model.train() player.reset() reward_sum = torch.zeros(player.num_agents).to(device) count_eps = 0 print('Start training...') while True: # sys to the shared model player.model.load_state_dict(shared_model.state_dict()) if player.done: player.reset() reward_sum = torch.zeros(player.num_agents).to(device) count_eps += 1 player.update_rnn_hidden() fps_counter = 0 t0 = time.time() for i in range(args.num_steps): player.action_train() reward_sum += player.reward fps_counter += 1 n_steps += 1 if player.done: for i, r_i in enumerate(reward_sum): # add for Pose Only if i not in player.env.random_ids: continue # writer.add_scalar('train/reward_' + str(i), r_i, n_steps) break fps = fps_counter / (time.time() - t0) policy_loss, value_loss, entropies, pred_loss, values0 = player.optimize(params, optimizer, shared_model, gpu_id) writer.add_scalar('train/policy_loss_sum', policy_loss.sum(), n_steps) writer.add_scalar('train/value_loss_sum', value_loss.sum(), n_steps) writer.add_scalar('train/entropies_sum', entropies.sum(), n_steps) writer.add_scalar('train/values0', values0.sum(), n_steps) writer.add_scalar('train/pred_R_loss', pred_loss, n_steps) writer.add_scalar('train/fps', fps, n_steps) # writer.add_scalar('train/lr', lr[0], n_iter) n_iter += 1 n_iters[rank] = n_iter if train_modes[rank] == -100: env.close() break
def train(rank, args, shared_model, optimizer, train_modes, n_iters, device, env=None): n_steps = 0 n_iter = 0 writer = SummaryWriter(os.path.join(args.log_dir, 'Agent:{}'.format(rank))) ptitle('Training Agent: {}'.format(rank)) torch.manual_seed(args.seed + rank) training_mode = args.train_mode env_name = args.env train_modes.append(training_mode) n_iters.append(n_iter) if env == None: env = create_env(env_name, args) params = shared_model.parameters() if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(params, lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(filter(lambda p: p.requires_grad, shared_model.parameters()), lr=args.lr) env.seed(args.seed + rank) player = Agent(None, env, args, None, None, device) player.model = build_model( player.env.observation_space, player.env.action_space, args, device).to(device) player.state = player.env.reset() if 'Unreal' in args.env: player.cam_pos = player.env.env.env.env.cam_pose player.collect_state = player.env.env.env.env.current_states player.set_cam_info() player.state = torch.from_numpy(player.state).float() player.state = player.state.to(device) player.model = player.model.to(device) player.model.train() reward_sum = torch.zeros(player.num_agents).to(device) count_eps = 0 cross_entropy_loss = nn.CrossEntropyLoss() while True: player.model.load_state_dict(shared_model.state_dict()) player.update_lstm() fps_counter = 0 t0 = time.time() for step in range(args.num_steps): player.action_train() n_steps += 1 reward_sum += player.reward if player.done: break update_steps = len(player.rewards) fps = fps_counter / (time.time() - t0) if player.done: for i in range(player.num_agents): writer.add_scalar('train/reward_'+str(i), reward_sum[i], n_steps) count_eps += 1 reward_sum = torch.zeros(player.num_agents).to(device) player.eps_len = 0 player.state = player.env.reset() player.set_cam_info() player.state = torch.from_numpy(player.state).float().to(device) R = torch.zeros(player.num_agents, 1).to(device) if not player.done: state = player.state value_multi, _, _, _, _, _, _, _ , _= player.model( (Variable(state, requires_grad=True), Variable((player.cam_info), requires_grad=True), player.H_multi, player.last_gate_ids, player.gt_gate)) for i in range(player.num_agents): R[i][0] = value_multi[i].data gates, gt_gates = [], [] for k1 in range(len(player.rewards)): for k2 in range(player.num_agents): gates.append(player.gates[k1][k2]) gt_gates.append(player.gate_gts[k1][k2]) gate_probs = torch.cat(gates).view(-1, 2).to(device) gate_gt_ids = torch.Tensor(gt_gates).view(1, -1).squeeze().long().to(device) gate_loss = cross_entropy_loss(gate_probs, gate_gt_ids) player.values.append(Variable(R).to(device)) policy_loss = torch.zeros(player.num_agents, 1).to(device) value_loss = torch.zeros(player.num_agents, 1).to(device) entropies = torch.zeros(player.num_agents, 1).to(device) w_entropies = torch.Tensor([[float(args.entropy)] for i in range(player.num_agents)]).to(device) R = Variable(R, requires_grad=True).to(device) gae = torch.zeros(1, 1).to(device) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ (player.log_probs[i] * Variable(gae)) - \ (w_entropies * player.entropies[i]) entropies += player.entropies[i] loss = policy_loss.sum() / update_steps / player.num_agents + 0.5 * value_loss.sum() / update_steps / player.num_agents + \ 5 * gate_loss player.model.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(params, 50) ensure_shared_grads(player.model, shared_model, gpu=args.gpu_ids[-1] >= 0) writer.add_scalar('train/policy_loss_sum', policy_loss.sum(), n_steps) writer.add_scalar('train/value_loss_sum', value_loss.sum(), n_steps) writer.add_scalar('train/entropies_sum', entropies.sum(), n_steps) writer.add_scalar('train/fps', fps, n_steps) writer.add_scalar('train/gate_loss', gate_loss, n_steps) n_iter += 1 n_iters[rank] = n_iter optimizer.step() player.clear_actions() if train_modes[rank] == -100: env.close() break
def train(rank, args, shared_model, optimizer, env_conf): ptitle('Train {0}: {1}'.format(args.env, rank)) print('Start training agent: ', rank) if rank == 0: logger = Logger(args.log_dir + '_losses/') train_step = 0 gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] env_conf["env_gpu"] = gpu_id torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = database_env(env_conf, seed=0) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) player = Agent(None, env, args, None, gpu_id) player.gpu_id = gpu_id player.model = get_model(args, args.model, env_conf["observation_shape"], args.features, env_conf["num_actions"], gpu_id=0, lstm_feats=args.lstm_feats) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() if rank == 0: eps_reward = 0 pinned_eps_reward = 0 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: player.eps_len = 0 if rank == 0: if train_step % args.train_log_period == 0 and train_step > 0: print("train: step", train_step, "\teps_reward", eps_reward) if train_step > 0: pinned_eps_reward = player.env.sum_reward eps_reward = 0 if args.lstm_feats: player.cx, player.hx = init_linear_lstm( args.lstm_feats, gpu_id) elif args.lstm_feats: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train() if rank == 0: eps_reward = player.env.sum_reward if player.done: break if player.done: if rank == 0: if train_step % args.train_log_period == 0 and train_step > 0: print("train: step", train_step, "\teps_reward", eps_reward) # print ("rewards: ", player.env.rewards) # print ("actions: ", player.actions) if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1, 1, 1) if not player.done: if args.lstm_feats: value, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) else: value, _ = player.model(Variable(player.state.unsqueeze(0))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1, 1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): if gpu_id >= 0: with torch.cuda.device(gpu_id): reward_i = torch.tensor(player.rewards[i]).cuda() else: reward_i = torch.tensor(player.rewards[i]) R = args.gamma * R + reward_i advantage = R - player.values[i] value_loss = value_loss + (0.5 * advantage * advantage).mean() delta_t = player.values[ i + 1].data * args.gamma + reward_i - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ (player.log_probs[i] * Variable(gae)).mean () - \ (args.entropy_alpha * player.entropies[i]).mean () player.model.zero_grad() sum_loss = (policy_loss + value_loss) sum_loss.backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions() if rank == 0: train_step += 1 if train_step % args.log_period == 0 and train_step > 0: log_info = { 'sum_loss': sum_loss, 'value_loss': value_loss, 'policy_loss': policy_loss, 'advanage': advantage, 'train eps reward': pinned_eps_reward, } for tag, value in log_info.items(): logger.scalar_summary(tag, value, train_step)
def train(rank, args, shared_model, optimizer): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] writer = SummaryWriter(log_dir=args.log_dir + 'tb_train') log = {} setup_logger('{}_train_log'.format(rank), r'{0}{1}_train_log'.format(args.log_dir, rank)) log['{}_train_log'.format(rank)] = logging.getLogger( '{}_train_log'.format(rank)) torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(env_id=rank, args=args, type='train') if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Clstm(player.env.observation_space.shape[2], player.env.action_space.n) player.state = player.env.reset() player.state = normalize_rgb_obs(player.state) player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() num_trains = 0 if not os.path.exists(args.log_dir + "images/"): os.makedirs(args.log_dir + "images/") while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) for step in range(args.num_steps): player.action_train() if player.done: break if player.done: num_trains += 1 log['{}_train_log'.format(rank)].info('entropy:{0}'.format( player.entropy.data[0])) writer.add_scalar("data/entropy_" + str(rank), player.entropy.data[0], num_trains) writer.add_image('FCN_' + str(rank), player.fcn, num_trains) writer.add_image('Depth_GroundTruth_' + str(rank), player.depth, num_trains) writer.add_image('RGB_' + str(rank), player.env.get_rgb(), num_trains) save_image( player.fcn.data, args.log_dir + "images/" + str(rank) + "_" + str(num_trains) + "_fcn.png") # print("player.fcn.data:", player.fcn.data) save_image( player.depth.data, args.log_dir + "images/" + str(rank) + "_" + str(num_trains) + "_depth.png") cv2.imwrite( args.log_dir + "images/" + str(rank) + "_" + str(num_trains) + "_rgb.png", player.env.get_rgb()) # print("player.depth.data:", player.depth.data) player.eps_len = 0 player.current_life = 0 state = player.env.reset() state = normalize_rgb_obs(state) player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: with torch.cuda.device(gpu_id): value, _, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx), Variable( torch.from_numpy(player.env.target).type( torch.FloatTensor).cuda()))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = args.gamma * player.values[ i + 1].data + player.rewards[i] - player.values[i].data gae = gae * args.gamma * args.tau + delta_t # policy_loss = policy_loss - \ # player.log_probs[i] * \ # Variable(gae) - 0.01 * player.entropies[i] \ # + player.fcn_losses[i] # FCN policy_loss = policy_loss - 1e-5*(player.log_probs[i] * Variable(gae)) - 1e-5*(0.01 * player.entropies[i]) \ + player.fcn_losses[i] * DEPTH_LOSS_DISCOUNT # FCN # policy_loss = policy_loss + player.fcn_losses[i] # FCN writer.add_scalar("data/value_loss_" + str(rank), value_loss, num_trains) writer.add_scalar("data/policy_loss_" + str(rank), policy_loss, num_trains) player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() torch.nn.utils.clip_grad_norm(player.model.parameters(), 40.0) ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(rank, args, shared_model, optimizer): # optimizer为shared_model的 init = True ptitle('Training Agent: {}'.format(rank)) torch.manual_seed(args.seed + rank) env = create_env(args.env, args.seed + rank) # env = gym.make(args.env) # env.seed(args.seed + rank) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr) player = Agent(None, env, args, None, rank) player.model = A3C_MLP(player.env.observation_space, player.env.action_space, args.stack_frames) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() player.model.train() #固定使用场景为train if rank == 1: # file = open(os.path.join(args.log_dir, 'TD_Error.txt'), 'w+') writer = SummaryWriter('8_27_train') local_step_counter = 0 while True: if init: # 初始化 shared_model.training_steps.weight.data \ .copy_(torch.Tensor([0])) shared_model.training_steps.bias.data \ .copy_(torch.Tensor([0])) init = False player.model.load_state_dict( shared_model.state_dict()) # synchronize parameters for step in range(args.num_steps): # print("thread", rank, local_step_counter, shared_model.training_steps.weight.data.cpu().numpy()) local_step_counter += 1 # update step counters shared_model.training_steps.weight.data \ .copy_(torch.Tensor([1]) + shared_model.training_steps.weight.data) # 总步骤(各个worker所走步数之和)T每次加一 player.action_train() # core if player.done: break terminal = False if player.done or player.eps_len >= args.max_episode_length: # 玩家完成或者超出最大迭代次数 terminal = True shared_model.done_nums += 1 if 'is_success' in player.info.keys(): shared_model.success_num += 1 R = torch.zeros(1) if not player.done: # 结算 state = player.state # A3C,value和policy net是用的同一个网络 value, _, _ = player.model(Variable(state)) R = value.data if terminal: #重置 shared_model.training_steps.bias.data \ .copy_(torch.Tensor([1]) + shared_model.training_steps.bias.data) # 总步数加一 player.eps_len = 0 state = player.env.reset() player.state = torch.from_numpy(state).float() player.reset_flag = True player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 R = Variable(R) gae = torch.zeros(1, 1) for i in reversed(range(len(player.rewards))): R = args.gamma * R + np.float(player.rewards[i]) # reward advantage = R - player.values[i] # advantage value_loss = value_loss + 0.5 * advantage.pow(2) # 公式(10) 更新w if rank == 1: # file.write(str(advantage.pow(2).data.cpu().numpy()[0])) # file.write(' ') # file.write( # str(int(shared_model.training_steps.weight.data.cpu().numpy()[0]))) # file.write('\n') writer.add_scalar( 'TD-error/train', advantage.pow(2).data.cpu().numpy()[0], shared_model.training_steps.weight.data.cpu().numpy()[0]) player.values[i] = player.values[i].float() player.values[i + 1] = player.values[i + 1].float() delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - \ player.values[i].data # a2c计算td-error # GAE算法 gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ (player.log_probs[i].sum() * Variable(gae)) - \ (0.01 * player.entropies[i].sum()) # 更新theta 公式(9) """ 每个线程和环境交互到一定量的数据后,就计算在自己线程里的神经网络损失函数的梯度, 但是这些梯度却并不更新自己线程里的神经网络,而是去更新公共的神经网络。 也就是n个线程会独立的使用累积的梯度分别更新公共部分的神经网络模型参数。 每隔一段时间,线程会将自己的神经网络的参数更新为公共神经网络的参数,进而指导后面的环境交互。 """ player.model.zero_grad() # policy_loss + 0.5 * value_loss即为loss if rank == 1: writer.add_scalar( 'VLoss/train', value_loss, shared_model.training_steps.weight.data.cpu().numpy()[0]) writer.add_scalar( 'PLoss/train', policy_loss, shared_model.training_steps.weight.data.cpu().numpy()[0]) (policy_loss + 0.5 * value_loss).backward() # 计算该worder的损失函数梯度 ensure_shared_grads(player.model, shared_model) # 该worker将自己的参数传给公用的模型 optimizer.step( ) # optimizer为shared_model的 step()将参数更新值施加到shared_model的parameters 上 player.clear_actions() if shared_model.training_steps.weight.data.cpu().numpy( ) > args.training_steps: print('num of success={0},training episodes={1},success_rate={2}'. format(shared_model.success_num, shared_model.done_nums, shared_model.success_num / shared_model.done_nums)) break
def test(args, shared_model, optimizer, env_conf): ptitle('Test Agent') gpu_id = args.gpu_ids[-1] start_time = datetime.now().strftime('%Y-%m-%d_%H_%M_%S') log = {} setup_logger('{}_log'.format(args.env), r'{0}{1}_{2}_log'.format( args.log_dir, args.env, start_time)) log['{}_log'.format(args.env)] = logging.getLogger('{}_log'.format( args.env)) d_args = vars(args) for k in d_args.keys(): log['{}_log'.format(args.env)].info('{0}: {1}'.format(k, d_args[k])) if not os.path.exists(args.save_model_dir): os.mkdir(args.save_model_dir) if args.seed: torch.manual_seed(args.seed) if gpu_id >= 0: torch.cuda.manual_seed(args.seed) env = atari_env(args.env, env_conf, args) reward_sum = 0 start = time.time() num_tests = 0 reward_total_sum = 0 player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Cff(player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() player.eps_len += 2 player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model = player.model.cuda() player.state = player.state.cuda() flag = True max_score = -10000 while True: p = optimizer.param_groups[0]['params'][0] step = optimizer.state[p]['step'] player.model.eval() if flag: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) flag = False with torch.no_grad(): if args.robust: #player.action_test_losses(args.epsilon_end) lin_coeff = min(1, (1.5*int(step)+1)/(args.total_frames/args.num_steps)) epsilon = lin_coeff*args.epsilon_end player.action_train(epsilon) else: player.action_train() #player.action_test_losses() reward_sum += player.noclip_reward if player.done and not player.info: state = player.env.reset() player.eps_len += 2 player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() elif player.info: # calculate losses for tracking R = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(R) gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) standard_loss = 0 worst_case_loss = 0 value_loss = 0 entropy = 0 for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss += 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t if args.robust: if advantage >= 0: worst_case_loss += - player.min_log_probs[i] * Variable(gae) else: worst_case_loss += - player.max_log_probs[i] * Variable(gae) standard_loss += -player.log_probs[i] * Variable(gae) entropy += player.entropies[i] standard_loss = standard_loss/len(player.rewards) worst_case_loss = worst_case_loss/len(player.rewards) value_loss = value_loss/len(player.rewards) entropy = entropy/len(player.rewards) player.clear_actions() flag = True num_tests += 1 reward_total_sum += reward_sum reward_mean = reward_total_sum / num_tests log['{}_log'.format(args.env)].info( ("Time {0}, steps {1}/{2}, ep reward {3}, ep length {4}, reward mean {5:.3f} \n"+ "Losses: Policy:{6:.3f}, Worst case: {7:.3f}, Value: {8:.3f}, Entropy: {9:.3f}"). format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - start)), int(step), args.total_frames/args.num_steps, reward_sum, player.eps_len, reward_mean, float(standard_loss), float(worst_case_loss), float(value_loss), float(entropy))) if args.save_max and reward_sum >= max_score: max_score = reward_sum if gpu_id >= 0: with torch.cuda.device(gpu_id): state_to_save = player.model.state_dict() torch.save(state_to_save, '{0}{1}_{2}_best.pt'.format( args.save_model_dir, args.env, start_time)) else: state_to_save = player.model.state_dict() torch.save(state_to_save, '{0}{1}_{2}_best.pt'.format( args.save_model_dir, args.env, start_time)) reward_sum = 0 player.eps_len = 0 state = player.env.reset() player.eps_len += 2 #stop after total steps gradient updates have passed if step >= args.total_frames/args.num_steps: if gpu_id >= 0: with torch.cuda.device(gpu_id): state_to_save = player.model.state_dict() torch.save(state_to_save, '{0}{1}_{2}_last.pt'.format( args.save_model_dir, args.env, start_time)) else: state_to_save = player.model.state_dict() torch.save(state_to_save, '{0}{1}_{2}_last.pt'.format( args.save_model_dir, args.env, start_time)) return time.sleep(10) player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda()
def train(rank, args, shared_models, optimizers, env_conf): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) env.seed(args.seed + rank) player = Agent(env, args, gpu_id) player.rank = rank player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.models[0].train() player.models[1].train() player.eps_len += 2 # player.test_models() while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): # player.model.load_state_dict(shared_model.state_dict()) player.models[0].load_state_dict(shared_models[0].state_dict()) player.models[1].load_state_dict(shared_models[1].state_dict()) else: # player.model.load_state_dict(shared_model.state_dict()) player.models[0].load_state_dict(shared_models[0].state_dict()) player.models[1].load_state_dict(shared_models[1].state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 512).cuda()) player.hx = Variable(torch.zeros(1, 512).cuda()) else: player.cx = Variable(torch.zeros(1, 512)) player.hx = Variable(torch.zeros(1, 512)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train() if player.done: break # if rank == 0: # print(player.episodic_reward) player.episodic_reward = 0 if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _, _ = player.models[player.curr_model_id]((Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() # player.values.append(Variable(R)) gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) # print("Length of values vector", len(player.values)) # print("Length of rewards vector", len(player.rewards)) # print("Length of model sequence vector", len(player.model_sequence)) next_val = Variable(R) last_val = next_val R_vec = [Variable(R), Variable(R)] # last_id = player.model_sequence[-1] active_flags = [False, False] policy_loss = [0, 0] value_loss = [0, 0] for reward, value, model_id, log_prob, entropy in zip( reversed(player.rewards), reversed(player.values), reversed(player.model_sequence), reversed(player.log_probs), reversed(player.entropies) ): active_flags[model_id] = True R_vec[model_id] = args.gamma * R_vec[model_id] + reward R_vec[(model_id+1)%2] *= args.gamma advantage = R_vec[model_id] - value value_loss[model_id] += 0.5 * advantage.pow(2) delta_t = reward + args.gamma * next_val.data - value.data gae = gae * args.gamma * args.tau + delta_t policy_loss[model_id] -= (log_prob * Variable(gae) + 0.01 * entropy) next_val = value try: if active_flags[0] is True: player.models[0].zero_grad() (policy_loss[0] + 0.5 * value_loss[0]).backward() ensure_shared_grads(player.models[0], shared_models[0], gpu = gpu_id >= 0) optimizers[0].step() if active_flags[1] is True: player.models[1].zero_grad() (policy_loss[1] + 0.5 * value_loss[1]).backward() ensure_shared_grads(player.models[1], shared_models[1], gpu = gpu_id >= 0) optimizers[1].step() except Exception as e: print("Exception caught. Ignoring") if rank == 1: print(rewards) print(model_sequence) player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) tp_weight = args.tp player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space, args.terminal_prediction, args.reward_prediction) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() # Below is where the cores are running episodes continously ... average_ep_length = 0 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 128).cuda()) player.hx = Variable(torch.zeros(1, 128).cuda()) else: player.cx = Variable(torch.zeros(1, 128)) player.hx = Variable(torch.zeros(1, 128)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.eps_len += 1 player.action_train() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _, _, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 reward_pred_loss = 0 terminal_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) # TODO why this is here? for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * player.values[ i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - player.log_probs[i] * Variable( gae) - 0.01 * player.entropies[i] if args.reward_prediction: reward_pred_loss = reward_pred_loss + ( player.reward_predictions[i] - player.rewards[i]).pow(2) if args.terminal_prediction: # new way of using emprical episode length as a proxy for current length. if player.average_episode_length is None: end_predict_labels = np.arange( player.eps_len - len(player.terminal_predictions), player.eps_len) / player.eps_len # heuristic else: end_predict_labels = np.arange( player.eps_len - len(player.terminal_predictions), player.eps_len) / player.average_episode_length for i in range(len(player.terminal_predictions)): terminal_loss = terminal_loss + ( player.terminal_predictions[i] - end_predict_labels[i]).pow(2) terminal_loss = terminal_loss / len(player.terminal_predictions) player.model.zero_grad() #print(f"policy loss {policy_loss} and value loss {value_loss} and terminal loss {terminal_loss} and reward pred loss {reward_pred_loss}") total_loss = policy_loss + 0.5 * value_loss + tp_weight * terminal_loss + 0.5 * reward_pred_loss total_loss.backward() # will free memory ... # Visualize Computation Graph #graph = make_dot(total_loss) #from graphviz import Source #Source.view(graph) ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions() if player.done: if player.average_episode_length is None: # initial one player.average_episode_length = player.eps_len else: player.average_episode_length = int( 0.99 * player.average_episode_length + 0.01 * player.eps_len) #print(player.average_episode_length, 'current one is ', player.eps_len) player.eps_len = 0 # reset here
def train_func(rank, args, shared_model, optimizer, env_conf, datasets): if args.deploy: return ptitle('Train {0}'.format(rank)) print('Start training agent: ', rank) if rank == 0: logger = Logger(args.log_dir[:-1] + '_losses/') train_step = 0 gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] env_conf["env_gpu"] = gpu_id torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = Debug_env(datasets, env_conf, seed=args.seed + rank) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) player = Agent(None, env, args, None) player.gpu_id = gpu_id nChan = 3 if args.is3D: nChan = 4 if args.alpha_only: nChan = 1 if not args.is3D: player.model = get_model(args, "ENet", input_shape=env_conf["obs_shape"], num_actions=args.num_actions * nChan) elif not args.obs3D: player.model = get_model(args, "ENet", input_shape=env_conf["obs_shape"], num_actions=args.num_actions * nChan) elif args.obs3D: player.model = get_model(args, "Net3D", input_shape=env_conf["obs_shape"], num_actions=args.num_actions * nChan) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() if rank == 0: eps_reward = 0 pinned_eps_reward = 0 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: player.eps_len = 0 if rank == 0: if train_step % args.train_log_period == 0 and train_step > 0: print("train: step", train_step, "\teps_reward", eps_reward) if train_step > 0: pinned_eps_reward = player.env.sum_rewards.mean() eps_reward = 0 for step in range(args.num_steps): player.action_train() if rank == 0: eps_reward = player.env.sum_rewards.mean() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() if not args.alpha_only: if not args.is3D: R = torch.zeros(1, 1, args.num_actions * 3) else: R = torch.zeros(1, 1, args.num_actions * 4) else: R = torch.zeros(1, 1, args.num_actions) if not player.done: value, _ = player.model(Variable(player.state.unsqueeze(0))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 if not args.alpha_only: if not args.is3D: gae = torch.zeros(1, 1, args.num_actions * 3) else: gae = torch.zeros(1, 1, args.num_actions * 4) else: gae = torch.zeros(1, 1, args.num_actions) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): if gpu_id >= 0: with torch.cuda.device(gpu_id): reward_i = torch.tensor(player.rewards[i]).cuda() else: reward_i = torch.tensor(player.rewards[i]) R = args.gamma * R + reward_i advantage = R - player.values[i] value_loss = value_loss + (0.5 * advantage * advantage).mean() delta_t = player.values[ i + 1].data * args.gamma + reward_i - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ (player.log_probs[i] * Variable(gae)).mean () - \ (args.entropy_alpha * player.entropies[i]).mean () player.model.zero_grad() sum_loss = (policy_loss + value_loss) curtime = time.time() sum_loss.backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) curtime = time.time() optimizer.step() player.clear_actions() if rank == 0: train_step += 1 if train_step % args.log_period * 10 == 0 and train_step > 0: log_info = { 'train: value_loss': value_loss, 'train: policy_loss': policy_loss, 'train: eps reward': pinned_eps_reward, } for tag, value in log_info.items(): logger.scalar_summary(tag, value, train_step)
def train_robust(rank, args, shared_model, optimizer, env_conf): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] if args.seed: torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) if args.seed: env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Cff(player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() player.eps_len += 2 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) p = optimizer.param_groups[0]['params'][0] step = optimizer.state[p]['step'] if step >= (args.total_frames / args.num_steps): return #increase linearly until 2/3 through halfway lin_coeff = min(1, (1.5 * int(step) + 1) / (args.total_frames / args.num_steps)) epsilon = lin_coeff * args.epsilon_end kappa = args.kappa_end #(1-lin_coeff)*1 + lin_coeff*args.kappa_end for step in range(args.num_steps): player.action_train(bound_epsilon=epsilon) if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _ = player.model(Variable(player.state.unsqueeze(0))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t if gae >= 0: worst_case_loss = -player.min_log_probs[i] * Variable(gae) else: worst_case_loss = -player.max_log_probs[i] * Variable(gae) standard_loss = -player.log_probs[i] * Variable(gae) policy_loss = policy_loss + kappa * standard_loss + ( 1 - kappa) * worst_case_loss - 0.01 * player.entropies[i] #print(policy_loss + 0.5 * value_loss) player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = Environment() # 創建環境 if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) # env.seed(args.seed + rank) player = Agent(None, env, args, None) # 創建代理人 player.gpu_id = gpu_id num_actions = env.get_num_actions() player.model = A3Clstm( Config.STACKED_FRAMES, # A3C模型 num_actions) player.state, available = player.env.reset() # 初始環境 player.state = torch.from_numpy(player.state).float() player.available = torch.from_numpy(available).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.available = player.available.cuda() player.model.train() # 訓練模式 player.eps_len += 1 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) # 更新網路 if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 512).cuda()) player.hx = Variable(torch.zeros(1, 512).cuda()) else: player.cx = Variable(torch.zeros(1, 512)) player.hx = Variable(torch.zeros(1, 512)) # 完成一次訓練 初始化LSTM else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): # T-max = 20 player.action_train() if player.done: break if player.done: state, available = player.env.reset() player.state = torch.from_numpy(state).float() player.available = torch.from_numpy(available).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.available = player.available.cuda() R = torch.zeros(1, 1) # if done : R_t-max = 0 if not player.done: value, _, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data # R_t-max = V(s) if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf, iters, checkpoint_path): iters = dill.loads(iters) if args.enable_gavel_iterator and rank == 0: iters._init_logger() ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() player.eps_len += 2 elapsed_time = 0 start_time = time.time() for i in iters: if i % 100 == 0: print('GPU %d finished step %d' % (rank, i), flush=True) if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 512).cuda()) player.hx = Variable(torch.zeros(1, 512).cuda()) else: player.cx = Variable(torch.zeros(1, 512)) player.hx = Variable(torch.zeros(1, 512)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions() elapsed_time += time.time() - start_time start_time = time.time() if (args.throughput_estimation_interval is not None and i % args.throughput_estimation_interval == 0 and rank == 0): print('[THROUGHPUT_ESTIMATION]\t%s\t%d' % (time.time(), i)) if (args.max_duration is not None and elapsed_time >= args.max_duration): break if args.enable_gavel_iterator and rank == 0: state = shared_model.state_dict() iters.save_checkpoint(state, checkpoint_path) iters.complete()
def train(rank, args, shared_model, optimizer, train_modes, n_iters, env=None): n_iter = 0 writer = SummaryWriter(os.path.join(args.log_dir, 'Agent:{}'.format(rank))) ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) training_mode = args.train_mode env_name = args.env train_modes.append(training_mode) n_iters.append(n_iter) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) device = torch.device('cuda:' + str(gpu_id)) if len(args.gpu_ids) > 1: device_share = torch.device('cpu') else: device_share = torch.device('cuda:' + str(args.gpu_ids[-1])) else: device = device_share = torch.device('cpu') if env is None: env = create_env(env_name, args) if args.train_mode == 0: params = shared_model.player0.parameters() elif args.train_mode == 1: params = shared_model.player1.parameters() else: params = shared_model.parameters() if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(params, lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(params, lr=args.lr) env.seed(args.seed) player = Agent(None, env, args, None, device) player.w_entropy_target = args.entropy_target player.gpu_id = gpu_id # prepare model player.model = build_model(player.env.observation_space, player.env.action_space, args, device) player.model = player.model.to(device) player.model.train() player.reset() reward_sum = torch.zeros(player.num_agents).to(device) reward_sum_org = np.zeros(player.num_agents) ave_reward = np.zeros(2) ave_reward_longterm = np.zeros(2) count_eps = 0 while True: # sys to the shared model player.model.load_state_dict(shared_model.state_dict()) if player.done: player.reset() reward_sum = torch.zeros(player.num_agents).to(device) reward_sum_org = np.zeros(player.num_agents) count_eps += 1 player.update_rnn_hiden() t0 = time.time() for i in range(args.num_steps): player.action_train() reward_sum += player.reward reward_sum_org += player.reward_org if player.done: for j, r_i in enumerate(reward_sum): writer.add_scalar('train/reward_' + str(j), r_i, player.n_steps) break fps = i / (time.time() - t0) # cfg training mode # 0: tracker 1: target -1:joint all training_mode = train_modes[rank] policy_loss, value_loss, entropies, pred_loss = player.optimize( params, optimizer, shared_model, training_mode, device_share) for i in range(min(player.num_agents, 3)): writer.add_scalar('train/policy_loss_' + str(i), policy_loss[i].mean(), player.n_steps) writer.add_scalar('train/value_loss_' + str(i), value_loss[i], player.n_steps) writer.add_scalar('train/entropies' + str(i), entropies[i].mean(), player.n_steps) writer.add_scalar('train/pred_R_loss', pred_loss, player.n_steps) writer.add_scalar('train/ave_reward', ave_reward[0] - ave_reward_longterm[0], player.n_steps) writer.add_scalar('train/mode', training_mode, player.n_steps) writer.add_scalar('train/fps', fps, player.n_steps) n_iter += 1 n_iters[rank] = n_iter if train_modes[rank] == -100: env.close() break
def train(rank, args, shared_model, optimizer, env_conf, shared_counter, targ_shared): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] device = torch.device('cuda:{}'.format(gpu_id) if gpu_id >= 0 else 'cpu') torch.manual_seed(args.seed + rank) torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None, gpu_id=gpu_id) player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space) player.model.apply(weights_init) player.state = player.env.reset() player.state = torch.from_numpy(player.state).to(torch.float32) player.state = player.state.to(device) player.model = player.model.to(device) #player.targ_model = copy.deepcopy(player.model) player.model.train() #player.targ_model.eval() player.eps_len += 2 while True: player.model.load_state_dict(shared_model.state_dict()) #player.targ_model.load_state_dict(targ_shared.state_dict()) if player.done: player.cx = torch.zeros(1, 512).to(device) player.hx = torch.zeros(1, 512).to(device) #player.targ_cx = copy.deepcopy(player.cx).detach() #player.targ_hx = copy.deepcopy(player.hx).detach() else: player.cx = player.cx.detach() player.hx = player.hx.detach() for step in range(args.num_steps): player.action_train() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).to(torch.float32) player.state = player.state.to(device) #alpha = player.model.log_alpha.exp().detach() alpha = .01 #alpha = 0 x_R = torch.zeros(1, 1) if not player.done: with torch.no_grad(): action, value, logit, q_value, _ = player.model( (player.state.unsqueeze(0), (player.hx, player.cx))) x_R = q_value[1].detach() - alpha * F.log_softmax( logit, -1).gather(-1, action) x_R = x_R.to(device) policy_loss = 0 adv_gae_loss = 0 for i in reversed(range(len(player.rewards))): x_R = args.gamma * x_R + player.rewards[i] adv_gae_loss = adv_gae_loss + (player.tra_adv_gae[i][1] - x_R.detach()).pow(2) * .5 #policy_loss = policy_loss - player.log_probs[i] * player.tra_adv_gae[i][0].detach() + alpha * player.log_probs[i] * player.log_probs[i].detach() policy_loss = policy_loss - (F.softmax( player.values[i], -1) * player.tra_adv_gae[i][0].detach()).sum( -1) - alpha * player.entropies[i].unsqueeze(-1) #policy_loss = policy_loss - player.log_probs[i] * (x_R - (F.softmax(player.values[i], -1) * # player.tra_adv_gae[i][0]).sum(-1) - alpha * player.entropies[i]).detach() + alpha * player.log_probs[i] * player.log_probs[i].detach() #prob = F.softmax(player.values[i], -1) #ent_alpha = alpha * player.entropies[i].unsqueeze(-1) #advs = (player.tra_adv_gae[i][0] - # ((player.tra_adv_gae[i][0] * prob).sum(-1, True) + # ent_alpha)).detach() #policy_loss = policy_loss - (prob * advs).sum(-1) - ent_alpha x_R = x_R - alpha * player.log_probs[i].detach() player.model.zero_grad() (policy_loss + .5 * adv_gae_loss).backward(retain_graph=False) ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions() with shared_counter.get_lock(): shared_counter.value += len(player.rewards) if shared_counter.value > args.interact_steps: break
def train(rank, args, shared_model, optimizer, optimizer_r, env_conf, lock, counter): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() player.eps_len += 2 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = [ Variable(torch.zeros(1, 512).cuda()), Variable(torch.zeros(1, 512).cuda()) ] player.hx = [ Variable(torch.zeros(1, 512).cuda()), Variable(torch.zeros(1, 512).cuda()) ] else: player.cx = [ Variable(torch.zeros(1, 512)), Variable(torch.zeros(1, 512)) ] player.hx = [ Variable(torch.zeros(1, 512)), Variable(torch.zeros(1, 512)) ] else: player.cx = [ Variable(player.cx[0].data), Variable(player.cx[1].data) ] player.hx = [ Variable(player.hx[0].data), Variable(player.cx[1].data) ] # 测试rnet的更新有没有影响到这里 # ps = list(player.model.r_net.named_parameters()) # n, v = ps[6] # print(v.sum()) for step in range(args.num_steps): player.action_train() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1, 1) if not player.done: value, _, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx[0], player.cx[0]), (player.hx[1], player.cx[1]))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] with lock: counter.value += 1 # rnet player.model.r_net.zero_grad() (args.actor_weight * policy_loss + (1 - args.actor_weight) * value_loss).backward(retain_graph=True) ensure_shared_grads(player.model.r_net, shared_model.r_net, gpu=gpu_id >= 0) optimizer_r.step() player.model.zero_grad() (policy_loss + 0.5 * value_loss).backward() player.model.r_net.zero_grad() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf, num_tau_samples=32, num_tau_prime_samples=32, kappa=1.0, num_quantiles=32): ptitle('Training Agent: {}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam( shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3Clstm(player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() player.eps_len += 2 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 512).cuda()) player.hx = Variable(torch.zeros(1, 512).cuda()) else: player.cx = Variable(torch.zeros(1, 512)) player.hx = Variable(torch.zeros(1, 512)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train() if player.done: break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() R = torch.zeros(1,num_tau_prime_samples) if not player.done: logit, _, _ = player.model((Variable( player.state.unsqueeze(0)), (player.hx, player.cx))) q_vals = torch.mean(logit,0) _, action = torch.max(q_vals,0) logit, _, _ = player.model((Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = logit[:,action] if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() #R = R.detach() R = Variable(R) value_loss = 0 for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R.repeat(num_tau_samples,1) - player.logits_array[i].repeat(1, num_tau_prime_samples) #print("Ad: ",advantage) loss = (torch.abs(advantage) <= kappa).float() * 0.5 * advantage ** 2 #print("loss: ",loss.sum(0).sum(0), loss) loss += (torch.abs(advantage) > kappa).float() * kappa * (torch.abs(advantage) - 0.5 * kappa) #print("loss: ",loss.sum(0).sum(0), loss) step_loss = torch.abs(player.quantiles_array[i].cuda() - (advantage.detach()<0).float()) * loss/kappa value_loss += step_loss.sum(0).mean(0) player.model.zero_grad() value_loss.backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions()
def train_func(rank, args, shared_model, optimizer, env_conf, datasets=None, shared_dict=None): if args.deploy: return ptitle('Train {0}'.format(rank)) print('Start training agent: ', rank) if rank == 0: logger = Logger(args.log_dir[:-1] + '_losses/') train_step = 0 gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] env_conf["env_gpu"] = gpu_id torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) raw_list, gt_lbl_list = datasets env = EM_env(raw_list, env_conf, type="train", gt_lbl_list=gt_lbl_list, seed=args.seed + rank) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad) player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = get_model(args, args.model, env.observation_space.shape, args.features, atrous_rates=args.atr_rate, num_actions=2, split=args.data_channel, gpu_id=gpu_id, multi=args.multi) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() if rank == 0: eps_reward = 0 pinned_eps_reward = 0 while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: player.eps_len = 0 if rank == 0: if train_step % args.train_log_period == 0 and train_step > 0: print("train: step", train_step, "\teps_reward", eps_reward) if train_step > 0: pinned_eps_reward = player.env.sum_reward.mean() eps_reward = 0 if args.lstm_feats: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx, player.hx = player.model.lstm.init_hidden( batch_size=1, use_cuda=True) else: player.cx, player.hx = player.model.lstm.init_hidden( batch_size=1, use_cuda=False) elif args.lstm_feats: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): if rank < args.lbl_agents: player.action_train(use_lbl=True) else: player.action_train() if rank == 0: eps_reward = player.env.sum_reward.mean() if player.done: break if player.done: state = player.env.reset(player.model, gpu_id) player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() if "3D" in args.data: R = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1], env_conf["size"][2]) else: R = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1]) if args.lowres: R = torch.zeros(1, 1, env_conf["size"][0] // 2, env_conf["size"][1] // 2) if not player.done: if args.lstm_feats: value, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) else: value, _ = player.model(Variable(player.state.unsqueeze(0))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 if "3D" in args.data: gae = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1], env_conf["size"][2]) else: gae = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1]) if args.rew_drop: keep_map = torch.tensor(player.env.keep_map) if args.lowres: gae = torch.zeros(1, 1, env_conf["size"][0] // 2, env_conf["size"][1] // 2) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() if args.rew_drop: keep_map = keep_map.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): if gpu_id >= 0: with torch.cuda.device(gpu_id): reward_i = torch.tensor(player.rewards[i]).cuda() else: reward_i = torch.tensor(player.rewards[i]) R = args.gamma * R + reward_i if args.rew_drop: advantage = R - player.values[i] value_loss = value_loss + (0.5 * advantage * advantage * keep_map).mean() delta_t = player.values[ i + 1].data * args.gamma + reward_i - player.values[i].data gae = gae * args.gamma * args.tau + delta_t else: advantage = R - player.values[i] value_loss = value_loss + (0.5 * advantage * advantage).mean() delta_t = player.values[ i + 1].data * args.gamma + reward_i - player.values[i].data gae = gae * args.gamma * args.tau + delta_t if args.noisy: policy_loss = policy_loss - \ (player.log_probs[i] * Variable(gae)).mean () else: if args.rew_drop: policy_loss = policy_loss - \ (player.log_probs[i] * Variable(gae) * keep_map).mean () - \ (args.entropy_alpha * player.entropies[i] * keep_map).mean () else: policy_loss = policy_loss - \ (player.log_probs[i] * Variable(gae)).mean () - \ (args.entropy_alpha * player.entropies[i]).mean () player.model.zero_grad() sum_loss = (policy_loss + value_loss) curtime = time.time() # print ("backward curtime:", curtime) sum_loss.backward() # print ("backward done", time.time () - curtime) ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) curtime = time.time() # print ("optim curtime:", curtime) optimizer.step() # print ("optim done", time.time () - curtime) player.clear_actions() if args.wctrl == "s2m": player.env.config["spl_w"] = shared_dict["spl_w"] player.env.config["mer_w"] = shared_dict["mer_w"] if rank == 0: train_step += 1 if train_step % args.log_period == 0 and train_step > 0: log_info = { 'train: value_loss': value_loss, 'train: policy_loss': policy_loss, 'train: eps reward': pinned_eps_reward, } if "EX" in args.model: log_info["cell_prob_loss"] = cell_prob_loss for tag, value in log_info.items(): logger.scalar_summary(tag, value, train_step)
def train(rank, args, shared_model, optimizer, env_conf): torch.manual_seed(args.seed + rank) env = atari_env(args.env, env_conf) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr) env.seed(args.seed + rank) player = Agent(None, env, args, None) player.model = A3Clstm( player.env.observation_space.shape[0], player.env.action_space) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() player.model.train() while True: player.model.load_state_dict(shared_model.state_dict()) for step in range(args.num_steps): player.action_train() if args.count_lives: player.check_state() if player.done: break if player.done: player.eps_len = 0 player.current_life = 0 state = player.env.reset() player.state = torch.from_numpy(state).float() R = torch.zeros(1, 1) if not player.done: value, _, _ = player.model( (Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 R = Variable(R) gae = torch.zeros(1, 1) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] optimizer.zero_grad() (policy_loss + 0.5 * value_loss).backward() torch.nn.utils.clip_grad_norm(player.model.parameters(), 40) ensure_shared_grads(player.model, shared_model) optimizer.step() player.clear_actions()
def train (rank, args, shared_model, optimizer, env_conf, datasets=None): ptitle('Training Agent: {}'.format(rank)) print ('Start training agent: ', rank) if rank == 0: logger = Logger (args.log_dir) train_step = 0 gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] env_conf ["env_gpu"] = gpu_id torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) if "EM_env" in args.env: raw, lbl, prob, gt_lbl = datasets env = EM_env (raw, lbl, prob, env_conf, 'train', gt_lbl) else: env = Voronoi_env (env_conf) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop (shared_model.parameters (), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam (shared_model.parameters (), lr=args.lr, amsgrad=args.amsgrad) # env.seed (args.seed + rank) if not args.continuous: player = Agent (None, env, args, None) else: player = Agent_continuous (None, env, args, None) player.gpu_id = gpu_id if not args.continuous: player.model = A3Clstm (env.observation_space.shape, env_conf["num_action"], args.hidden_feat) else: player.model = A3Clstm_continuous (env.observation_space.shape, env_conf["num_action"], args.hidden_feat) player.state = player.env.reset () player.state = torch.from_numpy (player.state).float () old_score = player.env.old_score final_score = 0 if gpu_id >= 0: with torch.cuda.device (gpu_id): player.state = player.state.cuda () player.model = player.model.cuda () player.model.train () if rank == 0: eps_reward = 0 pinned_eps_reward = 0 mean_log_prob = 0 # print ("rank: ", rank) while True: if gpu_id >= 0: with torch.cuda.device (gpu_id): player.model.load_state_dict (shared_model.state_dict ()) else: player.model.load_state_dict (shared_model.state_dict ()) if player.done: player.eps_len = 0 if rank == 0: if 0 <= (train_step % args.train_log_period) < args.max_episode_length: print ("train: step", train_step, "\teps_reward", eps_reward, "\timprovement", final_score - old_score) old_score = player.env.old_score pinned_eps_reward = eps_reward eps_reward = 0 mean_log_prob = 0 if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, args.hidden_feat).cuda()) player.hx = Variable(torch.zeros(1, args.hidden_feat).cuda()) else: player.cx = Variable(torch.zeros(1, args.hidden_feat)) player.hx = Variable(torch.zeros(1, args.hidden_feat)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.action_train () if rank == 0: # if 0 <= (train_step % args.train_log_period) < args.max_episode_length: # print ("train: step", train_step, "\taction = ", player.action) eps_reward += player.reward # print (eps_reward) mean_log_prob += player.log_probs [-1] / env_conf ["T"] if player.done: break if player.done: # if rank == 0: # print ("----------------------------------------------") final_score = player.env.old_score state = player.env.reset () player.state = torch.from_numpy (state).float () if gpu_id >= 0: with torch.cuda.device (gpu_id): player.state = player.state.cuda () R = torch.zeros (1, 1) if not player.done: if not args.continuous: value, _, _ = player.model((Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) else: value, _, _, _ = player.model((Variable(player.state.unsqueeze(0)), (player.hx, player.cx))) R = value.data if gpu_id >= 0: with torch.cuda.device(gpu_id): R = R.cuda() player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 gae = torch.zeros(1, 1) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = gae.cuda() R = Variable(R) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) delta_t = player.values[i + 1].data * args.gamma + player.rewards[i] - \ player.values[i].data gae = gae * args.gamma * args.tau + delta_t # print (player.rewards [i]) if not args.continuous: policy_loss = policy_loss - \ player.log_probs[i] * \ Variable(gae) - 0.01 * player.entropies[i] else: policy_loss = policy_loss - \ player.log_probs[i].sum () * Variable(gae) - \ 0.01 * player.entropies[i].sum () player.model.zero_grad () sum_loss = (policy_loss + value_loss) sum_loss.backward () ensure_shared_grads (player.model, shared_model, gpu=gpu_id >= 0) optimizer.step () player.clear_actions () if rank == 0: train_step += 1 if train_step % args.log_period == 0: log_info = { # 'train: sum_loss': sum_loss, 'train: value_loss': value_loss, 'train: policy_loss': policy_loss, 'train: advanage': advantage, # 'train: entropy': entropy, 'train: eps reward': pinned_eps_reward, # 'train: mean log prob': mean_log_prob } for tag, value in log_info.items (): logger.scalar_summary (tag, value, train_step)
def train(rank, args, shared_model, optimizer): ptitle('TrainingAgent{}'.format(rank)) gpu_id = args.gpu_ids[rank % len(args.gpu_ids)] num_trained_episodes = 0 torch.manual_seed(args.seed + rank) if gpu_id >= 0: torch.cuda.manual_seed(args.seed + rank) env = create_env(args.env, args) if optimizer is None: if args.optimizer == 'RMSprop': optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr) if args.optimizer == 'Adam': optimizer = optim.Adam(shared_model.parameters(), lr=args.lr) env.seed(args.seed + rank) tp_weight = args.tp player = Agent(None, env, args, None) player.gpu_id = gpu_id player.model = A3C_CONV(args.stack_frames, player.env.action_space, args.terminal_prediction, args.reward_prediction) player.state = player.env.reset() player.state = torch.from_numpy(player.state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() player.model = player.model.cuda() player.model.train() while True: if gpu_id >= 0: with torch.cuda.device(gpu_id): player.model.load_state_dict(shared_model.state_dict()) else: player.model.load_state_dict(shared_model.state_dict()) if player.done: #sys._debugmallocstats() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.cx = Variable(torch.zeros(1, 128).cuda()) player.hx = Variable(torch.zeros(1, 128).cuda()) else: player.cx = Variable(torch.zeros(1, 128)) player.hx = Variable(torch.zeros(1, 128)) else: player.cx = Variable(player.cx.data) player.hx = Variable(player.hx.data) for step in range(args.num_steps): player.eps_len += 1 player.action_train() if player.done: num_trained_episodes += 1 break if player.done: state = player.env.reset() player.state = torch.from_numpy(state).float() if gpu_id >= 0: with torch.cuda.device(gpu_id): player.state = player.state.cuda() if gpu_id >= 0: with torch.cuda.device(gpu_id): R = torch.zeros(1, 1).cuda() else: R = torch.zeros(1, 1) if not player.done: state = player.state state = state.unsqueeze(0) value, _, _, _, _, _ = player.model((Variable(state), (player.hx, player.cx))) R = value.data player.values.append(Variable(R)) policy_loss = 0 value_loss = 0 terminal_loss = 0 reward_pred_loss = 0 R = Variable(R) if gpu_id >= 0: with torch.cuda.device(gpu_id): gae = torch.zeros(1, 1).cuda() else: gae = torch.zeros(1, 1) for i in reversed(range(len(player.rewards))): R = args.gamma * R + player.rewards[i] advantage = R - player.values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion # print(player.rewards[i]) delta_t = player.rewards[i] + args.gamma * \ player.values[i + 1].data - player.values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - (player.log_probs[i].sum() * Variable(gae)) - (0.01 * player.entropies[i].sum()) if args.reward_prediction: reward_pred_loss = reward_pred_loss + (player.reward_predictions[i] - player.rewards[i]).pow(2) if args.terminal_prediction: # new way of using emprical episode length as a proxy for current length. if player.average_episode_length is None: end_predict_labels = np.arange(player.eps_len-len(player.terminal_predictions), player.eps_len) / player.eps_len # heuristic else: end_predict_labels = np.arange(player.eps_len-len(player.terminal_predictions), player.eps_len) / player.average_episode_length for i in range(len(player.terminal_predictions)): terminal_loss = terminal_loss + (player.terminal_predictions[i] - end_predict_labels[i]).pow(2) terminal_loss /= len(player.terminal_predictions) player.model.zero_grad() total_loss = policy_loss + 0.5 * value_loss + tp_weight*terminal_loss + 0.5*reward_pred_loss # Visualize Computation Graph #graph = make_dot(total_loss) #from graphviz import Source #Source.view(graph) total_loss.backward() ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0) optimizer.step() player.clear_actions() if player.done: #print(f' CPU {rank} -> train episode count is {num_trained_episodes}') if player.average_episode_length is None: # initial one player.average_episode_length = player.eps_len else: player.average_episode_length = int(0.99 * player.average_episode_length + 0.01 * player.eps_len) #print(player.average_episode_length, 'current one is ', player.eps_len) player.eps_len = 0 # reset here