A major (performance) update on the submodule: srl-zoo

.gitignore

@@ -1,5 +1,6 @@
 # Custom
 data/
+figures/
 .idea/
 *.npy
 *.npz
@@ -10,7 +11,7 @@ dev/
 .ropeproject/
 .pytest_cache/
 _build/
-
+figures/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

evaluation/enjoy_latent.py

@@ -35,7 +35,7 @@ def getImage(srl_model, state, device):
         net_out = srl_model.decode(state)
         img = detachToNumpy(net_out)[0].T
 
-    img = deNormalize(img, mode="image_net")
+    img = deNormalize(img, mode="tf")
     return img[:, :, ::-1]
 
 
@@ -61,14 +61,18 @@ def main():
     parser = argparse.ArgumentParser(description="latent space enjoy")
     parser.add_argument('--log-dir', default='', type=str, help='directory to load model')
     parser.add_argument('--no-cuda', default=False, action="store_true")
-
+    parser.add_argument('--gpu-num', type=int, default=0, help='CUDA visible device (use CPU if -1, default: 0)')
+    parser.add_argument('--img-shape', type=str, default="(3,128,128)", help='image shape (default "(3,128,128)"')
     args = parser.parse_args()
-    use_cuda = not args.no_cuda
-    device = th.device("cuda" if th.cuda.is_available() and use_cuda else "cpu")
-
-    srl_model, exp_config = SRL4robotics.loadSavedModel(args.log_dir, VALID_MODELS, cuda=use_cuda)
+    # use_cuda = not args.no_cuda
+    # device = th.device("cuda" if th.cuda.is_available() and use_cuda else "cpu")
+    if args.img_shape is None:
+        img_shape = None #(3,224,224)
+    else:
+        img_shape = tuple(map(int, args.img_shape[1:-1].split(",")))
+    srl_model, exp_config = SRL4robotics.loadSavedModel(args.log_dir, VALID_MODELS, cuda=args.gpu_num, img_shape=img_shape)
     # Retrieve the pytorch model
-    srl_model = srl_model.model
+    srl_model = srl_model.module
     losses = exp_config['losses']
     state_dim = exp_config['state-dim']
 
@@ -83,7 +87,11 @@ def main():
 
     if len(loss_dims) == 0:
         print(losses)
-        loss_dims = {losses[0]: state_dim}
+        ## HACK: GAN is currently a model_type not a losses (TODO)
+        if len(losses) == 0:
+            loss_dims = {'gan': state_dim}
+        else:
+            loss_dims = {losses[0]: state_dim}
 
     # Load all the states and images
     data = json.load(open(args.log_dir + 'image_to_state.json'))

evaluation/predict_reward.py

@@ -15,6 +15,7 @@
 parser = argparse.ArgumentParser(description='Predict Reward from Ground Truth')
 parser.add_argument('--epochs', type=int, default=10,
                     help='number of epochs to train (default: 10)')
+parser.add_argument('--img-shape', type=tuple, default=None, help='image shape (default (3,224,224))')
 parser.add_argument('--seed', type=int, default=1, help='random seed (default: 1)')
 parser.add_argument('-bs', '--batch-size', type=int, default=32, help='batch_size (default: 256)')
 parser.add_argument('--training-set-size', type=int, default=-1,
@@ -80,8 +81,12 @@
 mean_val = np.mean(states, axis=0, keepdims=True)
 
 datasets = {'train': TensorDataset(X_train), 'val': TensorDataset(X_val)}
-train_loader = DataLoader(datasets['train'], batch_size=args.batch_size, shuffle=True)
-val_loader = DataLoader(datasets['val'], batch_size=args.batch_size, shuffle=False)
+if args.img_shape is None:
+    img_shape = None #(3,224,224)
+else:
+    img_shape = tuple(map(int, args.img_shape[1:-1].split(",")))
+train_loader = DataLoader(datasets['train'], img_shape=img_shape, batch_size=args.batch_size, shuffle=True)
+val_loader = DataLoader(datasets['val'], img_shape=img_shape, batch_size=args.batch_size, shuffle=False)
 dataloaders = {'train': train_loader, 'val': val_loader}
 
 start_time = time.time()

losses/losses.py

@@ -6,8 +6,13 @@
 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
+try:
+    # absolute import
+    from models.priors import ReverseLayerF
+except:
+    # relative import
+    from ..models.priors import ReverseLayerF
 
-from models.priors import ReverseLayerF
 from .utils import correlationMatrix
 
 try:
@@ -36,7 +41,7 @@ def addToLosses(self, name, weight, loss_value):
         """
         :param name: (str)
         :param weight: (float)
-        :param loss_value: (FloatTensor)
+        :param loss_value: (torch.FloatTensor)
         :return:
         """
         self.names.append(name)
@@ -60,7 +65,7 @@ def resetLosses(self):
 
 
 def roboticPriorsLoss(states, next_states, minibatch_idx,
-            dissimilar_pairs, same_actions_pairs, weight, loss_manager):
+                      dissimilar_pairs, same_actions_pairs, weight, loss_manager):
     """
     Computing the 4 Robotic priors: Temporal coherence, Causality, Proportionality, Repeatability
     :param states: (th.Tensor)
@@ -77,17 +82,17 @@ def roboticPriorsLoss(states, next_states, minibatch_idx,
 
     state_diff = next_states - states
     state_diff_norm = state_diff.norm(2, dim=1)
-    similarity = lambda x, y: th.exp(-(x - y).norm(2, dim=1) ** 2)
+    def similarity(x, y): return th.exp(-(x - y).norm(2, dim=1) ** 2)
     temp_coherence_loss = (state_diff_norm ** 2).mean()
     causality_loss = similarity(states[dissimilar_pairs[:, 0]],
                                 states[dissimilar_pairs[:, 1]]).mean()
     proportionality_loss = ((state_diff_norm[same_actions_pairs[:, 0]] -
                              state_diff_norm[same_actions_pairs[:, 1]]) ** 2).mean()
 
     repeatability_loss = (
-            similarity(states[same_actions_pairs[:, 0]], states[same_actions_pairs[:, 1]]) *
-            (state_diff[same_actions_pairs[:, 0]] - state_diff[same_actions_pairs[:, 1]]).norm(2,
-                                                                                               dim=1) ** 2).mean()
+        similarity(states[same_actions_pairs[:, 0]], states[same_actions_pairs[:, 1]]) *
+        (state_diff[same_actions_pairs[:, 0]] - state_diff[same_actions_pairs[:, 1]]).norm(2,
+                                                                                           dim=1) ** 2).mean()
     weights = [1, 1, 1, 1]
     names = ['temp_coherence_loss', 'causality_loss', 'proportionality_loss', 'repeatability_loss']
     losses = [temp_coherence_loss, causality_loss, proportionality_loss, repeatability_loss]
@@ -155,20 +160,41 @@ def l2Loss(params, weight, loss_manager):
     return weight * l2_loss
 
 
-def rewardModelLoss(rewards_pred, rewards_st, weight, loss_manager):
+def rewardModelLoss(rewards_pred, rewards_st, weight, loss_manager, label_weights, ignore_index=-1):
     """
     Categorical Reward prediction Loss (Cross-entropy)
     :param rewards_pred: predicted reward - categorical (th.Tensor)
     :param rewards_st: (th.Tensor)
     :param weight: coefficient to weight the loss
     :param loss_manager: loss criterion needed to log the loss value (LossManager)
+    :param label_weights (torch tensor) specifies the loss weight for each label. (need to be a torch tensor)
+    :param ignore_index (int, optional) Specifies a target value that is ignored and does not 
+            contribute to the input gradient. 
     :return:
     """
-    loss_fn = nn.CrossEntropyLoss()
+    loss_fn = nn.CrossEntropyLoss(weight=label_weights, ignore_index=ignore_index)
     reward_loss = loss_fn(rewards_pred, target=rewards_st)
     loss_manager.addToLosses('reward_loss', weight, reward_loss)
     return weight * reward_loss
 
+
+def spclsLoss(cls_pred, cls_gt, weight, loss_manager):
+    """
+    Categorical Reward prediction Loss (Cross-entropy)
+    :param rewards_pred: predicted reward - categorical (th.Tensor)
+    :param rewards_st: (th.Tensor)
+    :param weight: coefficient to weight the loss
+    :param loss_manager: loss criterion needed to log the loss value (LossManager)
+    :param label_weights (torch tensor) specifies the loss weight for each label. (need to be a torch tensor)
+    :param ignore_index (int, optional) Specifies a target value that is ignored and does not 
+            contribute to the input gradient. 
+    :return:
+    """
+    loss_fn = nn.CrossEntropyLoss()
+    cls_loss = loss_fn(cls_pred, target=cls_gt)
+    loss_manager.addToLosses('spcls_loss', weight, cls_loss)
+    return weight * cls_loss
+
 def reconstructionLoss(input_image, target_image):
     """
     Reconstruction Loss for Autoencoders
@@ -215,7 +241,7 @@ def generationLoss(decoded, next_decoded, obs, next_obs, weight, loss_manager):
 
 
 def perceptualSimilarityLoss(encoded_real, encoded_prediction, next_encoded_real, next_encoded_prediction,
-                            weight, loss_manager):
+                             weight, loss_manager):
     """
     Perceptual similarity Loss for VAE as in
     # "DARLA: Improving Zero-Shot Transfer in Reinforcement Learning", Higgins et al.
@@ -272,14 +298,14 @@ def mutualInformationLoss(states, rewards_st, weight, loss_manager):
     I = 0
     eps = 1e-10
     p_x = float(1 / np.sqrt(2 * np.pi)) * \
-          th.exp(-th.pow(th.norm((X - th.mean(X, dim=0)) / (th.std(X, dim=0) + eps), 2, dim=1), 2) / 2) + eps
+        th.exp(-th.pow(th.norm((X - th.mean(X, dim=0)) / (th.std(X, dim=0) + eps), 2, dim=1), 2) / 2) + eps
     p_y = float(1 / np.sqrt(2 * np.pi)) * \
-          th.exp(-th.pow(th.norm((Y - th.mean(Y, dim=0)) / (th.std(Y, dim=0) + eps), 2, dim=1), 2) / 2) + eps
+        th.exp(-th.pow(th.norm((Y - th.mean(Y, dim=0)) / (th.std(Y, dim=0) + eps), 2, dim=1), 2) / 2) + eps
     for x in range(X.shape[0]):
         for y in range(Y.shape[0]):
             p_xy = float(1 / np.sqrt(2 * np.pi)) * \
-                   th.exp(-th.pow(th.norm((th.cat([X[x], Y[y]]) - th.mean(th.cat([X, Y], dim=1), dim=0)) /
-                                          (th.std(th.cat([X, Y], dim=1), dim=0) + eps), 2), 2) / 2) + eps
+                th.exp(-th.pow(th.norm((th.cat([X[x], Y[y]]) - th.mean(th.cat([X, Y], dim=1), dim=0)) /
+                                       (th.std(th.cat([X, Y], dim=1), dim=0) + eps), 2), 2) / 2) + eps
             I += p_xy * th.log(p_xy / (p_x[x] * p_y[y]))
 
     mutual_info_loss = th.exp(-I)
@@ -374,3 +400,31 @@ def tripletLoss(states, p_states, n_states, weight, loss_manager, alpha=0.2):
     tcn_triplet_loss = tcn_triplet_loss.mean()
     loss_manager.addToLosses('triplet_loss', weight, tcn_triplet_loss)
     return weight * tcn_triplet_loss
+
+
+def ganNonSaturateLoss(img_rating, label, weight, loss_manager, name="non_saturate_loss"):
+    binary_crossentropy = th.nn.BCELoss()(img_rating, label)
+    loss_manager.addToLosses(name, weight, binary_crossentropy)
+    return weight * binary_crossentropy
+
+
+def ganBCEaccuracy(output, label=1):
+    """
+    label (int): 0 or 1
+    """
+    pred = output > 0.5
+    pred = pred.type(th.float)
+    if label:
+        acc = pred.sum() / pred.numel()
+    else:
+        acc = 1. - pred.sum() / pred.numel()
+    return acc
+
+
+def AEboundLoss(state_pred, weight, loss_manager, name='bonud_state_loss', max_val=50):
+    # state_pred of shape (batch_size, state_dim)
+    A = state_pred ** 2
+    norm_inf, _ = th.max(A, 1)
+    bound_loss = th.mean(th.relu(norm_inf-max_val))
+    loss_manager.addToLosses(name, weight, bound_loss)
+    return weight * bound_loss

losses/utils.py

@@ -1,9 +1,17 @@
-from pipeline import NO_PAIRS_ERROR
-from utils import printRed
+try:
+    # absolute import
+    from pipeline import NO_PAIRS_ERROR
+    from utils import printRed
+except:
+    # relative import
+    from ..pipeline import NO_PAIRS_ERROR
+    from ..utils import printRed
+
 
 import torch as th
 import numpy as np
 
+
 def overSampling(batch_size, m_list, pairs, function_on_pairs, actions, rewards):
     """
     Look for minibatches missing pairs of observations with the similar/dissimilar rewards (see params)

models/__init__.py

@@ -1,6 +1,6 @@
-from .models import CustomCNN
-from .modules import SRLModules, SRLModulesSplit
-from .priors import SRLConvolutionalNetwork, SRLDenseNetwork, SRLCustomCNN, Discriminator
-from .supervised import ConvolutionalNetwork, DenseNetwork, CustomCNN
+from .base_models import CustomCNN, UNet
+from .modules import SRLModules
+from .supervised import ConvolutionalNetwork, DenseNetwork
 from .autoencoders import LinearAutoEncoder, DenseAutoEncoder, CNNAutoEncoder
 from .vae import DenseVAE, CNNVAE
+