Add graphs

conda_env.yml

@@ -1,4 +1,4 @@
-name: pytorch_sac_ae
+name: pytorch_sac_ae2
 channels:
   - defaults
 dependencies:

encoder.py

@@ -1,9 +1,8 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 
-def tie_weights(src, trg):  
+def tie_weights(src, trg):
     assert type(src) == type(trg)
     trg.weight = src.weight
     trg.bias = src.bias
@@ -11,7 +10,7 @@ def tie_weights(src, trg):
 
 OUT_DIM = {2: 39, 4: 35, 6: 31}
 
-
+'''
 class PixelEncoder(nn.Module):
     """Convolutional encoder of pixels observations."""
     def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32):
@@ -29,8 +28,8 @@ class PixelEncoder(nn.Module):
             self.convs.append(nn.Conv2d(num_filters, num_filters, 3, stride=1))
 
         out_dim = OUT_DIM[num_layers]
-        self.fc = nn.Linear(num_filters * out_dim * out_dim, self.feature_dim * 2)
-        self.ln = nn.LayerNorm(self.feature_dim * 2)
+        self.fc = nn.Linear(num_filters * out_dim * out_dim, self.feature_dim)
+        self.ln = nn.LayerNorm(self.feature_dim)
 
         self.outputs = dict()
 
@@ -64,17 +63,97 @@ class PixelEncoder(nn.Module):
 
         h_norm = self.ln(h_fc)
         self.outputs['ln'] = h_norm
-        
-        h_tan = torch.tanh(h_norm)
 
-        mu, logstd = torch.chunk(h_tan, 2, dim=-1)
+        out = torch.tanh(h_norm)
+        self.outputs['tanh'] = out
+
+        return out
+
+    def copy_conv_weights_from(self, source):
+        """Tie convolutional layers"""
+        # only tie conv layers
+        for i in range(self.num_layers):
+            tie_weights(src=source.convs[i], trg=self.convs[i])
+
+    def log(self, L, step, log_freq):
+        if step % log_freq != 0:
+            return
+
+        for k, v in self.outputs.items():
+            L.log_histogram('train_encoder/%s_hist' % k, v, step)
+            if len(v.shape) > 2:
+                L.log_image('train_encoder/%s_img' % k, v[0], step)
+
+        for i in range(self.num_layers):
+            L.log_param('train_encoder/conv%s' % (i + 1), self.convs[i], step)
+        L.log_param('train_encoder/fc', self.fc, step)
+        L.log_param('train_encoder/ln', self.ln, step)
+'''
+
+class PixelEncoder(nn.Module):
+    """Convolutional encoder of pixels observations."""
+    def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32):
+        super().__init__()
+
+        assert len(obs_shape) == 3
+
+        self.feature_dim = feature_dim
+        self.num_layers = num_layers
+
+        self.convs = nn.ModuleList(
+            [nn.Conv2d(obs_shape[0], num_filters, 3, stride=2)]
+        )
+        for i in range(num_layers - 1):
+            self.convs.append(nn.Conv2d(num_filters, num_filters, 3, stride=1))
+
+        out_dim = OUT_DIM[num_layers]
+        self.fc = nn.Linear(num_filters * out_dim * out_dim, self.feature_dim * 2)
+        self.ln = nn.LayerNorm(self.feature_dim * 2)
+        self.combine = nn.Linear(self.feature_dim + 6, self.feature_dim)
+
+        self.outputs = dict()
+
+    def reparameterize(self, mu, logstd):
+        std = torch.exp(logstd)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward_conv(self, obs):
+        obs = obs / 255.
+        self.outputs['obs'] = obs
+
+        conv = torch.relu(self.convs[0](obs))
+        self.outputs['conv1'] = conv
+
+        for i in range(1, self.num_layers):
+            conv = torch.relu(self.convs[i](conv))
+            self.outputs['conv%s' % (i + 1)] = conv
+
+        h = conv.view(conv.size(0), -1)
+        return h
+
+    def forward(self, obs, detach=False):
+        h = self.forward_conv(obs)
+
+        if detach:
+            h = h.detach()
+
+        h_fc = self.fc(h)
+        self.outputs['fc'] = h_fc
+
+        h_norm = self.ln(h_fc)
+        self.outputs['ln'] = h_norm
+
+        #out = torch.tanh(h_norm)
+
+        mu, logstd = torch.chunk(h_norm, 2, dim=-1)
+        logstd = torch.tanh(logstd)
         self.outputs['mu'] = mu
         self.outputs['logstd'] = logstd
-
-        std = torch.tanh(h_norm)
-        self.outputs['std'] = std
+        self.outputs['std'] = logstd.exp()
 
         out = self.reparameterize(mu, logstd)
+        self.outputs['tanh'] = out
         return out, mu, logstd
 
     def copy_conv_weights_from(self, source):
@@ -97,7 +176,6 @@ class PixelEncoder(nn.Module):
         L.log_param('train_encoder/fc', self.fc, step)
         L.log_param('train_encoder/ln', self.ln, step)
 
-
 class IdentityEncoder(nn.Module):
     def __init__(self, obs_shape, feature_dim, num_layers, num_filters):
         super().__init__()
@@ -115,6 +193,27 @@ class IdentityEncoder(nn.Module):
         pass
 
 
+_AVAILABLE_ENCODERS = {'pixel': PixelEncoder, 'identity': IdentityEncoder}
+
+
+def make_encoder(
+    encoder_type, obs_shape, feature_dim, num_layers, num_filters
+):
+    assert encoder_type in _AVAILABLE_ENCODERS
+    return _AVAILABLE_ENCODERS[encoder_type](
+        obs_shape, feature_dim, num_layers, num_filters
+    )
+
+def club_loss(x_samples, x_mu, x_logvar, y_samples):        
+        sample_size = x_samples.shape[0]
+        random_index = torch.randperm(sample_size).long()
+        
+        positive = -(x_mu - y_samples)**2 / x_logvar.exp()
+        negative = - (x_mu - y_samples[random_index])**2 / x_logvar.exp()
+        upper_bound = (positive.sum(dim = -1) - negative.sum(dim = -1)).mean()
+        return upper_bound/2.
+
+
 class TransitionModel(nn.Module):
     def __init__(self, state_size, hidden_size, action_size, history_size):
         super().__init__()
@@ -126,13 +225,11 @@ class TransitionModel(nn.Module):
         self.act_fn = nn.ELU()
 
         self.fc_state_action = nn.Linear(state_size + action_size, hidden_size)
+        self.fc_hidden = nn.Linear(hidden_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
 
@@ -143,7 +240,6 @@ class TransitionModel(nn.Module):
             
     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):        
@@ -161,29 +257,21 @@ class TransitionModel(nn.Module):
         return dict(
             sample = torch.reshape(state[name], (state[name].shape[0]* state[name].shape[1], *state[name].shape[2:])))
 
-    def transition_step(self, prev_state, prev_action, prev_hist, prev_not_done):
-        prev_state = prev_state.detach() * prev_not_done
-        prev_hist = prev_hist * prev_not_done
+    def transition_step(self, state, action, hist, not_done):
+        state = state * not_done
+        hist = hist * not_done
 
-        state_action_enc = self.fc_state_action(torch.cat([prev_state, prev_action], dim=-1))
-        state_action_enc = self.act_fn(self.batch_norm(state_action_enc))
+        state_action_enc = self.act_fn(self.fc_state_action(torch.cat([state, action], dim=-1)))
+        state_action_enc = self.act_fn(self.fc_hidden(state_action_enc))
+        state_action_enc = self.act_fn(self.fc_hidden(state_action_enc))
+        state_action_enc = self.act_fn(self.fc_hidden(state_action_enc))
 
-        current_hist = self.history_cell(state_action_enc, prev_hist)
-        state_mu = self.act_fn(self.fc_state_mu(torch.cat([state_action_enc, prev_hist], dim=-1)))
-        state_sigma = F.softplus(self.fc_state_sigma(torch.cat([state_action_enc, prev_hist], dim=-1)))
-        sample_state = state_mu + torch.randn_like(state_mu) * state_sigma
+        current_hist = self.history_cell(state_action_enc, hist)
+        next_state_mu = self.act_fn(self.fc_state_mu(torch.cat([state_action_enc, hist], dim=-1)))
+        next_state_sigma = torch.tanh(self.fc_state_sigma(torch.cat([state_action_enc, hist], dim=-1)))
+        next_state = next_state_mu + torch.randn_like(next_state_mu) * next_state_sigma.exp()
 
-        state_enc = {"mean": state_mu, "std": state_sigma, "sample": sample_state, "history": current_hist}
-        return state_enc
-
-    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}
+        state_enc = {"mean": next_state_mu, "logvar": next_state_sigma, "sample": next_state, "history": current_hist}
         return state_enc
 
     def observe_rollout(self, rollout_states, rollout_actions, init_history, nonterms):
@@ -197,12 +285,14 @@ class TransitionModel(nn.Module):
             observed_rollout.append(state_enc)
         observed_rollout = self.stack_states(observed_rollout, dim=0)
         return observed_rollout
+    
+    def forward(self, state, action, hist, not_done):
+        return self.transition_step(state, action, hist, not_done)
 
-    def reparemeterize(self, mean, std):
+    def reparameterize(self, mean, std):
         eps = torch.randn_like(mean)
         return mean + eps * std
     
-
 def club_loss(x_samples, x_mu, x_logvar, y_samples):        
         sample_size = x_samples.shape[0]
         random_index = torch.randperm(sample_size).long()
@@ -210,15 +300,4 @@ def club_loss(x_samples, x_mu, x_logvar, y_samples):
         positive = -(x_mu - y_samples)**2 / x_logvar.exp()
         negative = - (x_mu - y_samples[random_index])**2 / x_logvar.exp()
         upper_bound = (positive.sum(dim = -1) - negative.sum(dim = -1)).mean()
-        return upper_bound/2.
-
-_AVAILABLE_ENCODERS = {'pixel': PixelEncoder, 'identity': IdentityEncoder}
-
-
-def make_encoder(
-    encoder_type, obs_shape, feature_dim, num_layers, num_filters
-):
-    assert encoder_type in _AVAILABLE_ENCODERS
-    return _AVAILABLE_ENCODERS[encoder_type](
-        obs_shape, feature_dim, num_layers, num_filters
-    )
+        return upper_bound/2.0

graphs_plot.py

@@ -0,0 +1,86 @@
+import os
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
+
+"""
+def tabulate_events(dpath):
+    files = os.listdir(dpath)[0]
+    summary_iterators = [EventAccumulator(os.path.join(dpath, files)).Reload()]
+
+    tags = summary_iterators[0].Tags()['scalars']
+
+    for it in summary_iterators:
+        assert it.Tags()['scalars'] == tags
+
+    out = {t: [] for t in tags}
+    steps = []
+
+    for tag in tags:
+        steps = [e.step for e in summary_iterators[0].Scalars(tag)]
+
+        for events in zip(*[acc.Scalars(tag) for acc in summary_iterators]):
+            assert len(set(e.step for e in events)) == 1
+
+            out[tag].append([e.value for e in events])
+    
+    return out, steps
+
+events, steps = tabulate_events('/home/vedant/pytorch_sac_ae/log/runs')
+
+data = []
+
+for tag, values in events.items():
+    for run_idx, run_values in enumerate(values):
+        for step_idx, value in enumerate(run_values):
+            data.append({
+                'tag': tag,
+                'run': run_idx,
+                'step': steps[step_idx],
+                'value': value,
+            })
+
+df = pd.DataFrame(data)
+print(df.head())
+exit()
+
+plt.figure(figsize=(10,6))
+sns.lineplot(data=df, x='step', y='value', hue='tag', ci='sd')
+plt.show()
+"""
+
+
+from tensorboard.backend.event_processing import event_accumulator
+
+
+
+def data_from_tb(files):
+    all_steps, all_rewards = [], []
+    for file in files:
+        ea = event_accumulator.EventAccumulator(file, size_guidance={'scalars': 0})
+        ea.Reload()
+
+        episode_rewards = ea.Scalars('train/episode_reward')
+        steps = [event.step for event in episode_rewards][:990000]
+        rewards = [event.value for event in episode_rewards][:990000]
+        all_steps.append(steps)
+        all_rewards.append(rewards)
+    return all_steps, all_rewards
+
+
+files = ['/home/vedant/pytorch_sac_ae/log/runs/tb_21_05_2023-13_19_36/events.out.tfevents.1684667976.cpswkstn6-nvidia4090.1749060.0',
+         '/home/vedant/pytorch_sac_ae/log/runs/tb_22_05_2023-09_56_30/events.out.tfevents.1684742190.cpswkstn6-nvidia4090.1976229.0']
+
+all_steps, all_rewards = data_from_tb(files)
+mean_rewards = np.mean(all_rewards, axis=0)
+std_rewards = np.std(all_rewards, axis=0)
+mean_steps = np.mean(all_steps, axis=0) 
+
+df = pd.DataFrame({'Steps': mean_steps,'Rewards': mean_rewards,'Standard Deviation': std_rewards})
+
+sns.relplot(x='Steps', y='Rewards', kind='line', data=df, ci="sd")
+plt.fill_between(df['Steps'], df['Rewards'] - df['Standard Deviation'], df['Rewards'] + df['Standard Deviation'], color='b', alpha=.1)
+plt.title("Mean Rewards vs Steps with Standard Deviation")
+plt.show()

logger.py

@@ -163,4 +163,4 @@ class Logger(object):
 
     def dump(self, step):
         self._train_mg.dump(step, 'train')
-        self._eval_mg.dump(step, 'eval')
+        self._eval_mg.dump(step, 'eval')

sac_ae.py

@@ -70,8 +70,10 @@ class Actor(nn.Module):
         self.outputs = dict()
         self.apply(weight_init)
 
-    def forward(self, obs, compute_pi=True, compute_log_pi=True, detach_encoder=False):
-        obs, _, _ = self.encoder(obs, detach=detach_encoder)
+    def forward(
+        self, obs, compute_pi=True, compute_log_pi=True, detach_encoder=False
+    ):
+        obs,_,_ = self.encoder(obs, detach=detach_encoder)
 
         mu, log_std = self.trunk(obs).chunk(2, dim=-1)
 
@@ -98,6 +100,7 @@ class Actor(nn.Module):
             log_pi = None
 
         mu, pi, log_pi = squash(mu, pi, log_pi)
+
         return mu, pi, log_pi, log_std
 
     def log(self, L, step, log_freq=LOG_FREQ):
@@ -156,7 +159,7 @@ class Critic(nn.Module):
 
     def forward(self, obs, action, detach_encoder=False):
         # detach_encoder allows to stop gradient propogation to encoder
-        obs, _ , _ = self.encoder(obs, detach=detach_encoder)
+        obs,_,_ = self.encoder(obs, detach=detach_encoder)
 
         q1 = self.Q1(obs, action)
         q2 = self.Q2(obs, action)
@@ -179,39 +182,12 @@ class Critic(nn.Module):
             L.log_param('train_critic/q1_fc%d' % i, self.Q1.trunk[i * 2], step)
             L.log_param('train_critic/q2_fc%d' % i, self.Q2.trunk[i * 2], step)
 
-class CURL(nn.Module):
-    """
-    CURL
-    """
-
-    def __init__(self, obs_shape, z_dim, a_dim, batch_size, critic, critic_target, output_type="continuous"):
-        super(CURL, self).__init__()
-        self.batch_size = batch_size
-
-        self.encoder = critic.encoder
-
-        self.encoder_target = critic_target.encoder 
+class LBLoss(nn.Module):
+    def __init__(self, z_dim):
+        super(LBLoss, self).__init__()
+        self.z_dim = z_dim
 
         self.W = nn.Parameter(torch.rand(z_dim, z_dim))
-        self.combine = nn.Linear(z_dim + a_dim, z_dim)
-        self.output_type = output_type
-
-    def encode(self, x, a=None, detach=False, ema=False):
-        """
-        Encoder: z_t = e(x_t)
-        :param x: x_t, x y coordinates
-        :return: z_t, value in r2
-        """
-        if ema:
-            with torch.no_grad():
-                z_out = self.encoder_target(x)[0]
-                z_out = self.combine(torch.concat((z_out,a), dim=-1))
-        else:
-            z_out = self.encoder(x)[0]
-
-        if detach:
-            z_out = z_out.detach()
-        return z_out
 
     def compute_logits(self, z_a, z_pos):
         """
@@ -225,7 +201,8 @@ class CURL(nn.Module):
         logits = torch.matmul(z_a, Wz)  # (B,B)
         logits = logits - torch.max(logits, 1)[0][:, None]
         return logits
-        
+
+
 class SacAeAgent(object):
     """SAC+AE algorithm."""
     def __init__(
@@ -267,12 +244,6 @@ class SacAeAgent(object):
         self.critic_target_update_freq = critic_target_update_freq
         self.decoder_update_freq = decoder_update_freq
         self.decoder_latent_lambda = decoder_latent_lambda
-        
-        self.transition_model = TransitionModel(
-            encoder_feature_dim, 
-            hidden_dim, 
-            action_shape[0], 
-            encoder_feature_dim).to(device)
 
         self.actor = Actor(
             obs_shape, action_shape, hidden_dim, encoder_type,
@@ -289,7 +260,13 @@ class SacAeAgent(object):
             obs_shape, action_shape, hidden_dim, encoder_type,
             encoder_feature_dim, num_layers, num_filters
         ).to(device)
+        
+        self.transition_model = TransitionModel(
+            encoder_feature_dim, hidden_dim,  action_shape[0], history_size=256
+        ).to(device)
 
+        self.lb_loss = LBLoss(encoder_feature_dim).to(device)
+        
         self.critic_target.load_state_dict(self.critic.state_dict())
 
         # tie encoders between actor and critic
@@ -300,11 +277,6 @@ class SacAeAgent(object):
         # set target entropy to -|A|
         self.target_entropy = -np.prod(action_shape)
 
-        self.CURL = CURL(obs_shape, encoder_feature_dim, action_shape[0],
-                         obs_shape[0], self.critic,self.critic_target, output_type='continuous').to(self.device)
-        
-        self.cross_entropy_loss = nn.CrossEntropyLoss()
-
         self.decoder = None
         if decoder_type != 'identity':
             # create decoder
@@ -316,7 +288,10 @@ class SacAeAgent(object):
 
             # optimizer for critic encoder for reconstruction loss
             self.encoder_optimizer = torch.optim.Adam(
-                self.critic.encoder.parameters(), lr=encoder_lr
+                list(self.critic.encoder.parameters()) +
+                list(self.transition_model.parameters()), #+
+                #list(self.lb_loss.parameters()), 
+                lr=encoder_lr
             )
 
             # optimizer for decoder
@@ -335,10 +310,6 @@ class SacAeAgent(object):
             self.critic.parameters(), lr=critic_lr, betas=(critic_beta, 0.999)
         )
 
-        self.cpc_optimizer = torch.optim.Adam(
-                self.CURL.parameters(), lr=encoder_lr
-        )
-        
         self.log_alpha_optimizer = torch.optim.Adam(
             [self.log_alpha], lr=alpha_lr, betas=(alpha_beta, 0.999)
         )
@@ -387,6 +358,7 @@ class SacAeAgent(object):
                                  target_Q) + F.mse_loss(current_Q2, target_Q)
         L.log('train_critic/loss', critic_loss, step)
 
+
         # Optimize the critic
         self.critic_optimizer.zero_grad()
         critic_loss.backward()
@@ -423,74 +395,76 @@ class SacAeAgent(object):
         alpha_loss.backward()
         self.log_alpha_optimizer.step()
 
-    def update_decoder(self, last_obs, last_action, last_reward, curr_obs, last_not_done, action, reward, next_obs, not_done, target_obs, L, step):
-        h_curr, mu_h_curr, std_h_curr = self.critic.encoder(curr_obs)
-
+    def update_decoder(self, obs, target_obs, L, step, obs_list, action_list, reward_list, next_obs_list, not_done_list):
         with torch.no_grad():
-            h_last, _, _ = self.critic.encoder(last_obs)
-            self.transition_model.init_states(last_obs.shape[0], self.device)
-            curr_state = self.transition_model.transition_step(h_last, last_action, self.transition_model.prev_history, last_not_done)
-
-            hist = curr_state["history"]
-            next_state = self.transition_model.transition_step(h_curr, action, hist, not_done)
-
-            next_state_mu = next_state["mean"]
-            next_state_sigma = next_state["std"]
-            next_state_sample = next_state["sample"]
-            pred_dist = torch.distributions.Normal(next_state_mu, next_state_sigma)
+            hist = torch.zeros((target_obs.shape[0], 256)).to(self.device)
+            for i in range(len(obs_list)-1):
+                state, _, _ = self.critic.encoder(obs_list[i])
+                action = action_list[i]
+                not_done = not_done_list[i]
+                state_enc = self.transition_model(state, action, hist, not_done)
+                hist = state_enc["history"]
         
-        h, mu_h_next, logstd_h_next = self.critic.encoder(next_obs)
-        std_h_next = torch.exp(logstd_h_next)
-        enc_dist = torch.distributions.Normal(mu_h_next, std_h_next)
-        enc_loss = torch.mean(torch.distributions.kl.kl_divergence(enc_dist,pred_dist)) * 0.1
+        h, h_mu, h_logvar = self.critic.encoder(obs_list[-1])
+        h_clone = h.clone() 
+        
+        action = action_list[-1]
+        not_done = not_done_list[-1]
+        state_enc = self.transition_model(h, action, hist, not_done)
+        mean, std = state_enc["mean"], state_enc["logvar"].exp()
 
-        z_pos = self.CURL.encode(next_obs, action.detach(), ema=True)
-        logits = self.CURL.compute_logits(h_curr, z_pos)
+        h_dist_enc = torch.distributions.Normal(h_mu, h_logvar.exp())
+        h_dist_pred = torch.distributions.Normal(mean, std)
+        enc_loss = torch.distributions.kl.kl_divergence(h_dist_enc, h_dist_pred).mean() * 1e-2
+        
+        with torch.no_grad():
+            z_pos, _ , _ = self.critic_target.encoder(next_obs_list[-1])
+            z_out = self.critic_target.encoder.combine(torch.concat((z_pos, action), dim=-1))
+        logits = self.lb_loss.compute_logits(h, z_out)
         labels = torch.arange(logits.shape[0]).long().to(self.device)
-        lb_loss = self.cross_entropy_loss(logits, labels) * 0.1
-
-        ub_loss = club_loss(h, mu_h_next, logstd_h_next, next_state_sample) * 0.1
-    
+        lb_loss = nn.CrossEntropyLoss()(logits, labels) * 1e-2
+        
+        #with torch.no_grad():
+        #    z_pos, _ , _ = self.critic.encoder(next_obs_list[-1])
+        #ub_loss = club_loss(state_enc["sample"], mean, state_enc["logvar"], h) * 1e-1
+        
         if target_obs.dim() == 4:
             # preprocess images to be in [-0.5, 0.5] range
             target_obs = utils.preprocess_obs(target_obs)
-
-        rec_obs = self.decoder(h)
+        rec_obs = self.decoder(h_clone)
         rec_loss = F.mse_loss(target_obs, rec_obs)
 
-        # add L2 penalty on latent representation
-        # see https://arxiv.org/pdf/1903.12436.pdf
-        latent_loss = (0.5 * h.pow(2).sum(1)).mean()
-
-        loss = rec_loss + enc_loss + lb_loss + ub_loss #self.decoder_latent_lambda * latent_loss
+        ub_loss = torch.tensor(0.0)
+        #enc_loss = torch.tensor(0.0)
+        #lb_loss = torch.tensor(0.0)
+        #rec_loss = torch.tensor(0.0)
+        loss = rec_loss + enc_loss + lb_loss + ub_loss
         self.encoder_optimizer.zero_grad()
         self.decoder_optimizer.zero_grad()
-        self.cpc_optimizer.zero_grad()
         loss.backward()
-        
-        self.encoder_optimizer.step()   
+
+        self.encoder_optimizer.step()
         self.decoder_optimizer.step()
-        self.cpc_optimizer.step()
+
+        #enc_loss = torch.tensor(0.0)
         L.log('train_ae/ae_loss', loss, step)
+        L.log('train_ae/rec_loss', rec_loss, step)
+        L.log('train_ae/enc_loss', enc_loss, step)
         L.log('train_ae/lb_loss', lb_loss, step)
         L.log('train_ae/ub_loss', ub_loss, step)
-        L.log('train_ae/enc_loss', enc_loss, step)
-        L.log('train_ae/dec_loss', rec_loss, step)
 
         self.decoder.log(L, step, log_freq=LOG_FREQ)
 
     def update(self, replay_buffer, L, step):
-        last_obs, last_action, last_reward, curr_obs, last_not_done, action, reward, next_obs, not_done = replay_buffer.sample()
-        #obs, action, reward, next_obs, not_done = replay_buffer.sample()
+        obs_list, action_list, reward_list, next_obs_list, not_done_list = replay_buffer.sample()
+        obs, action, reward, next_obs, not_done = obs_list[-1], action_list[-1], reward_list[-1], next_obs_list[-1], not_done_list[-1]  
 
-        L.log('train/batch_reward', last_reward.mean(), step)
+        L.log('train/batch_reward', reward.mean(), step)
 
-        #self.update_critic(last_obs, last_action, last_reward, curr_obs, last_not_done, L, step)
-        self.update_critic(curr_obs, action, reward, next_obs, not_done, L, step)
+        self.update_critic(obs, action, reward, next_obs, not_done, L, step)
 
         if step % self.actor_update_freq == 0:
-            #self.update_actor_and_alpha(last_obs, L, step)
-            self.update_actor_and_alpha(curr_obs, L, step)
+            self.update_actor_and_alpha(obs, L, step)
 
         if step % self.critic_target_update_freq == 0:
             utils.soft_update_params(
@@ -505,7 +479,7 @@ class SacAeAgent(object):
             )
 
         if self.decoder is not None and step % self.decoder_update_freq == 0:
-            self.update_decoder(last_obs, last_action, last_reward, curr_obs, last_not_done, action, reward, next_obs, not_done, next_obs, L, step)
+            self.update_decoder(obs, obs, L, step, obs_list, action_list, reward_list, next_obs_list, not_done_list)
 
     def save(self, model_dir, step):
         torch.save(

train.py

@@ -28,36 +28,37 @@ def parse_args():
     parser.add_argument('--frame_stack', default=3, type=int)
     parser.add_argument('--img_source', default=None, type=str, choices=['color', 'noise', 'images', 'video', 'none'])
     parser.add_argument('--resource_files', type=str)
+    parser.add_argument('--resource_files_test', type=str)
     parser.add_argument('--total_frames', default=10000, type=int)
     # replay buffer
-    parser.add_argument('--replay_buffer_capacity', default=1000000, type=int)
+    parser.add_argument('--replay_buffer_capacity', default=100000, type=int)
     # train
     parser.add_argument('--agent', default='sac_ae', type=str)
     parser.add_argument('--init_steps', default=1000, type=int)
-    parser.add_argument('--num_train_steps', default=1000000, type=int)
-    parser.add_argument('--batch_size', default=512, type=int)
+    parser.add_argument('--num_train_steps', default=2000000, type=int)
+    parser.add_argument('--batch_size', default=32, type=int)
     parser.add_argument('--hidden_dim', default=1024, type=int)
     # eval
     parser.add_argument('--eval_freq', default=10000, type=int)
     parser.add_argument('--num_eval_episodes', default=10, type=int)
     # critic
-    parser.add_argument('--critic_lr', default=1e-3, type=float)
+    parser.add_argument('--critic_lr', default=1e-4, type=float)
     parser.add_argument('--critic_beta', default=0.9, type=float)
     parser.add_argument('--critic_tau', default=0.01, type=float)
     parser.add_argument('--critic_target_update_freq', default=2, type=int)
     # actor
-    parser.add_argument('--actor_lr', default=1e-3, type=float)
+    parser.add_argument('--actor_lr', default=1e-4, 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)
     # encoder/decoder
     parser.add_argument('--encoder_type', default='pixel', type=str)
-    parser.add_argument('--encoder_feature_dim', default=50, type=int)
-    parser.add_argument('--encoder_lr', default=1e-3, type=float)
+    parser.add_argument('--encoder_feature_dim', default=250, type=int)
+    parser.add_argument('--encoder_lr', default=1e-4, type=float)
     parser.add_argument('--encoder_tau', default=0.05, type=float)
     parser.add_argument('--decoder_type', default='pixel', type=str)
-    parser.add_argument('--decoder_lr', default=1e-3, type=float)
+    parser.add_argument('--decoder_lr', default=1e-4, type=float)
     parser.add_argument('--decoder_update_freq', default=1, type=int)
     parser.add_argument('--decoder_latent_lambda', default=1e-6, type=float)
     parser.add_argument('--decoder_weight_lambda', default=1e-7, type=float)
@@ -153,9 +154,25 @@ def main():
     )
     env.seed(args.seed)
 
+    env_test = dmc2gym.make(
+        domain_name=args.domain_name,
+        task_name=args.task_name,
+        seed=args.seed,
+        visualize_reward=False,
+        from_pixels=(args.encoder_type == 'pixel'),
+        height=args.image_size,
+        width=args.image_size,
+        frame_skip=args.action_repeat,
+        img_source=args.img_source,
+        resource_files=args.resource_files_test,
+        total_frames=args.total_frames
+    )
+    env_test.seed(args.seed)
+
     # stack several consecutive frames together
     if args.encoder_type == 'pixel':
         env = utils.FrameStack(env, k=args.frame_stack)
+        env_test = utils.FrameStack(env_test, k=args.frame_stack)
 
     utils.make_dir(args.work_dir)
     video_dir = utils.make_dir(os.path.join(args.work_dir, 'video'))
@@ -202,7 +219,7 @@ def main():
             # evaluate agent periodically
             if step % args.eval_freq == 0:
                 L.log('eval/episode', episode, step)
-                evaluate(env, agent, video, args.num_eval_episodes, L, step)
+                evaluate(env_test, agent, video, args.num_eval_episodes, L, step)
                 if args.save_model:
                     agent.save(model_dir, step)
                 if args.save_buffer:
@@ -218,65 +235,28 @@ def main():
 
             L.log('train/episode', episode, step)
 
-        if episode_step == 0:
-            last_obs = obs
-            # sample action for data collection
-            if step < args.init_steps:
-                last_action = env.action_space.sample()
-            else:
-                with utils.eval_mode(agent):
-                    last_action = agent.sample_action(last_obs)
-            
-            curr_obs, last_reward, last_done, _ = env.step(last_action)
-
-            # allow infinit bootstrap
-            last_done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(last_done)
-            episode_reward += last_reward
-
-            # sample action for data collection
-            if step < args.init_steps:
-                action = env.action_space.sample()
-            else:
-                with utils.eval_mode(agent):
-                    action = agent.sample_action(curr_obs)
-            
-            next_obs, reward, done, _ = env.step(action)
-
-            # allow infinit bootstrap
-            done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(done)
-            episode_reward += reward
-
-            replay_buffer.add(last_obs, last_action, last_reward, curr_obs, last_done_bool, action, reward, next_obs, done_bool)
-
-        last_obs = curr_obs
-        last_action = action
-        last_reward = reward
-        last_done = done
-        curr_obs = next_obs
-
         # sample action for data collection
         if step < args.init_steps:
             action = env.action_space.sample()
         else:
             with utils.eval_mode(agent):
-                action = agent.sample_action(curr_obs)
+                action = agent.sample_action(obs)
 
-        
         # run training update
         if step >= args.init_steps:
-            #num_updates = args.init_steps if step == args.init_steps else 1
-            num_updates = 1 if step == args.init_steps else 1
+            num_updates = args.init_steps if step == args.init_steps else 1
             for _ in range(num_updates):
                 agent.update(replay_buffer, L, step)
 
         next_obs, reward, done, _ = env.step(action)
 
         # allow infinit bootstrap
-        done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(done)
+        done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(
+            done
+        )
         episode_reward += reward
 
-        #replay_buffer.add(obs, action, reward, next_obs, done_bool)
-        replay_buffer.add(last_obs, last_action, last_reward, curr_obs, last_done_bool, action, reward, next_obs, done_bool)
+        replay_buffer.add(obs, action, reward, next_obs, done_bool)
 
         obs = next_obs
         episode_step += 1

utils.py

@@ -75,26 +75,18 @@ class ReplayBuffer(object):
         # the proprioceptive obs is stored as float32, pixels obs as uint8
         obs_dtype = np.float32 if len(obs_shape) == 1 else np.uint8
 
-        self.last_obses = np.empty((capacity, *obs_shape), dtype=obs_dtype)
-        self.curr_obses = np.empty((capacity, *obs_shape), dtype=obs_dtype)
+        self.obses = np.empty((capacity, *obs_shape), dtype=obs_dtype)
         self.next_obses = np.empty((capacity, *obs_shape), dtype=obs_dtype)
-        self.last_actions = np.empty((capacity, *action_shape), dtype=np.float32)
         self.actions = np.empty((capacity, *action_shape), dtype=np.float32)
-        self.last_rewards = np.empty((capacity, 1), dtype=np.float32)
         self.rewards = np.empty((capacity, 1), dtype=np.float32)
-        self.last_not_dones = np.empty((capacity, 1), dtype=np.float32)
         self.not_dones = np.empty((capacity, 1), dtype=np.float32)
 
         self.idx = 0
         self.last_save = 0
         self.full = False
 
-    def add(self, last_obs, last_action, last_reward, curr_obs, last_done, action, reward, next_obs, done):
-        np.copyto(self.last_obses[self.idx], last_obs)
-        np.copyto(self.last_actions[self.idx], last_action)
-        np.copyto(self.last_rewards[self.idx], last_reward)
-        np.copyto(self.curr_obses[self.idx], curr_obs)
-        np.copyto(self.last_not_dones[self.idx], not last_done)
+    def add(self, obs, action, reward, next_obs, done):
+        np.copyto(self.obses[self.idx], obs)
         np.copyto(self.actions[self.idx], action)
         np.copyto(self.rewards[self.idx], reward)
         np.copyto(self.next_obses[self.idx], next_obs)
@@ -104,35 +96,29 @@ class ReplayBuffer(object):
         self.full = self.full or self.idx == 0
 
     def sample(self):
+        begin = 2
         idxs = np.random.randint(
-            0, self.capacity if self.full else self.idx, size=self.batch_size
+            begin, self.capacity if self.full else self.idx, size=self.batch_size
         )
+        past_idxs = idxs - begin
 
-        last_obses = torch.as_tensor(self.last_obses[idxs], device=self.device).float()
-        last_actions = torch.as_tensor(self.last_actions[idxs], device=self.device)
-        last_rewards = torch.as_tensor(self.last_rewards[idxs], device=self.device)
-        curr_obses = torch.as_tensor(self.curr_obses[idxs], device=self.device).float()
-        last_not_dones = torch.as_tensor(self.last_not_dones[idxs], device=self.device)
-        actions = torch.as_tensor(self.actions[idxs], device=self.device)
-        rewards = torch.as_tensor(self.rewards[idxs], device=self.device)
-        next_obses = torch.as_tensor(self.next_obses[idxs], device=self.device).float()
-        not_dones = torch.as_tensor(self.not_dones[idxs], device=self.device)
+        obses = torch.as_tensor(np.swapaxes(np.asarray([self.obses[past_idxs:idxs] for past_idxs, idxs in zip(past_idxs, idxs)]),0,1), device=self.device).float()
+        actions = torch.as_tensor(np.swapaxes(np.asarray([self.actions[past_idxs:idxs] for past_idxs, idxs in zip(past_idxs, idxs)]),0,1), device=self.device)
+        rewards = torch.as_tensor(np.swapaxes(np.asarray([self.rewards[past_idxs:idxs] for past_idxs, idxs in zip(past_idxs, idxs)]),0,1), device=self.device)
+        next_obses = torch.as_tensor(np.swapaxes(np.asarray([self.next_obses[past_idxs:idxs] for past_idxs, idxs in zip(past_idxs, idxs)]),0,1), device=self.device).float()
+        not_dones = torch.as_tensor(np.swapaxes(np.asarray([self.not_dones[past_idxs:idxs] for past_idxs, idxs in zip(past_idxs, idxs)]),0,1), device=self.device)
 
-        return last_obses, last_actions, last_rewards, curr_obses, last_not_dones, actions, rewards, next_obses, not_dones
+        return obses, actions, rewards, next_obses, not_dones
 
     def save(self, save_dir):
         if self.idx == self.last_save:
             return
         path = os.path.join(save_dir, '%d_%d.pt' % (self.last_save, self.idx))
         payload = [
-            self.last_obses[self.last_save:self.idx],
-            self.last_actions[self.last_save:self.idx],
-            self.last_rewards[self.last_save:self.idx],
-            self.curr_obses[self.last_save:self.idx],
-            self.last_not_dones[self.last_save:self.idx],
+            self.obses[self.last_save:self.idx],
+            self.next_obses[self.last_save:self.idx],
             self.actions[self.last_save:self.idx],
             self.rewards[self.last_save:self.idx],
-            self.next_obses[self.last_save:self.idx],
             self.not_dones[self.last_save:self.idx]
         ]
         self.last_save = self.idx
@@ -146,14 +132,10 @@ class ReplayBuffer(object):
             path = os.path.join(save_dir, chunk)
             payload = torch.load(path)
             assert self.idx == start
-            self.last_obses[start:end] = payload[0]
-            self.last_actions[start:end] = payload[1]
-            self.last_rewards[start:end] = payload[2]
-            self.curr_obses[start:end] = payload[3]
-            self.last_not_dones[start:end] = payload[4]
+            self.obses[start:end] = payload[0]
+            self.next_obses[start:end] = payload[1]
             self.actions[start:end] = payload[2]
             self.rewards[start:end] = payload[3]
-            self.next_obses[start:end] = payload[4]
             self.not_dones[start:end] = payload[4]
             self.idx = end