DPI/models.py

@@ -39,7 +39,6 @@ class ObservationEncoder(nn.Module):
         dist = self.get_dist(mean, std)
 
         # Sampling via reparameterization Trick
-        #x = dist.rsample()
         x = self.reparameterize(mean, std)
 
         encoded_output = {"sample": x, "distribution": dist}
@@ -166,9 +165,9 @@ class RewardModel(nn.Module):
         )
 
     def forward(self, state):
-        reward = self.reward_model(state)
-        return torch.distributions.independent.Independent(
-                torch.distributions.Normal(reward, 1), 1)
+        reward = self.reward_model(state).squeeze(dim=1)
+        return reward
+
 
 class TransitionModel(nn.Module):
     def __init__(self, state_size, hidden_size, action_size, history_size):
@@ -311,7 +310,7 @@ class CLUBSample(nn.Module):  # Sampled version of the CLUB estimator
 
     def get_mu_var_samples(self, state_dict):
         dist = state_dict["distribution"]
-        sample = state_dict["sample"] #dist.sample()  # Use state_dict["sample"] if you want to use the same sample for all the losses
+        sample = dist.sample()  # Use state_dict["sample"] if you want to use the same sample for all the losses
         mu = dist.mean
         var = dist.variance
         return mu, var, sample
@@ -327,12 +326,8 @@ class CLUBSample(nn.Module):  # Sampled version of the CLUB estimator
         mu_curr, var_curr, _ = self.get_mu_var_samples(self.current_states)
         mu_neg, var_neg, _ = self.get_mu_var_samples(self.negative_current_states)
         
-        sample_size = pred_sample.shape[0]
-        random_index = torch.randperm(sample_size).long()
-
         pos = (-(mu_curr - pred_sample)**2 /var_curr).sum(dim=1).mean(dim=0)
-        neg = (-(mu_curr - pred_sample[random_index])**2 /var_curr).sum(dim=1).mean(dim=0)
-        #neg = (-(mu_neg - pred_sample)**2 /var_neg).sum(dim=1).mean(dim=0)
+        neg = (-(mu_neg - pred_sample)**2 /var_neg).sum(dim=1).mean(dim=0)
         upper_bound = pos - neg
 
         return upper_bound/2

DPI/train.py

@@ -65,7 +65,6 @@ def parse_args():
     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=2, type=int)
-    parser.add_argument('--td_lambda', default=0.95, type=int)
     # reward
     parser.add_argument('--reward_lr', default=1e-4, type=float)
     # actor
@@ -78,7 +77,6 @@ def parse_args():
     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=1e-3, 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('--encoder_tau', default=0.005, type=float)
     parser.add_argument('--encoder_stride', default=1, type=int)
@@ -96,7 +94,6 @@ def parse_args():
     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('--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')
@@ -105,7 +102,6 @@ def parse_args():
     parser.add_argument('--transition_model_type', default='', type=str, choices=['', 'deterministic', 'probabilistic', 'ensemble'])
     parser.add_argument('--render', default=False, action='store_true')
     parser.add_argument('--port', default=2000, type=int)
-    parser.add_argument('--num_likelihood_updates', default=5, type=int)
     args = parser.parse_args()
     return args
 
@@ -149,6 +145,12 @@ class DPI:
                                         seq_len=self.args.episode_length, 
                                         batch_size=args.batch_size,
                                         args=self.args)
+        self.data_buffer_clean = ReplayBuffer(size=self.args.replay_buffer_capacity, 
+                                              obs_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size),
+                                              action_size=self.env.action_space.shape[0], 
+                                              seq_len=self.args.episode_length, 
+                                              batch_size=args.batch_size,
+                                              args=self.args)
 
         # create work directory
         utils.make_dir(self.args.work_dir)
@@ -228,13 +230,11 @@ class DPI:
         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.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.value_model, self.transition_model, self.prjoection_head]
@@ -269,7 +269,7 @@ class DPI:
                 next_obs, rew, done, _ = self.env.step(action)
                 #next_obs_clean, _, done, _ = self.env_clean.step(action)
 
-                self.data_buffer.add(obs, action, next_obs, rew, episode_count+1, done)
+                self.data_buffer.add(obs, action, next_obs, episode_count+1, done)
                 #self.data_buffer_clean.add(obs_clean, action, next_obs_clean, episode_count+1, done)
                 
                 if args.save_video:
@@ -293,12 +293,11 @@ class DPI:
         self.collect_sequences(self.args.batch_size)
 
         # Group observations and next_observations by steps from past to present
-        last_observations = torch.tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float()[:self.args.episode_length-1]
+        last_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float()[:self.args.episode_length-1]
         current_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()[:self.args.episode_length-1]
         next_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()[1:]
         actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()[:self.args.episode_length-1]
         next_actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()[1:]
-        rewards = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"rewards",obs=False)).float()[1:]
 
         # Initialize transition model states
         self.transition_model.init_states(self.args.batch_size, device="cpu") # (N,128) 
@@ -307,7 +306,6 @@ class DPI:
         # Train encoder
         step = 0
         total_steps = 10000
-        metrics = {}
         while step < total_steps:
             for i in range(self.args.episode_length-1):
                 if i > 0:
@@ -316,7 +314,6 @@ class DPI:
                     self.current_states_dict = self.get_features(current_observations[i])
                     self.next_states_dict = self.get_features(next_observations[i], momentum=True)
                     self.action = actions[i] # (N,6)
-                    self.next_action = next_actions[i] # (N,6)
                     history = self.transition_model.prev_history
                     
                     # Encode negative observations
@@ -330,14 +327,14 @@ class DPI:
                     predicted_current_state_dict = self.transition_model.imagine_step(last_states_sample, self.action, self.history)
                     self.history = predicted_current_state_dict["history"]
 
+
+                    
                     # Calculate upper bound loss
-                    likeli_loss, ub_loss = self._upper_bound_minimization(self.last_states_dict, 
+                    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 = torch.tensor(likeli_loss.numpy(),dtype=torch.float32, requires_grad=True)
-                    #ikeli_loss = likeli_loss.mean()
                     
                     # Calculate encoder loss
                     encoder_loss = self._past_encoder_loss(self.current_states_dict,
@@ -359,83 +356,31 @@ class DPI:
                     
 
                     # update models
-                    """
-                    print(likeli_loss)
-                    for i in range(self.args.num_likelihood_updates):
-                        self.past_transition_opt.zero_grad()
-                        print(likeli_loss)
-                        likeli_loss.backward()
-                        nn.utils.clip_grad_norm_(self.past_transition_parameters, self.args.grad_clip_norm)
-                        self.past_transition_opt.step()
-                    print(encoder_loss, ub_loss, lb_loss, step)
-                    """
-
-                    world_model_loss = encoder_loss + ub_loss + lb_loss
+                    world_model_loss = encoder_loss + 1e-1 * ub_loss + lb_loss #1e-1 * ub_loss + 1e-5 * encoder_loss + 1e-1 * lb_loss
+                    print("ub_loss: {:.4f}, encoder_loss: {:.4f}, lb_loss: {:.4f}".format(ub_loss, encoder_loss, lb_loss))
+                    print("world_model_loss: {:.4f}".format(world_model_loss))
                     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()
 
-                    """
-                    if step % self.args.logging_freq:
-                        metrics['Upper Bound Loss'] = ub_loss.item()
-                        metrics['Encoder Loss'] = encoder_loss.item()
-                        metrics["Lower Bound Loss"] = lb_loss.item()
-                        metrics["World Model Loss"] = world_model_loss.item()
-                        wandb.log(metrics)
-                    """
                     
                     # 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)
+                        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(), 
+                                                                                self.action, self.history.detach(), 
                                                                                 imagine_horizon)
-                        #print(imagined_rollout["sample"].shape, imagined_rollout["distribution"][0].sample().shape)
-                    
-                    # 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()
-                    
-                    self.target_returns = self._compute_lambda_return(imag_rews[:-1], 
-                                                               target_imag_vals[:-1], 
-                                                               discounts[:-1] ,
-                                                               self.args.td_lambda, 
-                                                               target_imag_vals[-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)
-
-                    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_dist = self.value_model(value_feat)  
-                    value_loss = -torch.mean(self.discounts * value_dist.log_prob(value_targ).unsqueeze(-1))
-
-                    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()
+                        print(imagined_rollout["sample"].shape, imagined_rollout["distribution"][0].sample().shape)
+                    #exit()
                     
                     step += 1
 
                     if step>total_steps:
                         print("Training finished")
                         break
+                    #exit()
+
                     #print(total_ub_loss, total_encoder_loss)
 
 
@@ -445,9 +390,8 @@ class DPI:
                                  current_states,
                                  negative_current_states,
                                  predicted_current_states)
-        likelihood_loss = club_sample.learning_loss()
-        club_loss = club_sample()
-        return likelihood_loss, club_loss
+        club_loss = club_sample.loglikeli()
+        return club_loss
     
     def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict):
         # current state distribution
@@ -479,19 +423,6 @@ class DPI:
                 x = self.obs_encoder(x)
         return x
 
-    def _compute_lambda_return(self, rewards, values, discounts, td_lam, last_value):
-        next_values = torch.cat([values[1:], last_value.unsqueeze(0)],0)  
-        targets = rewards + discounts * next_values * (1-td_lam)
-        rets =[]
-        last_rew = last_value
-
-        for t in range(rewards.shape[0]-1, -1, -1):
-            last_rew = targets[t] + discounts[t] * td_lam *(last_rew)
-            rets.append(last_rew)
-
-        returns = torch.flip(torch.stack(rets), [0])
-        return returns
-
 if __name__ == '__main__':
     args = parse_args()
 

DPI/utils.py

@@ -120,17 +120,15 @@ class ReplayBuffer:
         self.args = args
         self.observations = np.empty((size, *obs_shape), dtype=np.uint8)
         self.actions = np.empty((size, action_size), dtype=np.float32)
-        self.rewards = np.empty((size,1), dtype=np.float32)
         self.next_observations = np.empty((size, *obs_shape), dtype=np.uint8)
         self.episode_count = np.zeros((size,), dtype=np.uint8) 
         self.terminals = np.empty((size,), dtype=np.float32)
         self.steps, self.episodes = 0, 0
     
-    def add(self, obs, ac, next_obs, rew, episode_count, done):
+    def add(self, obs, ac, next_obs, episode_count, done):
         self.observations[self.idx] = obs
         self.actions[self.idx] = ac
         self.next_observations[self.idx] = next_obs
-        self.rewards[self.idx] = rew
         self.episode_count[self.idx] = episode_count
         self.terminals[self.idx] = done
         self.idx = (self.idx + 1) % self.size