DPI/models.py

@@ -23,24 +23,23 @@ class ObservationEncoder(nn.Module):
 
         self.convs = nn.Sequential(*layers)
 
-        self.fc = nn.Linear(256 * obs_shape[0], 2 * state_size)    # 9 if 3 frames stacked
+        self.fc = nn.Linear(256 * 3 * 3, 2 * state_size)
 
     def forward(self, x):
-        x_reshaped = x.reshape(-1, *x.shape[-3:])
-        x_embed = self.convs(x_reshaped)
-        x_embed = torch.reshape(x_embed, (*x.shape[:-3], -1))
-        x = self.fc(x_embed)
+        x = self.convs(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
         
         # Mean and standard deviation
         mean, std = torch.chunk(x, 2, dim=-1)
-        mean = nn.ELU()(mean)
         std = F.softplus(std)
-        std = torch.clamp(std, min=0.0, max=1e1)
+        std = torch.clamp(std, min=0.0, max=1e5)
 
         # Normal Distribution
         dist = self.get_dist(mean, std)
 
         # Sampling via reparameterization Trick
+        #x = dist.rsample()
         x = self.reparameterize(mean, std)
 
         encoded_output = {"sample": x, "distribution": dist}
@@ -64,7 +63,7 @@ class ObservationDecoder(nn.Module):
         self.output_shape = output_shape
         self.input_size = 256 * 3 * 3
         self.in_channels = [self.input_size, 256, 128, 64]
-        self.out_channels = [256, 128, 64, 9]
+        self.out_channels = [256, 128, 64, 3]
 
         if output_shape[1] == 84:
             self.kernels = [5, 7, 5, 6]
@@ -95,50 +94,43 @@ class ObservationDecoder(nn.Module):
 
 
 class Actor(nn.Module):
-    def __init__(self, state_size, hidden_size, action_size, num_layers=4, min_std=1e-4, init_std=5, mean_scale=5):
+    def __init__(self, state_size, hidden_size, action_size, num_layers=5):
         super().__init__()
         self.state_size = state_size    
         self.hidden_size = hidden_size
         self.action_size = action_size
         self.num_layers = num_layers
 
-        self._min_std = min_std
-        self._init_std = init_std
-        self._mean_scale = mean_scale
+        self._min_std=torch.Tensor([1e-4])[0]
+        self._init_std=torch.Tensor([5])[0]
+        self._mean_scale=torch.Tensor([5])[0]
 
         layers = []
         for i in range(self.num_layers):
             input_channels = state_size if i == 0 else self.hidden_size
-            layers.append(nn.Linear(input_channels, self.hidden_size))
-            layers.append(nn.ReLU())
-        layers.append(nn.Linear(self.hidden_size, 2*self.action_size))
+            output_channels = self.hidden_size if i!= self.num_layers-1 else 2*action_size
+            layers.append(nn.Linear(input_channels, output_channels))
+            layers.append(nn.LeakyReLU())
         self.action_model = nn.Sequential(*layers)
 
     def get_dist(self, mean, std):
         distribution = torch.distributions.Normal(mean, std)
         distribution = torch.distributions.transformed_distribution.TransformedDistribution(distribution, TanhBijector())
         distribution = torch.distributions.independent.Independent(distribution, 1)
-        distribution = SampleDist(distribution)
         return distribution
     
-    def add_exploration(self, action, action_noise=0.3):
-        return torch.clamp(torch.distributions.Normal(action, action_noise).rsample(), -1, 1)
-    
     def forward(self, features):
         out = self.action_model(features)
         mean, std = torch.chunk(out, 2, dim=-1) 
         
-        raw_init_std = np.log(np.exp(self._init_std) - 1)
+        raw_init_std = torch.log(torch.exp(self._init_std) - 1)
         action_mean = self._mean_scale * torch.tanh(mean / self._mean_scale)
         action_std = F.softplus(std + raw_init_std) + self._min_std
 
         dist = self.get_dist(action_mean, action_std)
-        sample = dist.rsample() #self.reparameterize(action_mean, action_std)
+        sample = dist.rsample()
         return sample
         
-    def reparameterize(self, mu, std):
-        eps = torch.randn_like(std)
-        return mu + eps * std
 
 class ValueModel(nn.Module):
     def __init__(self, state_size, hidden_size, num_layers=4):
@@ -148,12 +140,11 @@ class ValueModel(nn.Module):
         self.num_layers = num_layers
 
         layers = []
-        for i in range(self.num_layers-1):
+        for i in range(self.num_layers):
             input_channels = state_size if i == 0 else self.hidden_size
-            output_channels = self.hidden_size
+            output_channels = self.hidden_size if i!= self.num_layers-1 else 1
             layers.append(nn.Linear(input_channels, output_channels))
             layers.append(nn.LeakyReLU())
-        layers.append(nn.Linear(self.hidden_size, int(np.prod(1))))
         self.value_model = nn.Sequential(*layers) 
 
     def forward(self, state):
@@ -178,7 +169,6 @@ class RewardModel(nn.Module):
         return torch.distributions.independent.Independent(
                 torch.distributions.Normal(reward, 1), 1)
 
-"""
 class TransitionModel(nn.Module):
     def __init__(self, state_size, hidden_size, action_size, history_size):
         super().__init__()
@@ -190,7 +180,6 @@ class TransitionModel(nn.Module):
         self.act_fn = nn.LeakyReLU()
 
         self.fc_state_action = nn.Linear(state_size + action_size, hidden_size)
-        self.ln = nn.LayerNorm(hidden_size)
         self.history_cell = nn.GRUCell(hidden_size + history_size, history_size)
         self.fc_state_prior  = nn.Linear(history_size + state_size + action_size, 2 * state_size)
         self.fc_state_posterior = nn.Linear(history_size + state_size + action_size, 2 * state_size)
@@ -205,25 +194,12 @@ class TransitionModel(nn.Module):
         distribution = torch.distributions.independent.Independent(distribution, 1)
         return distribution
 
-    def stack_states(self, states, dim=0):        
-        s = dict(
-            mean = torch.stack([state['mean'] for state in states], dim=dim),
-            std  = torch.stack([state['std'] for state in states], dim=dim),
-            sample = torch.stack([state['sample'] for state in states], dim=dim),
-            history = torch.stack([state['history'] for state in states], dim=dim),)
-        if 'distribution' in states:
-            dist = dict(distribution = [state['distribution'] for state in states])
-            s.update(dist)
-        return s
+    def imagine_step(self, prev_state, prev_action, prev_history):
+        state_action = self.act_fn(self.fc_state_action(torch.cat([prev_state, prev_action], dim=-1)))
+        prev_hist = prev_history.detach()
+        history = self.history_cell(torch.cat([state_action, prev_hist], dim=-1), prev_hist)
 
-    def seq_to_batch(self, state, name):
-        return dict(
-            sample = torch.reshape(state[name], (state[name].shape[0]* state[name].shape[1], *state[name].shape[2:])))
-    
-    def imagine_step(self, state, action, history):
-        state_action = self.ln(self.act_fn(self.fc_state_action(torch.cat([state, action], dim=-1))))
-        imag_hist = self.history_cell(torch.cat([state_action, history], dim=-1), history)
-        state_prior = self.fc_state_prior(torch.cat([imag_hist, state, action], dim=-1))
+        state_prior = self.fc_state_prior(torch.cat([history, prev_state, prev_action], dim=-1))
         state_prior_mean, state_prior_std = torch.chunk(state_prior, 2, dim=-1)
         state_prior_std = F.softplus(state_prior_std)
 
@@ -232,103 +208,19 @@ class TransitionModel(nn.Module):
 
         # Sampling via reparameterization Trick
         sample_state_prior = self.reparemeterize(state_prior_mean, state_prior_std)
-        prior = {"mean": state_prior_mean, "std": state_prior_std, "sample": sample_state_prior, "history": imag_hist, "distribution": state_prior_dist}
+        prior = {"mean": state_prior_mean, "std": state_prior_std, "sample": sample_state_prior, "history": history, "distribution": state_prior_dist}
         return prior
     
-    def imagine_rollout(self, state, action, history, horizon):
-        imagined_priors = []
-        for i in range(horizon):
-            prior = self.imagine_step(state, action, history)
-            state = prior["sample"]
-            history = prior["history"]
-            imagined_priors.append(prior)
-        imagined_priors = self.stack_states(imagined_priors, dim=0)
-        return imagined_priors
-    
-    def observe_step(self, prev_state, prev_action, prev_history, nonterms):
-        state_action = self.ln(self.act_fn(self.fc_state_action(torch.cat([prev_state, prev_action], dim=-1))))
-        current_history = self.history_cell(torch.cat([state_action, prev_history], dim=-1), prev_history)
-        state_prior = self.fc_state_prior(torch.cat([prev_history, prev_state, prev_action], dim=-1))
-        state_prior_mean, state_prior_std = torch.chunk(state_prior*nonterms, 2, dim=-1)
-        state_prior_std = F.softplus(state_prior_std) + 0.1
-        sample_state_prior = state_prior_mean + torch.randn_like(state_prior_mean) * state_prior_std
-        prior = {"mean": state_prior_mean, "std": state_prior_std, "sample": sample_state_prior, "history": current_history}
-        return prior
-
-    def observe_rollout(self, rollout_states, rollout_actions, init_history, nonterms):
-        observed_rollout = []
-        for i in range(rollout_states.shape[0]):
-            actions = rollout_actions[i] * nonterms[i]
-            prior = self.observe_step(rollout_states[i], actions, init_history, nonterms[i])
-            init_history = prior["history"]
-            observed_rollout.append(prior)
-        observed_rollout = self.stack_states(observed_rollout, dim=0)
-        return observed_rollout
-
-    def reparemeterize(self, mean, std):
-        eps = torch.randn_like(std)
-        return mean + eps * std
-"""
-
-class TransitionModel(nn.Module):
-    def __init__(self, state_size, hidden_size, action_size, history_size):
-        super().__init__()
-
-        self.state_size = state_size
-        self.hidden_size = hidden_size
-        self.action_size = action_size
-        self.history_size = history_size
-        self.act_fn = nn.ELU()
-
-        self.fc_state_action = nn.Linear(state_size + action_size, hidden_size)
-        self.history_cell = nn.GRUCell(hidden_size, history_size)
-        self.fc_state_mu = nn.Linear(history_size + hidden_size, state_size)
-        self.fc_state_sigma = nn.Linear(history_size + hidden_size, state_size)
-
-        self.batch_norm = nn.BatchNorm1d(hidden_size)
-        self.batch_norm2 = nn.BatchNorm1d(state_size)
-
-        self.min_sigma = 1e-4
-        self.max_sigma = 1e0
-
-    def init_states(self, batch_size, device):
-        self.prev_state = torch.zeros(batch_size, self.state_size).to(device)
-        self.prev_action = torch.zeros(batch_size, self.action_size).to(device)
-        self.prev_history = torch.zeros(batch_size, self.history_size).to(device)
-            
-    def get_dist(self, mean, std):
-        distribution = torch.distributions.Normal(mean, std)
-        distribution = torch.distributions.independent.Independent(distribution, 1)
-        return distribution
-
     def stack_states(self, states, dim=0):        
         s = dict(
             mean = torch.stack([state['mean'] for state in states], dim=dim),
             std  = torch.stack([state['std'] for state in states], dim=dim),
             sample = torch.stack([state['sample'] for state in states], dim=dim),
             history = torch.stack([state['history'] for state in states], dim=dim),)
-        if 'distribution' in states:
         dist = dict(distribution = [state['distribution'] for state in states])
         s.update(dist)
         return s
     
-    def seq_to_batch(self, state, name):
-        return dict(
-            sample = torch.reshape(state[name], (state[name].shape[0]* state[name].shape[1], *state[name].shape[2:])))
-    
-    def imagine_step(self, state, action, history):
-        next_state_action_enc = self.act_fn(self.batch_norm(self.fc_state_action(torch.cat([state, action], dim=-1))))
-        imag_history = self.history_cell(next_state_action_enc, history)
-        next_state_mu = self.act_fn(self.batch_norm2(self.fc_state_mu(torch.cat([next_state_action_enc, imag_history], dim=-1))))
-        next_state_sigma = torch.sigmoid(self.fc_state_sigma(torch.cat([next_state_action_enc, imag_history], dim=-1)))
-        next_state_sigma = self.min_sigma + (self.max_sigma - self.min_sigma) * next_state_sigma
-
-        # Normal Distribution
-        next_state_dist = self.get_dist(next_state_mu, next_state_sigma)
-        next_state_sample = self.reparemeterize(next_state_mu, next_state_sigma)
-        prior = {"mean": next_state_mu, "std": next_state_sigma, "sample": next_state_sample, "history": imag_history, "distribution": next_state_dist}
-        return prior
-    
     def imagine_rollout(self, state, action, history, horizon):
         imagined_priors = []
         for i in range(horizon):
@@ -339,30 +231,8 @@ class TransitionModel(nn.Module):
         imagined_priors = self.stack_states(imagined_priors, dim=0)
         return imagined_priors
     
-    def observe_step(self, prev_state, prev_action, prev_history):
-        state_action_enc = self.act_fn(self.batch_norm(self.fc_state_action(torch.cat([prev_state, prev_action], dim=-1))))
-        current_history = self.history_cell(state_action_enc, prev_history)
-        state_mu = self.act_fn(self.batch_norm2(self.fc_state_mu(torch.cat([state_action_enc, prev_history], dim=-1))))
-        state_sigma = F.softplus(self.fc_state_sigma(torch.cat([state_action_enc, prev_history], dim=-1)))
-
-        sample_state = state_mu + torch.randn_like(state_mu) * state_sigma
-        state_enc = {"mean": state_mu, "std": state_sigma, "sample": sample_state, "history": current_history}
-        return state_enc
-
-    def observe_rollout(self, rollout_states, rollout_actions, init_history, nonterms):
-        observed_rollout = []
-        for i in range(rollout_states.shape[0]):
-            rollout_states_ = rollout_states[i]
-            rollout_actions_ = rollout_actions[i]
-            init_history_ = nonterms[i] * init_history
-            state_enc = self.observe_step(rollout_states_, rollout_actions_, init_history_)
-            init_history = state_enc["history"]
-            observed_rollout.append(state_enc)
-        observed_rollout = self.stack_states(observed_rollout, dim=0)
-        return observed_rollout
-
     def reparemeterize(self, mean, std):
-        eps = torch.randn_like(mean)
+        eps = torch.randn_like(std)
         return mean + eps * std
 
 
@@ -430,7 +300,6 @@ class ContrastiveHead(nn.Module):
         return logits
     
 
-"""
 class CLUBSample(nn.Module):  # Sampled version of the CLUB estimator
     def __init__(self, last_states, current_states, negative_current_states, predicted_current_states):
         super(CLUBSample, self).__init__()
@@ -444,7 +313,7 @@ class CLUBSample(nn.Module):  # Sampled version of the CLUB estimator
         sample = state_dict["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.detach(), var.detach(), sample.detach()
+        return mu, var, sample
         
     def loglikeli(self):
         _, _, pred_sample = self.get_mu_var_samples(self.predicted_current_states)
@@ -461,136 +330,15 @@ class CLUBSample(nn.Module):  # Sampled version of the CLUB estimator
         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_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)
         upper_bound = pos - neg
 
         return upper_bound/2
 
     def learning_loss(self):
         return - self.loglikeli()
-"""
-
-class CLUBSample(nn.Module):  # Sampled version of the CLUB estimator
-    def __init__(self, x_dim, y_dim, hidden_size):
-        super(CLUBSample, self).__init__()
-        self.p_mu = nn.Sequential(nn.Linear(x_dim, hidden_size//2),
-                                       nn.ReLU(),
-                                       nn.Linear(hidden_size//2, y_dim))
-
-        self.p_logvar = nn.Sequential(nn.Linear(x_dim, hidden_size//2),
-                                       nn.ReLU(),
-                                       nn.Linear(hidden_size//2, y_dim),
-                                       nn.Tanh())
-
-    def get_mu_logvar(self, x_samples):
-        mu = self.p_mu(x_samples)
-        logvar = self.p_logvar(x_samples)
-        return mu, logvar
-     
-        
-    def loglikeli(self, x_samples, y_samples):
-        mu, logvar = self.get_mu_logvar(x_samples)
-        return (-(mu - y_samples)**2 /logvar.exp()-logvar).sum(dim=1).mean(dim=0)
-    
-
-    def forward(self, x_samples, y_samples, y_negatives):
-        mu, logvar = self.get_mu_logvar(x_samples)
-        
-        sample_size = x_samples.shape[0]
-        #random_index = torch.randint(sample_size, (sample_size,)).long()
-        random_index = torch.randperm(sample_size).long()
-        
-        positive = -(mu - y_samples)**2 / logvar.exp()
-        #negative = - (mu - y_samples[random_index])**2 / logvar.exp()
-        negative = -(mu - y_negatives)**2 / logvar.exp()
-        upper_bound = (positive.sum(dim = -1) - negative.sum(dim = -1)).mean()
-        return upper_bound/2.
-
-    def learning_loss(self, x_samples, y_samples):
-        return -self.loglikeli(x_samples, y_samples)
-
-
-class QFunction(nn.Module):
-    """MLP for q-function."""
-    def __init__(self, obs_dim, action_dim, hidden_dim):
-        super().__init__()
-
-        self.trunk = nn.Sequential(
-            nn.Linear(obs_dim + action_dim, hidden_dim), nn.ReLU(),
-            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
-            nn.Linear(hidden_dim, 1)
-        )
-
-    def forward(self, obs, action):
-        assert obs.size(0) == action.size(0)
-
-        obs_action = torch.cat([obs, action], dim=1)
-        return self.trunk(obs_action)
-
-
-class Critic(nn.Module):
-    """Critic network, employes two q-functions."""
-    def __init__(
-        self, obs_shape, action_shape, hidden_dim, encoder_feature_dim):
-        super().__init__()
-
-        self.Q1 = QFunction(
-            self.encoder.feature_dim, action_shape[0], hidden_dim
-        )
-        self.Q2 = QFunction(
-            self.encoder.feature_dim, action_shape[0], hidden_dim
-        )
-
-        self.outputs = dict()
-
-    def forward(self, obs, action, detach_encoder=False):
-        # detach_encoder allows to stop gradient propogation to encoder
-        obs = self.encoder(obs, detach=detach_encoder)
-
-        q1 = self.Q1(obs, action)
-        q2 = self.Q2(obs, action)
-
-        self.outputs['q1'] = q1
-        self.outputs['q2'] = q2
-
-        return q1, q2
-
-class SampleDist:
-    def __init__(self, dist, samples=100):
-        self._dist = dist
-        self._samples = samples
-
-    @property
-    def name(self):
-        return 'SampleDist'
-
-    def __getattr__(self, name):
-        return getattr(self._dist, name)
-
-    def mean(self):
-        sample = self._dist.rsample(self._samples)
-        return torch.mean(sample, 0)
-
-    def mode(self):
-        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
-        sample = dist.rsample()
-        logprob = dist.log_prob(sample)
-        batch_size = sample.size(1)
-        feature_size = sample.size(2)
-        indices = torch.argmax(logprob, dim=0).reshape(1, batch_size, 1).expand(1, batch_size, feature_size)
-        return torch.gather(sample, 0, indices).squeeze(0)
-
-    def entropy(self):
-        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
-        sample = dist.rsample()
-        logprob = dist.log_prob(sample)
-        return -torch.mean(logprob, 0)
-
-    def sample(self):
-        return self._dist.sample()
     
 
 if "__name__ == __main__":
-    tr = TransitionModel(50, 512, 1, 256)
-
+    pass

DPI/replay_buffer.py

@@ -1,51 +0,0 @@
-import torch
-import numpy as np
-
-
-class ReplayBuffer:
-
-    def __init__(self, size, obs_shape, action_size, seq_len, batch_size):
-        self.size = size
-        self.obs_shape = obs_shape
-        self.action_size = action_size
-        self.seq_len = seq_len
-        self.batch_size = batch_size
-        self.idx = 0
-        self.full = False
-        self.observations = np.empty((size, *obs_shape), dtype=np.uint8) 
-        self.next_observations = np.empty((size, *obs_shape), dtype=np.uint8)
-        self.actions = np.empty((size, action_size), dtype=np.float32)
-        self.rewards = np.empty((size,), dtype=np.float32) 
-        self.terminals = np.empty((size,), dtype=np.float32)
-        self.steps, self.episodes = 0, 0
-    
-    def add(self, obs, ac, next_obs, rew, done):
-        self.observations[self.idx] = obs
-        self.next_observations[self.idx] = next_obs
-        self.actions[self.idx] = ac
-        self.rewards[self.idx] = rew
-        self.terminals[self.idx] = done
-        self.idx = (self.idx + 1) % self.size
-        self.full = self.full or self.idx == 0
-        self.steps += 1 
-        self.episodes = self.episodes + (1 if done else 0)
-
-    def _sample_idx(self, L):
-        valid_idx = False
-        while not valid_idx:
-            idx = np.random.randint(0, self.size if self.full else self.idx - L)
-            idxs = np.arange(idx, idx + L) % self.size
-            valid_idx = not self.idx in idxs[1:] 
-        return idxs
-
-    def _retrieve_batch(self, idxs, n, L):
-        vec_idxs = idxs.transpose().reshape(-1)  # Unroll indices
-        observations = self.observations[vec_idxs]
-        next_observations = self.next_observations[vec_idxs]
-        return observations.reshape(L, n, *observations.shape[1:]),self.actions[vec_idxs].reshape(L, n, -1), next_observations.reshape(L, n, *next_observations.shape[1:]), self.rewards[vec_idxs].reshape(L, n), self.terminals[vec_idxs].reshape(L, n)
-
-    def sample(self):
-        n = self.batch_size
-        l = self.seq_len
-        obs,acs,nxt_obs,rews,terms= self._retrieve_batch(np.asarray([self._sample_idx(l) for _ in range(n)]), n, l)
-        return obs,acs,nxt_obs,rews,terms

DPI/train.py

@@ -6,12 +6,11 @@ import wandb
 import random
 import argparse
 import numpy as np
-from collections import OrderedDict
 
 import utils
-from utils import ReplayBuffer, FreezeParameters, make_env, preprocess_obs, soft_update_params, save_image, shuffle_along_axis, Logger
-from replay_buffer import ReplayBuffer
+from utils import ReplayBuffer, FreezeParameters, make_env, preprocess_obs, soft_update_params, save_image
 from models import ObservationEncoder, ObservationDecoder, TransitionModel, Actor, ValueModel, RewardModel, ProjectionHead, ContrastiveHead, CLUBSample
+from logger import Logger
 from video import VideoRecorder
 from dmc2gym.wrappers import set_global_var
 
@@ -41,22 +40,19 @@ def parse_args():
     parser.add_argument('--resource_files', type=str)
     parser.add_argument('--eval_resource_files', type=str)
     parser.add_argument('--img_source', default=None, type=str, choices=['color', 'noise', 'images', 'video', 'none'])
-    parser.add_argument('--total_frames', default=5000, type=int)  # 10000
+    parser.add_argument('--total_frames', default=1000, type=int)  # 10000
     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)
     # train
     parser.add_argument('--agent', default='dpi', type=str, choices=['baseline', 'bisim', 'deepmdp', 'db', 'dpi', 'rpc'])
-    parser.add_argument('--init_steps', default=5000, type=int)
-    parser.add_argument('--num_train_steps', default=100000, type=int)
-    parser.add_argument('--update_steps', default=10, type=int)
-    parser.add_argument('--batch_size', default=64, type=int) 
-    parser.add_argument('--state_size', default=100, type=int)
-    parser.add_argument('--hidden_size', default=512, type=int)
-    parser.add_argument('--history_size', default=256, type=int)
-    parser.add_argument('--episode_collection', default=5, type=int)
-    parser.add_argument('--episodes_buffer', default=5, type=int, help='Initial number of episodes to store in the buffer')
+    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-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)
@@ -64,33 +60,42 @@ 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-6, type=float)
+    parser.add_argument('--value_lr', default=8e-5, 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-6, type=float)     
+    parser.add_argument('--actor_lr', default=8e-5, 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-6, type=float)
-    parser.add_argument('--decoder_lr', default=6e-6, type=float)
-    parser.add_argument('--reward_lr', default=8e-6, type=float)
-    parser.add_argument('--encoder_tau', default=0.005, type=float)
+    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('--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=2500, type=int)
+    parser.add_argument('--saving_interval', default=1000, 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')
@@ -121,7 +126,6 @@ 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
@@ -133,14 +137,14 @@ class DPI:
             self.env = utils.FrameStack(self.env, k=self.args.frame_stack)
             self.env = utils.ActionRepeat(self.env, self.args.action_repeat)
             self.env = utils.NormalizeActions(self.env)
-            self.env = utils.TimeLimit(self.env, 1000 // args.action_repeat)
 
         # create replay buffer
-        self.data_buffer = ReplayBuffer(self.args.replay_buffer_capacity, 
-                                        self.env.observation_space.shape, 
-                                        self.env.action_space.shape[0],
-                                        self.args.episode_length, 
-                                        self.args.batch_size)
+        self.data_buffer = 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)
@@ -157,19 +161,16 @@ class DPI:
                             obs_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size), # (9,84,84)
                             state_size=self.args.state_size # 128
                             ).to(device)
-        self.obs_encoder.apply(self.init_weights)
         
         self.obs_encoder_momentum = ObservationEncoder(
                             obs_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size), # (9,84,84)
                             state_size=self.args.state_size # 128
                             ).to(device)
-        self.obs_encoder_momentum.apply(self.init_weights)
 
         self.obs_decoder = ObservationDecoder(
                             state_size=self.args.state_size, # 128
-                            output_shape=(self.args.channels*self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84)
+                            output_shape=(self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84)
                             ).to(device)
-        self.obs_decoder.apply(self.init_weights)
 
         self.transition_model = TransitionModel(
                             state_size=self.args.state_size, # 128
@@ -177,7 +178,6 @@ class DPI:
                             action_size=self.env.action_space.shape[0], # 6
                             history_size=self.args.history_size, # 128
                             ).to(device)
-        self.transition_model.apply(self.init_weights)
         
         # Actor Model
         self.actor_model = Actor(
@@ -185,27 +185,22 @@ 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
                             hidden_size=self.args.hidden_size, # 256
                             ).to(device)
-        self.value_model.apply(self.init_weights)
         
         self.target_value_model = ValueModel(
                             state_size=self.args.state_size, # 128
                             hidden_size=self.args.hidden_size, # 256
                             ).to(device)
-        self.target_value_model.apply(self.init_weights)
 
         self.reward_model = RewardModel(
                             state_size=self.args.state_size, # 128
                             hidden_size=self.args.hidden_size, # 256
                             ).to(device)
-        self.reward_model.apply(self.init_weights)
         
         # Contrastive Models
         self.prjoection_head = ProjectionHead(
@@ -224,32 +219,23 @@ class DPI:
                             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.club_sample.parameters()) + \
-                                      list(self.contrastive_head.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.past_transition_parameters = self.transition_model.parameters()
 
         # optimizers
-        self.world_model_opt = torch.optim.Adam(self.world_model_parameters, self.args.world_model_lr,eps=1e-6)
-        self.value_opt = torch.optim.Adam(self.value_model.parameters(), self.args.value_lr,eps=1e-6)
-        self.actor_opt = torch.optim.Adam(self.actor_model.parameters(), self.args.actor_lr,eps=1e-6)
-        self.decoder_opt = torch.optim.Adam(self.obs_decoder.parameters(), self.args.decoder_lr,eps=1e-6)
-        self.reward_opt = torch.optim.Adam(self.reward_model.parameters(), self.args.reward_lr,eps=1e-6)
+        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.prjoection_head, self.transition_model, self.club_sample, self.contrastive_head, 
-                                    self.obs_encoder_momentum, self.prjoection_head_momentum]
+        self.world_model_modules = [self.obs_encoder, self.obs_decoder, self.reward_model, self.transition_model, self.prjoection_head]
         self.value_modules = [self.value_model]
         self.actor_modules = [self.actor_model]
-        self.decoder_modules = [self.obs_decoder]
-        self.reward_modules = [self.reward_model]
 
         if use_saved:
             self._use_saved_models(saved_model_dir)
@@ -259,432 +245,280 @@ class DPI:
         self.obs_decoder.load_state_dict(torch.load(os.path.join(saved_model_dir, 'obs_decoder.pt')))
         self.transition_model.load_state_dict(torch.load(os.path.join(saved_model_dir, 'transition_model.pt')))
 
-    def collect_random_sequences(self, seed_steps):
+    def collect_sequences(self, episodes,  random=True, actor_model=None, encoder_model=None):
         obs = self.env.reset()
         done = False
+
         all_rews = []
-        self.global_episodes += 1
+        #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)
+
             epi_reward = 0
-        for _ in tqdm.tqdm(range(seed_steps), desc='Collecting episodes'):
+            for i in range(self.args.episode_length):
+                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()
+                
                 next_obs, rew, done, _ = self.env.step(action)
-            self.data_buffer.add(obs, action, next_obs, rew, done)
-            obs = next_obs  
-            epi_reward += rew
-            if done:
+                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
-                all_rews.append(epi_reward)
-                epi_reward = 0
-        return all_rews
-    
-    def collect_sequences(self, collect_steps, actor_model):
-        obs = self.env.reset()
-        done = False
-        all_rews = []
-        self.global_episodes += 1
-        epi_reward = 0
-        for episode_count in tqdm.tqdm(range(collect_steps), desc='Collecting episodes'):
-            with torch.no_grad():
-                obs_  = torch.tensor(obs.copy(), dtype=torch.float32)
-                obs_ = preprocess_obs(obs_).to(device)
-                #state = self.get_features(obs_)["sample"].unsqueeze(0)
-                state = self.get_features(obs_)["distribution"].rsample()
-                action = actor_model(state)
-                action = actor_model.add_exploration(action)
-                action = action.cpu().numpy()[0]
-            next_obs, rew, done, _ = self.env.step(action)
-            self.data_buffer.add(obs, action, next_obs, rew, done)
-
-            if done:
-                obs = self.env.reset()
-                done = False
-                all_rews.append(epi_reward)
-                epi_reward = 0
                 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))
         return all_rews
 
     def train(self, step, total_steps):
-        # logger
-        logdir = os.path.dirname(os.path.realpath(__file__)) + "/log/logs/"
-        if not(os.path.exists(logdir)):
-            os.makedirs(logdir)
-        initial_logs = OrderedDict()
-        logger = Logger(logdir)
+        counter = 0
+        while step < total_steps:
             
-        episodic_rews = self.collect_random_sequences(self.args.init_steps//args.action_repeat)
-        self.global_step = self.data_buffer.steps
-        
-        initial_logs.update({
-                'train_avg_reward':np.mean(episodic_rews),  
-                'train_max_reward': np.max(episodic_rews),
-                'train_min_reward': np.min(episodic_rews),
-                'train_std_reward':np.std(episodic_rews),
-            })
-        logger.log_scalars(initial_logs, step=0)
-        logger.flush()
-        
-
-        while self.global_step  < total_steps:
-            logs = OrderedDict()
-            step += 1
-            for update_steps in range(self.args.update_steps):
-                model_loss, actor_loss, value_loss, actor_model = self.update((step-1)*args.update_steps + update_steps)
-            
-            initial_logs.update({
-                'model_loss' : model_loss,
-                'actor_loss': actor_loss,
-                'value_loss': value_loss,
-                'train_avg_reward':np.mean(episodic_rews),  
-                'train_max_reward': np.max(episodic_rews),
-                'train_min_reward': np.min(episodic_rews),
-                'train_std_reward':np.std(episodic_rews),
-            })
-            logger.log_scalars(logs, self.global_step)
-
-
-            print("########## Global Step:", self.global_step, " ##########")
-            for key, value in initial_logs.items():
-                print(key, " : ", value)
-            
-            episodic_rews = self.collect_sequences(1000//self.args.action_repeat, actor_model)
-
-            if  self.global_step % 3150 == 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()
-   
-            self.global_step = self.data_buffer.steps * self.args.action_repeat
-
-        """
             # collect experience
             if step !=0:
                 encoder = self.obs_encoder
                 actor = self.actor_model
-            all_rews = self.collect_sequences(self.args.episode_collection, actor_model=actor, encoder_model=encoder)
-        """
+                #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)
+            else:
+                all_rews = self.collect_sequences(self.args.batch_size, random=True)
 
-    def collect_batch(self):
-        obs_, acs_, nxt_obs_, rews_, terms_ = self.data_buffer.sample()      
+            # 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))
+            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)
 
-        obs = torch.tensor(obs_, dtype=torch.float32)[1:]
-        last_obs = torch.tensor(obs_, dtype=torch.float32)[:-1]
-        nxt_obs = torch.tensor(nxt_obs_, dtype=torch.float32)[1:]
-        acs = torch.tensor(acs_, dtype=torch.float32)[:-1].to(device)
-        nxt_acs = torch.tensor(acs_, dtype=torch.float32)[1:].to(device)
-        rews = torch.tensor(rews_, dtype=torch.float32)[:-1].to(device).unsqueeze(-1)
-        nonterms = torch.tensor((1.0-terms_), dtype=torch.float32)[:-1].to(device).unsqueeze(-1)
+            # 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)
 
-        last_obs = preprocess_obs(last_obs).to(device)
-        obs = preprocess_obs(obs).to(device)
-        nxt_obs = preprocess_obs(nxt_obs).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)
 
-        return last_obs, obs, nxt_obs, acs, rews, nxt_acs, nonterms
+            # Train encoder
+            if step == 0:
+                step += 1
+            for _ in range(self.args.collection_interval // self.args.episode_length+1):
+                counter += 1
+                for i in range(self.args.episode_length-1):
+                    if i > 0:
+                        # Encode observations and next_observations
+                        self.last_states_dict = self.get_features(last_observations[i])
+                        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
                         
-    def update(self, step):
-        last_observations, current_observations, next_observations, actions, rewards, next_actions, nonterms = self.collect_batch()
+                        # Encode negative observations
+                        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
+                        negative_current_observations = current_observations[random_time_index][idx]
+                        self.negative_current_states_dict = self.obs_encoder(negative_current_observations)
 
-        #last_observations, current_observations, next_observations, actions, next_actions, rewards = self.select_one_batch()
+                        # Predict current state from past state with transition model
+                        last_states_sample = self.last_states_dict["sample"]
+                        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, 
+                                                                            self.current_states_dict,
+                                                                            self.negative_current_states_dict,
+                                                                            predicted_current_state_dict
+                                                                            )
                         
-        world_loss, enc_loss, rew_loss, dec_loss, ub_loss, lb_loss = self.world_model_losses(last_observations, 
-                                                                                            current_observations, 
-                                                                                            next_observations, 
-                                                                                            actions, 
-                                                                                            next_actions, 
-                                                                                            rewards,
-                                                                                            nonterms)             
+                        # Calculate encoder loss
+                        encoder_loss = self._past_encoder_loss(self.current_states_dict,
+                                                            predicted_current_state_dict)
+
+                        # contrastive projection
+                        vec_anchor = predicted_current_state_dict["sample"]
+                        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()
+
+                        # contrastive loss
+                        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]))        
+
+                        # 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_loss.backward()
+                        world_model_loss.backward()
                         nn.utils.clip_grad_norm_(self.world_model_parameters, self.args.grad_clip_norm)
                         self.world_model_opt.step()
 
-        self.decoder_opt.zero_grad()
-        dec_loss.backward()
-        nn.utils.clip_grad_norm_(self.obs_decoder.parameters(), self.args.grad_clip_norm)
-        self.decoder_opt.step()
+                        # update momentum encoder
+                        soft_update_params(self.obs_encoder, self.obs_encoder_momentum, self.args.encoder_tau)
 
-        self.reward_opt.zero_grad()
-        rew_loss.backward()
-        nn.utils.clip_grad_norm_(self.reward_model.parameters(), self.args.grad_clip_norm)
-        self.reward_opt.step()
+                        # update momentum projection head
+                        soft_update_params(self.prjoection_head, self.prjoection_head_momentum, self.args.encoder_tau)
 
-        actor_loss = self.actor_model_losses()
+                        # 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)
+
+                        # 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 = self.value_model_losses()
+                        # 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))
+                        
+                        # update value
                         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)
+                        # update target value
+                        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))
 
-        # update target value networks
-        #if step % self.args.value_target_update_freq == 0:
-        #    self.target_value_model = copy.deepcopy(self.value_model)
                         
                         if step % self.args.logging_freq:
-            writer.add_scalar('World Loss/World Loss', world_loss.detach().item(), step)
-            writer.add_scalar('Main Models Loss/Encoder Loss', enc_loss.detach().item(), step)
-            writer.add_scalar('Main Models Loss/Decoder Loss', dec_loss.detach().item(), step)
+                            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', rew_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('Bound Loss/Lower Bound Loss', lb_loss.detach().item(), step)                            
                         
-        return world_loss.item(), actor_loss.item(), value_loss.item(), self.actor_model
+                        step += 1
+                        if step>total_steps:
+                            print("Training finished")
+                            break
                         
-    def world_model_losses(self, last_obs, curr_obs, nxt_obs, actions, nxt_actions, rewards, nonterms):   
-        # get features
-        self.last_state_feat = self.get_features(last_obs)
-        self.curr_state_feat = self.get_features(curr_obs)
-        self.nxt_state_feat = self.get_features(nxt_obs)
-        self.nxt_state_feat_lb = self.get_features(nxt_obs, momentum=True)
+                        # 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)
 
-        # states
-        self.last_state_enc = self.last_state_feat["distribution"].rsample() #self.last_state_feat["sample"]
-        self.curr_state_enc = self.curr_state_feat["distribution"].rsample() #self.curr_state_feat["sample"]
-        self.nxt_state_enc = self.nxt_state_feat["distribution"].rsample() #self.nxt_state_feat["sample"]
-        self.nxt_state_enc_lb = self.nxt_state_feat_lb["distribution"].rsample() #self.nxt_state_feat_lb["sample"]  
+                        #torch.cuda.empty_cache()    # memory leak issues
                 
-        # predict next states
-        self.transition_model.init_states(self.args.batch_size, device) # (N,128)
-        self.observed_rollout = self.transition_model.observe_rollout(self.last_state_enc, actions, self.transition_model.prev_history, nonterms)
-        self.pred_curr_state_dist = self.transition_model.get_dist(self.observed_rollout["mean"], self.observed_rollout["std"])
-        self.pred_curr_state_enc = self.pred_curr_state_dist.rsample()  #self.observed_rollout["sample"]
+                for j in range(len(all_rews)):
+                                writer.add_scalar('Rewards/Rewards', all_rews[j], count[j])
                 
-        # encoder loss
-        enc_loss = self._encoder_loss(self.curr_state_feat["distribution"], self.pred_curr_state_dist)
 
-        # reward loss
-        rew_dist = self.reward_model(self.curr_state_enc.detach())
-        #print(torch.cat([rew_dist.mean[0], rewards[0]],dim=-1))
-        rew_loss = -torch.mean(rew_dist.log_prob(rewards))
+    def evaluate(self, env, eval_episodes, render=False):
 
-        # decoder loss
-        dec_dist = self.obs_decoder(self.nxt_state_enc.detach())
-        dec_loss = -torch.mean(dec_dist.log_prob(nxt_obs))
+        episode_rew = np.zeros((eval_episodes))
 
-        # upper bound loss
-        past_ub_loss = 0
-        for i in range(self.curr_state_enc.shape[0]):
-            _, ub_loss = self._upper_bound_minimization(self.curr_state_enc[i],
-                                                        self.pred_curr_state_enc[i])
-            ub_loss = ub_loss + past_ub_loss
-            past_ub_loss = ub_loss
-        ub_loss = ub_loss / self.curr_state_enc.shape[0]
-        ub_loss = 1 * ub_loss
+        video_images = [[] for _ in range(eval_episodes)]
 
-        # lower bound loss
-        # contrastive projection
-        vec_anchor = self.pred_curr_state_enc.detach()
-        vec_positive = self.nxt_state_enc_lb.detach()
-        z_anchor = self.prjoection_head(vec_anchor, nxt_actions)
-        z_positive = self.prjoection_head_momentum(vec_positive, nxt_actions)
-
-        # contrastive loss
-        past_lb_loss = 0
-        for i in range(z_anchor.shape[0]):
-            logits = self.contrastive_head(z_anchor[i], z_positive[i])
-            labels = torch.arange(logits.shape[0]).long().to(device)
-            lb_loss = F.cross_entropy(logits, labels) + past_lb_loss
-            past_lb_loss = lb_loss.detach().item()
-        lb_loss = -0.01 * lb_loss/(z_anchor.shape[0])
-        
-        world_loss = enc_loss + ub_loss + lb_loss
-
-        return world_loss, enc_loss , rew_loss, dec_loss, ub_loss, lb_loss
-
-    def actor_model_losses(self):
-        with torch.no_grad():
-            #curr_state_enc = self.curr_state_enc.reshape(self.args.episode_length-1,-1) #self.transition_model.seq_to_batch(self.curr_state_feat, "sample")["sample"]
-            #curr_state_hist = self.observed_rollout["history"].reshape(self.args.episode_length-1,-1) #self.transition_model.seq_to_batch(self.observed_rollout, "history")["sample"]
-            curr_state_enc = self.curr_state_enc.reshape(-1, self.args.state_size)
-            curr_state_hist = self.observed_rollout["history"].reshape(-1, self.args.history_size)
-
-        with FreezeParameters(self.world_model_modules + self.decoder_modules + self.reward_modules + self.value_modules):
-            imagine_horizon = self.args.imagine_horizon
-            action = self.actor_model(curr_state_enc.detach())
-            self.imagined_rollout = self.transition_model.imagine_rollout(curr_state_enc, 
-                                                                     action, curr_state_hist.detach(),
-                                                                     imagine_horizon)
-            self.pred_nxt_state_dist = self.transition_model.get_dist(self.imagined_rollout["mean"], self.imagined_rollout["std"])
-            self.pred_nxt_state_enc = self.pred_nxt_state_dist.rsample() #self.transition_model.reparemeterize(self.imagined_rollout["mean"], self.imagined_rollout["std"])
-            
-        with FreezeParameters(self.world_model_modules + self.value_modules + self.decoder_modules + self.reward_modules):
-            imag_rewards_dist = self.reward_model(self.pred_nxt_state_enc)
-            imag_values_dist = self.value_model(self.pred_nxt_state_enc)
-            imag_rewards = imag_rewards_dist.mean
-            imag_values = imag_values_dist.mean
-            #print(torch.cat([imag_rewards[0], imag_values[0]],dim=-1))
-            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.returns)
-        return actor_loss
-
-    def value_model_losses(self):
-        with torch.no_grad():
-            value_feat = self.pred_nxt_state_enc[:-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))
-        return value_loss
-
-    def select_one_batch(self):
-        # 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]!=False)[0]
-        
-        # 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)
-
-        # 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)
-        
-        # select random starting points
-        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))
-        
-        # shuffle
-        random.shuffle(random_episode_number)
-        random.shuffle(init_index)
-        
-        # select first k elements
-        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)
-
-        return last_observations, current_observations, next_observations, actions, next_actions, rewards
-
-    def evaluate(self, eval_episodes=10):
-        path = path =  os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth"
-        self.restore_checkpoint(path)
-
-        obs = self.env.reset()
+        for i in range(eval_episodes):
+            obs = env.reset()
             done = False
+            prev_state = self.rssm.init_state(1, self.device)
+            prev_action = torch.zeros(1, self.action_size).to(self.device)
 
-        #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
             while not done:
                 with torch.no_grad():
-                    obs = torch.tensor(obs.copy(), dtype=torch.float32).unsqueeze(0)
-                    obs_processed = preprocess_obs(obs).to(device)
-                    state = self.get_features(obs_processed)["distribution"].rsample()
-                    action = self.actor_model(state).cpu().detach().numpy().squeeze()     
-                next_obs, rew, done, _ = self.env.step(action)
-                rewards += rew
+                    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)
+
+                episode_rew[i] += rew
+
+                if render:
+                    video_images[i].append(obs['image'].transpose(1,2,0).copy())
                 obs = next_obs
+        return episode_rew, np.array(video_images[:self.args.max_videos_to_save])
 
-                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 = self.env.reset()
-            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(array).float()
-        return torch.tensor(transposed_array).float()
-
-    def _upper_bound_minimization(self, current_states, predicted_current_states):
-        current_negative_states = shuffle_along_axis(current_states.clone(), axis=0)
-        current_negative_states = shuffle_along_axis(current_negative_states, axis=1)
-        club_loss = self.club_sample(current_states, predicted_current_states, current_negative_states)
-        likelihood_loss = 0
+    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()
         return likelihood_loss, club_loss
     
-    def _encoder_loss(self, curr_states_dist, predicted_curr_states_dist):
+    def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict):
+        # current state distribution
+        curr_states_dist = curr_states_dict["distribution"]
+
+        # predicted current state distribution
+        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.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist).mean()
 
         return loss
     
     def get_features(self, x, momentum=False):        
         if self.args.aug:
-            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)
+            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)
 
         with torch.no_grad():
             if momentum:
@@ -694,19 +528,6 @@ class DPI:
         return x
     
     def _compute_lambda_return(self, rewards, values, discounts, td_lam, last_value):
-        next_values = torch.cat([values[1:], last_value[None]], 0)
-        target = rewards + discounts * next_values * (1 - td_lam)
-        timesteps = list(range(rewards.shape[0] - 1, -1, -1))
-        outputs = []
-        accumulated_reward = last_value
-        for t in timesteps:
-            inp = target[t]
-            discount_factor = discounts[t]
-            accumulated_reward = inp + discount_factor * td_lam * accumulated_reward
-            outputs.append(accumulated_reward)
-        returns = torch.flip(torch.stack(outputs), [0])
-        return returns
-        """
         next_values = torch.cat([values[1:], last_value.unsqueeze(0)],0)  
         targets = rewards + discounts * next_values * (1-td_lam)
         rets =[]
@@ -718,25 +539,6 @@ class DPI:
 
         returns = torch.flip(torch.stack(rets), [0])
         return returns
-        """
-    
-    def lambda_return(self,imged_reward, value_pred, bootstrap, discount=0.99, lambda_=0.95):
-        # Setting lambda=1 gives a discounted Monte Carlo return.
-        # Setting lambda=0 gives a fixed 1-step return.
-        next_values = torch.cat([value_pred[1:], bootstrap[None]], 0)
-        discount_tensor = discount * torch.ones_like(imged_reward)  # pcont
-        inputs = imged_reward + discount_tensor * next_values * (1 - lambda_)
-        last = bootstrap
-        indices = reversed(range(len(inputs)))
-        outputs = []
-        for index in indices:
-            inp, disc = inputs[index], discount_tensor[index]
-            last = inp + disc * lambda_ * last
-            outputs.append(last)
-        outputs = list(reversed(outputs))
-        outputs = torch.stack(outputs, 0)
-        returns = outputs
-        return returns
     
     def save_models(self, save_path):
         torch.save(
@@ -749,17 +551,6 @@ class DPI:
             '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()
     
@@ -769,7 +560,6 @@ if __name__ == '__main__':
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
     step = 0
-    total_steps = 2000000
+    total_steps = 10000
     dpi = DPI(args)
     dpi.train(step,total_steps)
-    dpi.evaluate()

DPI/utils.py

@@ -1,13 +1,10 @@
 import os
 import random
-import pickle
 import numpy as np
 from collections import deque
 
 import torch
 import torch.nn as nn
-from torch.utils.tensorboard import SummaryWriter
-
 
 import gym
 import dmc2gym
@@ -63,6 +60,17 @@ def make_dir(dir_path):
     return dir_path
 
 
+def preprocess_obs(obs, bits=5):
+    """Preprocessing image, see https://arxiv.org/abs/1807.03039."""
+    bins = 2**bits
+    assert obs.dtype == torch.float32
+    if bits < 8:
+        obs = torch.floor(obs / 2**(8 - bits))
+    obs = obs / bins
+    obs = obs + torch.rand_like(obs) / bins
+    obs = obs - 0.5
+    return obs
+
 
 class FrameStack(gym.Wrapper):
     def __init__(self, env, k):
@@ -136,90 +144,8 @@ class NormalizeActions:
         original = np.where(self._mask, original, action)
         return self._env.step(original)
 
-class TimeLimit:
-
-    def __init__(self, env, duration):
-        self._env = env
-        self._duration = duration
-        self._step = None
-
-    def __getattr__(self, name):
-        return getattr(self._env, name)
-
-    def step(self, action):
-        assert self._step is not None, 'Must reset environment.'
-        obs, reward, done, info = self._env.step(action)
-        self._step += 1
-        if self._step >= self._duration:
-          done = True
-          if 'discount' not in info:
-            info['discount'] = np.array(1.0).astype(np.float32)
-          self._step = None
-        return obs, reward, done, info
-
-    def reset(self):
-        self._step = 0
-        return self._env.reset()
-
 
 class ReplayBuffer:
-
-    def __init__(self, size, obs_shape, action_size, seq_len, batch_size, args):
-
-        self.size = size
-        self.obs_shape = obs_shape
-        self.action_size = action_size
-        self.seq_len = seq_len
-        self.batch_size = batch_size
-        self.idx = 0
-        self.full = False
-        self.observations = np.empty((size, *obs_shape), dtype=np.uint8) 
-        self.next_observations = np.empty((size, *obs_shape), dtype=np.uint8)
-        self.actions = np.empty((size, action_size), dtype=np.float32)
-        self.rewards = np.empty((size,), dtype=np.float32) 
-        self.terminals = np.empty((size,), dtype=np.float32)
-        self.steps, self.episodes = 0, 0
-        self.episode_count = np.zeros((size,), dtype=np.int32)
-    
-    def add(self, obs, ac, next_obs, rew, done, episode_count):
-        self.observations[self.idx] = obs
-        self.next_observations[self.idx] = next_obs
-        self.actions[self.idx] = ac
-        self.rewards[self.idx] = rew
-        self.terminals[self.idx] = done
-        self.full = self.full or self.idx == 0
-        self.steps += 1
-        self.episodes = self.episodes + (1 if done else 0)
-        self.episode_count[self.idx] = episode_count
-        self.idx = (self.idx + 1) % self.size
-
-    def _sample_idx(self, L):
-        valid_idx = False
-        while not valid_idx:
-            idx = np.random.randint(0, self.size if self.full else self.idx - L)
-            idxs = np.arange(idx, idx + L) % self.size
-            valid_idx = not self.idx in idxs[1:] 
-        return idxs
-
-    def _retrieve_batch(self, idxs, n, L):        
-        vec_idxs = idxs.transpose().reshape(-1)  # Unroll indices
-        observations = self.observations[vec_idxs]
-        next_obs = self.next_observations[vec_idxs]
-        obs = observations.reshape(L, n, *observations.shape[1:])
-        next_obs = next_obs.reshape(L, n, *next_obs.shape[1:])
-        acs = self.actions[vec_idxs].reshape(L, n, -1)
-        rew = self.rewards[vec_idxs].reshape(L, n)
-        term = self.terminals[vec_idxs].reshape(L, n)
-        return obs, acs, next_obs, rew, term
-
-    def sample(self):
-        n = self.batch_size
-        l = self.seq_len
-        obs,acs,next_obs,rews,terms= self._retrieve_batch(np.asarray([self._sample_idx(l) for _ in range(n)]), n, l)
-        return obs,acs,next_obs,rews,terms
-    
-
-class ReplayBuffer1:
     def __init__(self, size, obs_shape, action_size, seq_len, batch_size, args):
         self.size = size
         self.obs_shape = obs_shape
@@ -273,11 +199,7 @@ class ReplayBuffer1:
     def group_steps(self, buffer, variable, obs=True):
         variable = getattr(buffer, variable)
         non_zero_indices = np.nonzero(buffer.episode_count)[0]
-        print(buffer.episode_count)
         variable = variable[non_zero_indices]
-        print(variable.shape)
-        exit()
-
         if obs:
             variable = variable.reshape(-1, self.args.episode_length,
                                         self.args.frame_stack*self.args.channels,
@@ -292,9 +214,8 @@ class ReplayBuffer1:
                                     self.args.image_size,self.args.image_size)
         return variable
     
-    def sample_random_idx(self, buffer_length, last=False):
-        init = 0 if last else buffer_length - self.args.batch_size
-        random_indices = random.sample(range(init, buffer_length), self.args.batch_size)
+    def sample_random_idx(self, buffer_length):
+        random_indices = random.sample(range(0, buffer_length), self.args.batch_size) 
         return random_indices
 
     def group_and_sample_random_batch(self, buffer, variable_name, device, random_indices, is_obs=True, offset=0):
@@ -326,23 +247,19 @@ def make_env(args):
         )
     return env
 
-def shuffle_along_axis(a, axis):
-    idx = np.random.rand(*a.shape).argsort(axis=axis)
-    return np.take_along_axis(a,idx,axis=axis)
-
 def preprocess_obs(obs):
-    obs = (obs/255.0) - 0.5
+    obs = obs/255.0 - 0.5
     return obs
 
 def soft_update_params(net, target_net, tau):
     for param, target_param in zip(net.parameters(), target_net.parameters()):
         target_param.data.copy_(
-            tau * param.detach().data + (1 - tau) * target_param.data
+            tau * param.data + (1 - tau) * target_param.data
         )
 
 def save_image(array, filename):
     array = array.transpose(1, 2, 0)
-    array = ((array+0.5) * 255).astype(np.uint8)
+    array = (array * 255).astype(np.uint8)
     image = Image.fromarray(array)
     image.save(filename)
 
@@ -363,20 +280,6 @@ def video_from_array(arr, high_noise, filename):
                 out.write(frame)
             out.release()
 
-def save_video(images):
-    """
-    Image shape is (T, C, H, W)
-    Example:(50, 3, 84, 84)
-    """
-    output_file = "output.avi"
-    fourcc = cv2.VideoWriter_fourcc(*'XVID')
-    fps = 2
-    height, width, channels = 84,84,3
-    out = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
-    for image in images:
-        image = np.uint8(image.transpose((1, 2, 0)))
-        out.write(image)
-    out.release()
 
 class CorruptVideos:
     def __init__(self, dir_path):
@@ -450,51 +353,3 @@ class FreezeParameters:
     def __exit__(self, exc_type, exc_val, exc_tb):
         for i, param in enumerate(get_parameters(self.modules)):
             param.requires_grad = self.param_states[i]
-
-class Logger:
-
-    def __init__(self, log_dir, n_logged_samples=10, summary_writer=None):
-        self._log_dir = log_dir
-        print('########################')
-        print('logging outputs to ', log_dir)
-        print('########################')
-        self._n_logged_samples = n_logged_samples
-        self._summ_writer = SummaryWriter(log_dir, flush_secs=1, max_queue=1)
-
-    def log_scalar(self, scalar, name, step_):
-        self._summ_writer.add_scalar('{}'.format(name), scalar, step_)
-
-    def log_scalars(self, scalar_dict, step):
-        for key, value in scalar_dict.items():
-            print('{} : {}'.format(key, value))
-            self.log_scalar(value, key, step)
-        self.dump_scalars_to_pickle(scalar_dict, step)
-
-    def log_videos(self, videos, step, max_videos_to_save=1, fps=20, video_title='video'):
-
-        # max rollout length
-        max_videos_to_save = np.min([max_videos_to_save, videos.shape[0]])
-        max_length = videos[0].shape[0]
-        for i in range(max_videos_to_save):
-            if videos[i].shape[0]>max_length:
-                max_length = videos[i].shape[0]
-
-        # pad rollouts to all be same length
-        for i in range(max_videos_to_save):
-            if videos[i].shape[0]<max_length:
-                padding = np.tile([videos[i][-1]], (max_length-videos[i].shape[0],1,1,1))
-                videos[i] = np.concatenate([videos[i], padding], 0)
-
-            clip = mpy.ImageSequenceClip(list(videos[i]), fps=fps)
-            new_video_title = video_title+'{}_{}'.format(step, i) + '.gif'
-            filename = os.path.join(self._log_dir, new_video_title)
-            video.write_gif(filename, fps =fps)
-
-
-    def dump_scalars_to_pickle(self, metrics, step, log_title=None):
-        log_path = os.path.join(self._log_dir, "scalar_data.pkl" if log_title is None else log_title)
-        with open(log_path, 'ab') as f:
-            pickle.dump({'step': step, **dict(metrics)}, f)
-
-    def flush(self):
-        self._summ_writer.flush()