diff --git a/DPI/train.py b/DPI/train.py index 272099b..cf26d3e 100644 --- a/DPI/train.py +++ b/DPI/train.py @@ -44,39 +44,41 @@ def parse_args(): parser.add_argument('--high_noise', action='store_true') # replay buffer parser.add_argument('--replay_buffer_capacity', default=50000, type=int) #50000 - parser.add_argument('--episode_length', default=51, type=int) + parser.add_argument('--episode_length', default=21, type=int) # 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=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('--batch_size', default=128, type=int) #512 + parser.add_argument('--state_size', default=30, type=int) parser.add_argument('--hidden_size', default=256, type=int) - parser.add_argument('--history_size', default=256, type=int) + parser.add_argument('--history_size', default=128, type=int) + parser.add_argument('--episode_collection', default=5, type=int) + parser.add_argument('--episodes_buffer', default=20, 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) + parser.add_argument('--imagine_horizon', default=10, type=str) parser.add_argument('--grad_clip_norm', type=float, default=100.0, help='Gradient clipping norm') # 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=1e-3, type=float) + parser.add_argument('--value_lr', default=8e-6, 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=1e-3, type=float) + parser.add_argument('--actor_lr', default=8e-6, 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('--world_model_lr', default=1e-3, type=float) - parser.add_argument('--encoder_tau', default=0.001, type=float) + parser.add_argument('--world_model_lr', default=1e-5, type=float) + parser.add_argument('--encoder_tau', default=0.005, type=float) parser.add_argument('--decoder_type', default='pixel', type=str, choices=['pixel', 'identity', 'contrastive', 'reward', 'inverse', 'reconstruction']) parser.add_argument('--num_layers', default=4, type=int) parser.add_argument('--num_filters', default=32, type=int) @@ -117,6 +119,7 @@ class DPI: self.env = make_env(self.args) #self.args.seed = np.random.randint(0, 1000) self.env.seed(self.args.seed) + self.global_episodes = 0 # noiseless environment setup self.args.version = 2 # env_id changes to v2 @@ -160,7 +163,7 @@ class DPI: self.obs_decoder = ObservationDecoder( state_size=self.args.state_size, # 128 - output_shape=(self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84) + output_shape=(self.args.channels*self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84) ).to(device) self.transition_model = TransitionModel( @@ -176,7 +179,9 @@ class DPI: hidden_size=self.args.hidden_size, # 256, action_size=self.env.action_space.shape[0], # 6 ).to(device) - + self.actor_model.apply(self.init_weights) + + # Value Models self.value_model = ValueModel( state_size=self.args.state_size, # 128 @@ -209,23 +214,33 @@ class DPI: self.contrastive_head = ContrastiveHead( hidden_size=self.args.hidden_size, # 256 ).to(device) + + self.club_sample = CLUBSample( + x_dim=self.args.state_size, # 128 + y_dim=self.args.state_size, # 128 + hidden_size=self.args.hidden_size, # 256 + ).to(device) # model 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()) + list(self.reward_model.parameters()) + list(self.club_sample.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.reward_opt = torch.optim.Adam(self.reward_model.parameters(), 1e-5) + #self.decoder_opt = torch.optim.Adam(self.obs_decoder.parameters(), 1e-4) # Create Modules - self.world_model_modules = [self.obs_encoder, self.prjoection_head, self.transition_model, self.obs_decoder, self.reward_model] + self.world_model_modules = [self.obs_encoder, self.prjoection_head, self.transition_model, self.obs_decoder, self.reward_model, self.club_sample] self.value_modules = [self.value_model] self.actor_modules = [self.actor_model] + #self.reward_modules = [self.reward_model] + #self.decoder_modules = [self.obs_decoder] if use_saved: self._use_saved_models(saved_model_dir) @@ -240,79 +255,98 @@ class DPI: done = False 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) - + self.global_episodes += 1 epi_reward = 0 - for i in range(self.args.episode_length): + while not done: if random: action = self.env.action_space.sample() else: 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() - + obs = torch.tensor(obs.copy(), dtype=torch.float32).unsqueeze(0) + obs_processed = preprocess_obs(obs).to(device) + state = self.obs_encoder(obs_processed)["distribution"].sample() + action = self.actor_model(state).cpu().numpy().squeeze() + #action = self.env.action_space.sample() + 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 done: #or i == self.args.episode_length-1: - obs = self.env.reset() - done=False - else: - obs = next_obs + self.data_buffer.add(obs, action, next_obs, rew, done, self.global_episodes) + obs = next_obs epi_reward += rew + + obs = self.env.reset() + done=False 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)) return all_rews def train(self, step, total_steps): counter = 0 + import matplotlib.pyplot as plt + fig, ax = plt.subplots() while step < total_steps: # collect experience if step !=0: encoder = self.obs_encoder actor = self.actor_model - all_rews = self.collect_sequences(10, random=False, actor_model=actor, encoder_model=encoder) + all_rews = self.collect_sequences(self.args.episode_collection, random=False, actor_model=actor, encoder_model=encoder) else: - all_rews = self.collect_sequences(self.args.batch_size, random=True) + all_rews = self.collect_sequences(self.args.episodes_buffer, 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 = 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) + # collect sequences + non_zero_indices = np.nonzero(self.data_buffer.episode_count)[0] + current_obs = self.data_buffer.observations[non_zero_indices] + next_obs = self.data_buffer.next_observations[non_zero_indices] + actions_raw = self.data_buffer.actions[non_zero_indices] + rewards = self.data_buffer.rewards[non_zero_indices] + self.terms = np.where(self.data_buffer.terminals[non_zero_indices]!=0)[0] - 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) - actions = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"actions", device=device, is_obs=False, random_indices=random_indices) - next_actions = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"actions", device=device, is_obs=False, offset=1, random_indices=random_indices) - rewards = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"rewards", device=device, is_obs=False, offset=1, random_indices=random_indices) - - # Preprocessing - last_observations = preprocess_obs(last_observations).to(device) - current_observations = preprocess_obs(current_observations).to(device) - next_observations = preprocess_obs(next_observations).to(device) - - # Initialize transition model states - self.transition_model.init_states(self.args.batch_size, device) # (N,128) - self.history = self.transition_model.prev_history # (N,128) + # Group by episodes + current_obs = self.grouped_arrays(current_obs) + next_obs = self.grouped_arrays(next_obs) + actions_raw = self.grouped_arrays(actions_raw) + rewards_ = self.grouped_arrays(rewards) # Train encoder if step == 0: step += 1 - for _ in range(1):#(self.args.collection_interval // self.args.episode_length+1): + + update_steps = 1 if step > 1 else 1 + #for _ in range(self.args.collection_interval // self.args.episode_length+1): + for _ in range(update_steps): counter += 1 - past_encoder_loss = 0 + + # Select random chunks of episodes + if current_obs.shape[0] < self.args.batch_size: + random_episode_number = np.random.randint(0, current_obs.shape[0], self.args.batch_size) + else: + random_episode_number = random.sample(range(current_obs.shape[0]), self.args.batch_size) + + if current_obs[0].shape[0]-self.args.episode_length < self.args.batch_size: + init_index = np.random.randint(0, current_obs[0].shape[0]-self.args.episode_length-2, self.args.batch_size) + else: + init_index = np.asarray(random.sample(range(current_obs[0].shape[0]-self.args.episode_length), self.args.batch_size)) + random.shuffle(random_episode_number) + random.shuffle(init_index) + last_observations = self.select_first_k(current_obs, init_index, random_episode_number)[:-1] + current_observations = self.select_first_k(current_obs, init_index, random_episode_number)[1:] + next_observations = self.select_first_k(next_obs, init_index, random_episode_number)[:-1] + actions = self.select_first_k(actions_raw, init_index, random_episode_number)[:-1].to(device) + next_actions = self.select_first_k(actions_raw, init_index, random_episode_number)[1:].to(device) + rewards = self.select_first_k(rewards_, init_index, random_episode_number)[1:].to(device) + + # Preprocessing + last_observations = preprocess_obs(last_observations).to(device) + current_observations = preprocess_obs(current_observations).to(device) + next_observations = preprocess_obs(next_observations).to(device) + + # Initialize transition model states + self.transition_model.init_states(self.args.batch_size, device) # (N,128) + self.history = self.transition_model.prev_history # (N,128) + + past_world_model_loss = 0 + past_action_loss = 0 + past_value_loss = 0 for i in range(self.args.episode_length-1): if i > 0: # Encode observations and next_observations @@ -323,11 +357,11 @@ class DPI: self.next_action = next_actions[i] # (N,6) history = self.transition_model.prev_history - # Encode negative observations + # Encode negative observations fro upper bound loss idx = torch.randperm(current_observations[i].shape[0]) # random permutation on batch - random_time_index = torch.randint(0, self.args.episode_length-2, (1,)).item() # random time index + random_time_index = torch.randint(0, current_observations.shape[0]-2, (1,)).item() # random time index negative_current_observations = current_observations[random_time_index][idx] - self.negative_current_states_dict = self.obs_encoder(negative_current_observations) + self.negative_current_states_dict = self.get_features(negative_current_observations) # Predict current state from past state with transition model last_states_sample = self.last_states_dict["sample"] @@ -335,25 +369,24 @@ class DPI: self.history = predicted_current_state_dict["history"] # Calculate upper bound loss - likeli_loss, ub_loss = self._upper_bound_minimization(self.last_states_dict, - self.current_states_dict, - self.negative_current_states_dict, - predicted_current_state_dict - ) + likeli_loss, ub_loss = self._upper_bound_minimization(self.last_states_dict["sample"].detach(), + self.current_states_dict["sample"].detach(), + self.negative_current_states_dict["sample"].detach(), + predicted_current_state_dict["sample"].detach(), + ) # Calculate encoder loss - encoder_loss = past_encoder_loss + self._past_encoder_loss(self.current_states_dict, predicted_current_state_dict) - past_encoder_loss = encoder_loss.item() + encoder_loss = self._past_encoder_loss(self.current_states_dict, predicted_current_state_dict) # 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) + nxt_obs = next_observations[i:i+horizon].reshape(-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,:,:])) - + decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon])) + # contrastive projection - vec_anchor = predicted_current_state_dict["sample"] + vec_anchor = predicted_current_state_dict["sample"].detach() vec_positive = self.next_states_dict["sample"].detach() z_anchor = self.prjoection_head(vec_anchor, self.action) z_positive = self.prjoection_head_momentum(vec_positive, next_actions[i]).detach() @@ -363,24 +396,25 @@ class DPI: labels = torch.arange(logits.shape[0]).long().to(device) lb_loss = F.cross_entropy(logits, labels) - # reward loss + # 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[i])) # world model loss - world_model_loss = encoder_loss + ub_loss + lb_loss + reward_loss + decoder_loss - + world_model_loss = (10*encoder_loss + 10*ub_loss + 1e-1*lb_loss + reward_loss + 1e-3*decoder_loss + past_world_model_loss) * 1e-3 + past_world_model_loss = world_model_loss.item() + # 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) - + imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"], + action, self.history, + 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 + imag_values = self.target_value_model(imagined_rollout["sample"]).mean discounts = self.args.discount * torch.ones_like(imag_rewards).detach() @@ -392,44 +426,26 @@ class DPI: 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.returns) + actor_loss = -torch.mean(self.discounts * self.returns) + past_action_loss + past_action_loss = actor_loss.item() # value loss with torch.no_grad(): value_feat = imagined_rollout["sample"][:-1].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 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) + value_dist = self.value_model(value_feat) - 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) + value_loss = -torch.mean(self.discounts * value_dist.log_prob(value_targ).unsqueeze(-1)) + past_value_loss + past_value_loss = value_loss.item() # 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)) - + #count = np.arange((counter-1) * (self.args.batch_size), (counter) * (self.args.batch_size)) + count = (counter-1) * (self.args.batch_size) if step % self.args.logging_freq: 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) @@ -438,27 +454,44 @@ class DPI: 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) - + writer.add_scalar('Bound Loss/Lower Bound Loss', lb_loss.detach().item(), self.data_buffer.steps) step += 1 - - - # 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() + print(world_model_loss, actor_loss, value_loss) + # update actor model + 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() + + # update world model + 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 value model + self.value_opt.zero_grad() + value_loss.backward() + nn.utils.clip_grad_norm_(self.value_model.parameters(), self.args.grad_clip_norm) + 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) + + rew_len = np.arange(count, count+self.args.episode_collection) if count != 0 else np.arange(0, self.args.batch_size) + for j in range(len(all_rews)): + writer.add_scalar('Rewards/Rewards', all_rews[j], rew_len[j]) + + print(step) + if step % 2850 == 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, eval_episodes=10): path = path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" @@ -476,9 +509,10 @@ class DPI: done = False while not done: 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() + obs = torch.tensor(obs.copy(), dtype=torch.float32).unsqueeze(0) + obs_processed = preprocess_obs(obs).to(device) + state = self.obs_encoder(obs_processed)["distribution"].sample() + action = self.actor_model(state).cpu().detach().numpy().squeeze() next_obs, rew, done, _ = self.env.step(action) rewards += rew @@ -491,16 +525,51 @@ class DPI: episodic_rewards.append(rewards) print("Episodic rewards: ", episodic_rewards) print("Average episodic reward: ", np.mean(episodic_rewards)) - + + def init_weights(self, m): + if isinstance(m, nn.Linear): + torch.nn.init.xavier_uniform_(m.weight) + m.bias.data.fill_(0.01) + def grouped_arrays(self,array): + indices = [0] + self.terms.tolist() + + def subarrays(): + for start, end in zip(indices[:-1], indices[1:]): + yield array[start:end] + + try: + subarrays = np.stack(list(subarrays()), axis=0) + except ValueError: + subarrays = np.asarray(list(subarrays())) + + return subarrays + + def select_first_k(self, array, init_index, episode_number): + term_index = init_index + self.args.episode_length + array = array[episode_number] + + array_list = [] + for i in range(array.shape[0]): + array_list.append(array[i][init_index[i]:term_index[i]]) + array = np.asarray(array_list) + + if array.ndim == 5: + transposed_array = np.transpose(array, (1, 0, 2, 3, 4)) + elif array.ndim == 4: + transposed_array = np.transpose(array, (1, 0, 2, 3)) + elif array.ndim == 3: + transposed_array = np.transpose(array, (1, 0, 2)) + elif array.ndim == 2: + transposed_array = np.transpose(array, (1, 0)) + else: + transposed_array = np.expand_dims(array, axis=0) + return torch.tensor(transposed_array).float() + def _upper_bound_minimization(self, last_states, current_states, negative_current_states, predicted_current_states): - club_sample = CLUBSample(last_states, - current_states, - negative_current_states, - predicted_current_states) - likelihood_loss = club_sample.learning_loss() - club_loss = club_sample() + club_loss = self.club_sample(current_states, predicted_current_states, negative_current_states) + likelihood_loss = 0 return likelihood_loss, club_loss def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict): @@ -511,7 +580,7 @@ class DPI: predicted_curr_states_dist = predicted_curr_states_dict["distribution"] # KL divergence loss - loss = torch.mean(torch.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist)) + loss = torch.mean(torch.distributions.kl.kl_divergence(curr_states_dist,predicted_curr_states_dist)) return loss @@ -573,7 +642,7 @@ if __name__ == '__main__': device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') step = 0 - total_steps = 200000 + total_steps = 500000 dpi = DPI(args) dpi.train(step,total_steps) dpi.evaluate() \ No newline at end of file