def train(rank, scene_scope, task_scope, args, shared_model, counter, lock, optimizer=None): torch.manual_seed(args.seed + rank) #env = create_atari_env(args.env_name) #env.seed(args.seed + rank) env = Environment({ 'scene_name': scene_scope, 'terminal_state_id': int(task_scope) }) model = ActorCriticFFNetwork(ACTION_SIZE) if optimizer is None: # TODO: Discount learning rate based on episode length optimizer = my_optim.SharedRMSprop(shared_model.parameters(), lr=args.lr, alpha=args.alpha, eps=args.eps) optimizer.share_memory() model.train() env.reset() state = torch.from_numpy(env.s_t) done = True episode_length = 0 for i in range(int(args.max_episode_length)): # Sync with the shared model model.load_state_dict(shared_model.state_dict()) ''' if done: cx = Variable(torch.zeros(1, 256)) hx = Variable(torch.zeros(1, 256)) else: cx = Variable(cx.data) hx = Variable(hx.data) ''' values = [] log_probs = [] rewards = [] entropies = [] for step in range(args.num_steps): print('Thread: ', rank, ', step: ', step, 'epochs:', i) episode_length += 1 logit, value = model(env.s_t, env.target) prob = F.softmax(logit, dim=1) log_prob = F.log_softmax(logit, dim=1) entropy = -(log_prob * prob).sum(1, keepdim=True) entropies.append(entropy) action = prob.multinomial(num_samples=1).data log_prob = log_prob.gather(1, Variable(action)) env.step(action) #state, reward, done, _ = env.step(action.numpy()) env.update() state = env.s_t reward = env.reward done = env.terminal done = done or episode_length >= args.max_episode_length reward = max(min(reward, 1), -1) with lock: if counter.value % 1000 == 0: print('Now saving data. Please wait.') torch.save(shared_model.state_dict(), CHECKPOINT_DIR + '/' + 'checkpoint.pth.tar') counter.value += 1 if done: episode_length = 0 if env.terminal: print('Task completed') counter.value += 1 if done: episode_length = 0 env.reset() state = env.s_t state = torch.from_numpy(state) values.append(value) log_probs.append(log_prob) rewards.append(reward) if done: break R = torch.zeros(1, 1) if not done: _, value = model(env.s_t, env.target) R = value.data values.append(Variable(R)) policy_loss = 0 value_loss = 0 R = Variable(R) gae = torch.zeros(1, 1) for i in reversed(range(len(rewards))): R = args.gamma * R + rewards[i] advantage = R - values[i] value_loss = value_loss + 0.5 * advantage.pow(2) # Generalized Advantage Estimataion delta_t = rewards[i] + args.gamma * \ values[i + 1].data - values[i].data gae = gae * args.gamma * args.tau + delta_t policy_loss = policy_loss - \ log_probs[i] * Variable(gae) - args.entropy_coef * entropies[i] optimizer.zero_grad() (policy_loss + args.value_loss_coef * value_loss).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) ensure_shared_grads(model, shared_model) optimizer.step()
def test(rank, scene_scope, task_scope, args, shared_model, counter): torch.manual_seed(args.seed + rank) env = Environment({ 'scene_name': scene_scope, 'terminal_state_id': int(task_scope) }) model = ActorCriticFFNetwork(ACTION_SIZE) model.eval() height, width, layers = env.observation.shape video = cv2.VideoWriter('video/' + task_scope + '.mp4',-1,1,(width,height)) env.reset() state = torch.from_numpy(env.s_t) reward_sum = 0 done = True start_time = time.time() # a quick hack to prevent the agent from stucking actions = deque(maxlen=100) episode_length = 0 img = cv2.cvtColor(env.observation, cv2.COLOR_BGR2RGB) video.write(img) for i in range(100): episode_length += 1 # Sync with the shared model if done: model.load_state_dict(shared_model.state_dict()) logit, value = model(env.s_t, env.target) prob = F.softmax(logit, dim=1) action = prob.max(1, keepdim=True)[1].data.numpy() env.step(action[0, 0]) env.update() img = cv2.cvtColor(env.observation, cv2.COLOR_BGR2RGB) video.write(img) reward = env.reward state = env.s_t done = env.terminal print(env.terminal_state_id, env.current_state_id) done = done or episode_length >= args.max_episode_length reward_sum += reward # a quick hack to prevent the agent from stucking actions.append(action[0, 0]) if actions.count(actions[0]) == actions.maxlen: done = True if done: print("Time {}, num steps {}, FPS {:.0f}, episode reward {}, episode length {}".format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - start_time)), counter.value, counter.value / (time.time() - start_time), reward_sum, episode_length)) reward_sum = 0 episode_length = 0 actions.clear() env.reset() state = env.s_t break state = torch.from_numpy(state) cv2.destroyAllWindows() video.release()
''' p = mp.Process(target=test, args=(args.num_processes, args, shared_model, counter)) p.start() processes.append(p) ''' branches = [] for scene in scene_scopes: for task in list_of_tasks[scene]: branches.append((scene, task)) NUM_TASKS = len(branches) if os.path.exists(CHECKPOINT_DIR + '/' + 'checkpoint.pth.tar'): checkpoint = torch.load(CHECKPOINT_DIR + '/' + 'checkpoint.pth.tar', map_location=lambda storage, loc: storage) # set global step shared_model.load_state_dict(checkpoint) print("Model loaded") else: print("Could not find old checkpoint") for rank in range(0, args.num_processes): scene, task = branches[rank % NUM_TASKS] p = mp.Process(target=test, args=(rank, scene, task, args, shared_model, counter)) #p = mp.Process(target=train, args=(rank ,scene, task, args, shared_model, counter, lock, optimizer)) p.start() processes.append(p) for p in processes: p.join() '''