diff --git a/DPI/train.py b/DPI/train.py index 4845d91..272099b 100644 --- a/DPI/train.py +++ b/DPI/train.py @@ -48,11 +48,11 @@ def parse_args(): # train parser.add_argument('--agent', default='dpi', type=str, choices=['baseline', 'bisim', 'deepmdp', 'db', 'dpi', 'rpc']) parser.add_argument('--init_steps', default=10000, type=int) - parser.add_argument('--num_train_steps', default=10000, type=int) - parser.add_argument('--batch_size', default=30, type=int) #512 - parser.add_argument('--state_size', default=256, type=int) - parser.add_argument('--hidden_size', default=128, type=int) - parser.add_argument('--history_size', default=128, type=int) + parser.add_argument('--num_train_steps', default=100000, type=int) + parser.add_argument('--batch_size', default=50, type=int) #512 + parser.add_argument('--state_size', default=512, type=int) + parser.add_argument('--hidden_size', default=256, type=int) + parser.add_argument('--history_size', default=256, type=int) parser.add_argument('--num-units', type=int, default=50, help='num hidden units for reward/value/discount models') parser.add_argument('--load_encoder', default=None, type=str) parser.add_argument('--imagine_horizon', default=15, type=str) @@ -60,42 +60,33 @@ def parse_args(): # eval parser.add_argument('--eval_freq', default=10, type=int) # TODO: master had 10000 parser.add_argument('--num_eval_episodes', default=20, type=int) + parser.add_argument('--evaluation_interval', default=10000, type=int) # TODO: master had 10000 # value - parser.add_argument('--value_lr', default=8e-5, type=float) + parser.add_argument('--value_lr', default=1e-3, type=float) parser.add_argument('--value_beta', default=0.9, type=float) parser.add_argument('--value_tau', default=0.005, type=float) parser.add_argument('--value_target_update_freq', default=100, type=int) parser.add_argument('--td_lambda', default=0.95, type=int) # actor - parser.add_argument('--actor_lr', default=8e-5, type=float) + parser.add_argument('--actor_lr', default=1e-3, type=float) parser.add_argument('--actor_beta', default=0.9, type=float) parser.add_argument('--actor_log_std_min', default=-10, type=float) parser.add_argument('--actor_log_std_max', default=2, type=float) parser.add_argument('--actor_update_freq', default=2, type=int) # world/encoder/decoder parser.add_argument('--encoder_type', default='pixel', type=str, choices=['pixel', 'pixelCarla096', 'pixelCarla098', 'identity']) - parser.add_argument('--encoder_feature_dim', default=50, type=int) - parser.add_argument('--world_model_lr', default=6e-4, type=float) - parser.add_argument('--past_transition_lr', default=1e-3, type=float) - parser.add_argument('--encoder_lr', default=1e-3, type=float) + parser.add_argument('--world_model_lr', default=1e-3, type=float) parser.add_argument('--encoder_tau', default=0.001, type=float) - parser.add_argument('--encoder_stride', default=1, type=int) parser.add_argument('--decoder_type', default='pixel', type=str, choices=['pixel', 'identity', 'contrastive', 'reward', 'inverse', 'reconstruction']) - parser.add_argument('--decoder_lr', default=1e-3, type=float) - parser.add_argument('--decoder_update_freq', default=1, type=int) - parser.add_argument('--decoder_weight_lambda', default=0.0, type=float) parser.add_argument('--num_layers', default=4, type=int) parser.add_argument('--num_filters', default=32, type=int) parser.add_argument('--aug', action='store_true') # sac parser.add_argument('--discount', default=0.99, type=float) - parser.add_argument('--init_temperature', default=0.01, type=float) - parser.add_argument('--alpha_lr', default=1e-3, type=float) - parser.add_argument('--alpha_beta', default=0.9, type=float) # misc parser.add_argument('--seed', default=1, type=int) parser.add_argument('--logging_freq', default=100, type=int) - parser.add_argument('--saving_interval', default=1000, type=int) + parser.add_argument('--saving_interval', default=2500, type=int) parser.add_argument('--work_dir', default='.', type=str) parser.add_argument('--save_tb', default=False, action='store_true') parser.add_argument('--save_model', default=False, action='store_true') @@ -221,19 +212,18 @@ class DPI: # model parameters - self.world_model_parameters = list(self.obs_encoder.parameters()) + list(self.obs_decoder.parameters()) + \ - list(self.value_model.parameters()) + list(self.transition_model.parameters()) + \ - list(self.prjoection_head.parameters()) + self.world_model_parameters = list(self.obs_encoder.parameters()) + list(self.prjoection_head.parameters()) + \ + list(self.transition_model.parameters()) + list(self.obs_decoder.parameters()) + \ + list(self.reward_model.parameters()) self.past_transition_parameters = self.transition_model.parameters() # optimizers self.world_model_opt = torch.optim.Adam(self.world_model_parameters, self.args.world_model_lr) self.value_opt = torch.optim.Adam(self.value_model.parameters(), self.args.value_lr) self.actor_opt = torch.optim.Adam(self.actor_model.parameters(), self.args.actor_lr) - self.past_transition_opt = torch.optim.Adam(self.past_transition_parameters, self.args.past_transition_lr) # Create Modules - self.world_model_modules = [self.obs_encoder, self.obs_decoder, self.reward_model, self.transition_model, self.prjoection_head] + self.world_model_modules = [self.obs_encoder, self.prjoection_head, self.transition_model, self.obs_decoder, self.reward_model] self.value_modules = [self.value_model] self.actor_modules = [self.actor_model] @@ -252,8 +242,8 @@ class DPI: all_rews = [] #video = VideoRecorder(self.video_dir if args.save_video else None, resource_files=args.resource_files) for episode_count in tqdm.tqdm(range(episodes), desc='Collecting episodes'): - if args.save_video: - self.env.video.init(enabled=True) + #if args.save_video: + # self.env.video.init(enabled=True) epi_reward = 0 for i in range(self.args.episode_length): @@ -263,24 +253,24 @@ class DPI: with torch.no_grad(): obs_torch = torch.unsqueeze(torch.tensor(obs).float(),0).to(device) state = self.obs_encoder(obs_torch)["distribution"].sample() - action = self.actor_model(state).cpu().detach().numpy().squeeze() + action = self.actor_model(state).cpu().detach().numpy().squeeze() next_obs, rew, done, _ = self.env.step(action) self.data_buffer.add(obs, action, next_obs, rew, episode_count+1, done) - if args.save_video: - self.env.video.record(self.env) + #if args.save_video: + # self.env.video.record(self.env) - if done or i == self.args.episode_length-1: + if done: #or i == self.args.episode_length-1: obs = self.env.reset() done=False else: obs = next_obs epi_reward += rew all_rews.append(epi_reward) - if args.save_video: - self.env.video.save('noisy/%d.mp4' % episode_count) - print("Collected {} random episodes".format(episode_count+1)) + #if args.save_video: + # self.env.video.save('noisy/%d.mp4' % episode_count) + #print("Collected {} random episodes".format(episode_count+1)) return all_rews def train(self, step, total_steps): @@ -291,14 +281,16 @@ class DPI: if step !=0: encoder = self.obs_encoder actor = self.actor_model - #all_rews = self.collect_sequences(self.args.batch_size, random=True) - all_rews = self.collect_sequences(self.args.batch_size, random=False, actor_model=actor, encoder_model=encoder) + all_rews = self.collect_sequences(10, random=False, actor_model=actor, encoder_model=encoder) else: all_rews = self.collect_sequences(self.args.batch_size, random=True) # Group by steps and sample random batch - random_indices = self.data_buffer.sample_random_idx(self.args.batch_size * ((step//self.args.collection_interval)+1)) # random indices for batch - #random_indices = np.arange(self.args.batch_size * ((step//self.args.collection_interval)),self.args.batch_size * ((step//self.args.collection_interval)+1)) + #random_indices = self.data_buffer.sample_random_idx(self.args.batch_size * ((step//self.args.collection_interval)+1)) # random indices for batch + #random_indices = self.data_buffer.sample_random_idx(self.data_buffer.steps//self.args.episode_length) + final_idx = self.data_buffer.group_steps(self.data_buffer, "observations").shape[1] + random_indices = self.data_buffer.sample_random_idx(final_idx, last=True) + last_observations = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"observations", "cpu", random_indices=random_indices) current_observations = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"next_observations", device="cpu", random_indices=random_indices) next_observations = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"next_observations", device="cpu", offset=1, random_indices=random_indices) @@ -318,8 +310,9 @@ class DPI: # Train encoder if step == 0: step += 1 - for _ in range(self.args.collection_interval // self.args.episode_length+1): + for _ in range(1):#(self.args.collection_interval // self.args.episode_length+1): counter += 1 + past_encoder_loss = 0 for i in range(self.args.episode_length-1): if i > 0: # Encode observations and next_observations @@ -349,8 +342,15 @@ class DPI: ) # Calculate encoder loss - encoder_loss = self._past_encoder_loss(self.current_states_dict, - predicted_current_state_dict) + encoder_loss = past_encoder_loss + self._past_encoder_loss(self.current_states_dict, predicted_current_state_dict) + past_encoder_loss = encoder_loss.item() + + # decoder loss + horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) + nxt_obs = next_observations[i:i+horizon].view(-1,9,84,84) + next_states_encodings = self.get_features(nxt_obs)["sample"].view(horizon,self.args.batch_size, -1) + obs_dist = self.obs_decoder(next_states_encodings) + decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon][:,:,:3,:,:])) # contrastive projection vec_anchor = predicted_current_state_dict["sample"] @@ -362,136 +362,137 @@ class DPI: logits = self.contrastive_head(z_anchor, z_positive) labels = torch.arange(logits.shape[0]).long().to(device) lb_loss = F.cross_entropy(logits, labels) - - # behaviour learning - with FreezeParameters(self.world_model_modules): - imagine_horizon = self.args.imagine_horizon #np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) - imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"].detach(), - self.next_action, self.history.detach(), - imagine_horizon) - - # decoder loss - horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) - obs_dist = self.obs_decoder(imagined_rollout["sample"][:horizon]) - decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon][:,:,:3,:,:])) # reward loss reward_dist = self.reward_model(self.current_states_dict["sample"]) - reward_loss = -torch.mean(reward_dist.log_prob(rewards[:-1])) + reward_loss = -torch.mean(reward_dist.log_prob(rewards[:-1])) - # update models - world_model_loss = encoder_loss + 100 * ub_loss + lb_loss + reward_loss + decoder_loss * 1e-2 - self.world_model_opt.zero_grad() - world_model_loss.backward() - nn.utils.clip_grad_norm_(self.world_model_parameters, self.args.grad_clip_norm) - self.world_model_opt.step() - - # update momentum encoder - soft_update_params(self.obs_encoder, self.obs_encoder_momentum, self.args.encoder_tau) - - # update momentum projection head - soft_update_params(self.prjoection_head, self.prjoection_head_momentum, self.args.encoder_tau) - - # actor loss - with FreezeParameters(self.world_model_modules + self.value_modules): - imag_rew_dist = self.reward_model(imagined_rollout["sample"]) - target_imag_val_dist = self.target_value_model(imagined_rollout["sample"]) - - imag_rews = imag_rew_dist.mean - target_imag_vals = target_imag_val_dist.mean - - discounts = self.args.discount * torch.ones_like(imag_rews).detach() + # world model loss + world_model_loss = encoder_loss + ub_loss + lb_loss + reward_loss + decoder_loss - self.target_returns = self._compute_lambda_return(imag_rews[:-1], - target_imag_vals[:-1], - discounts[:-1] , - self.args.td_lambda, - target_imag_vals[-1]) + # actor loss + with FreezeParameters(self.world_model_modules): + imagine_horizon = self.args.imagine_horizon #np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) + action = self.actor_model(self.current_states_dict["sample"]) + imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"].detach(), + action, self.history.detach(), + imagine_horizon) + + with FreezeParameters(self.world_model_modules + self.value_modules): + imag_rewards = self.reward_model(imagined_rollout["sample"]).mean + imag_values = self.value_model(imagined_rollout["sample"]).mean + + discounts = self.args.discount * torch.ones_like(imag_rewards).detach() + + self.returns = self._compute_lambda_return(imag_rewards[:-1], + imag_values[:-1], + discounts[:-1] , + self.args.td_lambda, + imag_values[-1]) discounts = torch.cat([torch.ones_like(discounts[:1]), discounts[1:-1]], 0) self.discounts = torch.cumprod(discounts, 0).detach() - actor_loss = -torch.mean(self.discounts * self.target_returns) + actor_loss = -torch.mean(self.discounts * self.returns) - # update actor - self.actor_opt.zero_grad() - actor_loss.backward() - nn.utils.clip_grad_norm_(self.actor_model.parameters(), self.args.grad_clip_norm) - self.actor_opt.step() - # value loss with torch.no_grad(): value_feat = imagined_rollout["sample"][:-1].detach() - value_targ = self.target_returns.detach() + value_targ = self.returns.detach() value_dist = self.value_model(value_feat) value_loss = -torch.mean(self.discounts * value_dist.log_prob(value_targ).unsqueeze(-1)) - - # update value + + # update models + self.world_model_opt.zero_grad() + self.actor_opt.zero_grad() self.value_opt.zero_grad() + + world_model_loss.backward() + actor_loss.backward() value_loss.backward() + + nn.utils.clip_grad_norm_(self.world_model_parameters, self.args.grad_clip_norm) + nn.utils.clip_grad_norm_(self.actor_model.parameters(), self.args.grad_clip_norm) nn.utils.clip_grad_norm_(self.value_model.parameters(), self.args.grad_clip_norm) + + self.world_model_opt.step() + self.actor_opt.step() self.value_opt.step() + # update momentum encoder and projection head + soft_update_params(self.obs_encoder, self.obs_encoder_momentum, self.args.encoder_tau) + soft_update_params(self.prjoection_head, self.prjoection_head_momentum, self.args.encoder_tau) + # update target value - if step % self.args.value_target_update_freq == 0: - self.target_value_model = copy.deepcopy(self.value_model) + #if step % self.args.value_target_update_freq == 0: + # self.target_value_model = copy.deepcopy(self.value_model) # counter for reward count = np.arange((counter-1) * (self.args.batch_size), (counter) * (self.args.batch_size)) - if step % self.args.logging_freq: - writer.add_scalar('World Loss/World Loss', world_model_loss.detach().item(), step) - writer.add_scalar('Main Models Loss/Encoder Loss', encoder_loss.detach().item(), step) - writer.add_scalar('Main Models Loss/Decoder Loss', decoder_loss, step) - writer.add_scalar('Actor Critic Loss/Actor Loss', actor_loss.detach().item(), step) - writer.add_scalar('Actor Critic Loss/Value Loss', value_loss.detach().item(), step) - writer.add_scalar('Actor Critic Loss/Reward Loss', reward_loss.detach().item(), step) - writer.add_scalar('Bound Loss/Upper Bound Loss', ub_loss.detach().item(), step) - writer.add_scalar('Bound Loss/Lower Bound Loss', lb_loss.detach().item(), step) + writer.add_scalar('World Loss/World Loss', world_model_loss.detach().item(), self.data_buffer.steps) + writer.add_scalar('Main Models Loss/Encoder Loss', encoder_loss.detach().item(), self.data_buffer.steps) + writer.add_scalar('Main Models Loss/Decoder Loss', decoder_loss, self.data_buffer.steps) + writer.add_scalar('Actor Critic Loss/Actor Loss', actor_loss.detach().item(), self.data_buffer.steps) + writer.add_scalar('Actor Critic Loss/Value Loss', value_loss.detach().item(), self.data_buffer.steps) + writer.add_scalar('Actor Critic Loss/Reward Loss', reward_loss.detach().item(), self.data_buffer.steps) + writer.add_scalar('Bound Loss/Upper Bound Loss', ub_loss.detach().item(), self.data_buffer.steps) + writer.add_scalar('Bound Loss/Lower Bound Loss', lb_loss.detach().item(), self.data_buffer.steps) step += 1 - if step>total_steps: - print("Training finished") - break - - # save model - if step % self.args.saving_interval == 0: - path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" - self.save_models(path) - #torch.cuda.empty_cache() # memory leak issues + + # save model + #if step % 500 == 0:#self.args.saving_interval == 0: + # print("Saving model") + # path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" + # self.save_models(path) for j in range(len(all_rews)): writer.add_scalar('Rewards/Rewards', all_rews[j], count[j]) + + #print(self.data_buffer.steps , ((self.args.episode_length-1) * self.args.batch_size * 5)) + if self.data_buffer.steps % 5100 == 0 and self.data_buffer.steps!=0: #self.args.evaluation_interval == 0: + print("Saving model") + path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" + self.save_models(path) + self.evaluate() - def evaluate(self, env, eval_episodes, render=False): + def evaluate(self, eval_episodes=10): + path = path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" + self.restore_checkpoint(path) - episode_rew = np.zeros((eval_episodes)) + obs = self.env.reset() + done = False - video_images = [[] for _ in range(eval_episodes)] - - for i in range(eval_episodes): - obs = env.reset() + #video = VideoRecorder(self.video_dir, resource_files=self.args.resource_files) + if self.args.save_video: + self.env.video.init(enabled=True) + episodic_rewards = [] + for episode in range(eval_episodes): + rewards = 0 done = False - prev_state = self.rssm.init_state(1, self.device) - prev_action = torch.zeros(1, self.action_size).to(self.device) - while not done: with torch.no_grad(): - posterior, action = self.act_with_world_model(obs, prev_state, prev_action) - action = action[0].cpu().numpy() - next_obs, rew, done, _ = env.step(action) - prev_state = posterior - prev_action = torch.tensor(action, dtype=torch.float32).to(self.device).unsqueeze(0) + obs_torch = torch.unsqueeze(torch.tensor(obs).float(),0).to(device) + state = self.obs_encoder(obs_torch)["distribution"].sample() + action = self.actor_model(state).cpu().detach().numpy().squeeze() + + next_obs, rew, done, _ = self.env.step(action) + rewards += rew - episode_rew[i] += rew - - if render: - video_images[i].append(obs['image'].transpose(1,2,0).copy()) + if self.args.save_video: + self.env.video.record(self.env) + self.env.video.save('/home/vedant/Curiosity/Curiosity/DPI/log/video/learned_model.mp4') obs = next_obs - return episode_rew, np.array(video_images[:self.args.max_videos_to_save]) + obs = self.env.reset() + episodic_rewards.append(rewards) + print("Episodic rewards: ", episodic_rewards) + print("Average episodic reward: ", np.mean(episodic_rewards)) + + def _upper_bound_minimization(self, last_states, current_states, negative_current_states, predicted_current_states): club_sample = CLUBSample(last_states, @@ -510,15 +511,16 @@ class DPI: predicted_curr_states_dist = predicted_curr_states_dict["distribution"] # KL divergence loss - loss = torch.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist).mean() + loss = torch.mean(torch.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist)) return loss def get_features(self, x, momentum=False): if self.args.aug: - x = T.RandomCrop((80, 80))(x) # (None,80,80,4) - x = T.functional.pad(x, (4, 4, 4, 4), "symmetric") # (None,88,88,4) - x = T.RandomCrop((84, 84))(x) # (None,84,84,4) + crop_transform = T.RandomCrop(size=80) + cropped_x = torch.stack([crop_transform(x[i]) for i in range(x.size(0))]) + padding = (2, 2, 2, 2) + x = F.pad(cropped_x, padding) with torch.no_grad(): if momentum: @@ -550,6 +552,17 @@ class DPI: 'actor_optimizer': self.actor_opt.state_dict(), 'value_optimizer': self.value_opt.state_dict(), 'world_model_optimizer': self.world_model_opt.state_dict(),}, save_path) + + def restore_checkpoint(self, ckpt_path): + checkpoint = torch.load(ckpt_path) + self.transition_model.load_state_dict(checkpoint['rssm']) + self.actor_model.load_state_dict(checkpoint['actor']) + self.reward_model.load_state_dict(checkpoint['reward_model']) + self.obs_encoder.load_state_dict(checkpoint['obs_encoder']) + self.obs_decoder.load_state_dict(checkpoint['obs_decoder']) + self.world_model_opt.load_state_dict(checkpoint['world_model_optimizer']) + self.actor_opt.load_state_dict(checkpoint['actor_optimizer']) + self.value_opt.load_state_dict(checkpoint['value_optimizer']) if __name__ == '__main__': args = parse_args() @@ -560,6 +573,7 @@ if __name__ == '__main__': device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') step = 0 - total_steps = 10000 + total_steps = 200000 dpi = DPI(args) - dpi.train(step,total_steps) \ No newline at end of file + dpi.train(step,total_steps) + dpi.evaluate() \ No newline at end of file