dataset and agent initialization generalized

contrastive_learning/configs/pli_train.yaml

@@ -15,14 +15,16 @@ experiment: pli
 seed: 42
 device: cuda
 
+agent_type: pli # Used for agent_inits
+dataset_type: state
 train_epochs: 1000
 save_frequency: 10 # Frequency to save the model - there will be a test in each step
 train_dset_split: 0.8
 
 # Hyperparameters to be used everywhere
-batch_size: 64
+batch_size: 128
 lr: 1e-2
-weight_decay: 1e-5
+weight_decay: 1e-6
 action_dim: 2
 pos_dim: 8 # Number of numbers for one detected marker (4 * 2)
 hidden_dim: 64

contrastive_learning/configs/train.yaml

@@ -1,25 +1,37 @@
 defaults:
-    - _self_
-    - agent: cpn # forward model will also be added in the future - TODO: maybe add a wrapper agent class?
-    - encoder: resnet18
-    - trans: mlp
+    # - _self_
+    # - agent: cpn # forward model will also be added in the future - TODO: maybe add a wrapper agent class?
+    # - encoder: resnet18
+    # - trans: mlp
+    # - optimizer: adam
+    - agent: pli
     - optimizer: adam
 
+model:
+  _target_: contrastive_learning.models.custom_models.LinearInverse
+  input_dim: ??? # They will be set according to the rest of the config
+  action_dim: ???
+  hidden_dim: ???
+
 # this needs to be specified manually
-experiment: resnet18_mlp_pretrained_false
+experiment: pli
 
 seed: 42
 device: cuda
 
+agent_type: pli # cpn / pli
+dataset_type: state # State or visual for now
 train_epochs: 1000
 save_frequency: 10 # Frequency to save the model - there will be a test in each step
 train_dset_split: 0.8
 
 # Hyperparameters to be used everywhere
-batch_size: 16
-lr: 5e-4
+batch_size: 64
+lr: 1e-2
 weight_decay: 1e-5
-z_dim: 1000 # resnet18 gives out 1000 dimensions
+z_dim: 1000 # resnet18 gives out 1000 dimensions (needed for cpn)
+pos_dim: 8 # this is needed for pli
+hidden_dim: 64
 action_dim: 2
 
 distributed: true

contrastive_learning/datasets/dataloaders.py

@@ -0,0 +1,32 @@
+import torch
+import torch.utils.data as data 
+from omegaconf import DictConfig, OmegaConf
+
+from contrastive_learning.datasets.state_dataset import StateDataset
+from contrastive_learning.datasets.visual_dataset import VisualDataset
+
+# Script to return dataloaders
+
+def get_dataloaders(cfg : DictConfig):
+    # Load dataset - splitting will be done with random splitter
+    if cfg.dataset_type == 'state':
+        dataset = StateDataset(data_dir=cfg.data_dir)
+    else:
+        dataset = VisualDataset(data_dir=cfg.data_dir, frame_interval=cfg.frame_interval, video_type=cfg.video_type)
+
+    train_dset_size = int(len(dataset) * cfg.train_dset_split)
+    test_dset_size = len(dataset) - train_dset_size
+
+    # Random split the train and validation datasets
+    train_dset, test_dset = data.random_split(dataset, 
+                                             [train_dset_size, test_dset_size],
+                                             generator=torch.Generator().manual_seed(cfg.seed))
+    train_sampler = data.DistributedSampler(train_dset, drop_last=True, shuffle=True) if cfg.distributed else None
+    test_sampler = data.DistributedSampler(test_dset, drop_last=True, shuffle=False) if cfg.distributed else None # val will not be shuffled
+
+    train_loader = data.DataLoader(train_dset, batch_size=cfg.batch_size, shuffle=train_sampler is None,
+                                    num_workers=cfg.num_workers, sampler=train_sampler)
+    test_loader = data.DataLoader(test_dset, batch_size=cfg.batch_size, shuffle=test_sampler is None,
+                                    num_workers=cfg.num_workers, sampler=test_sampler)
+
+    return train_loader, test_loader, dataset

contrastive_learning/datasets/pli_dataset.py → contrastive_learning/datasets/state_dataset.py

@@ -13,7 +13,7 @@
 from contrastive_learning.datasets.preprocess import smoothen_corners, dump_pos_corners
 # import contrastive_learning.utils.utils as utils
 
-class Dataset:
+class StateDataset:
     def __init__(self, data_dir: str) -> None:
 
         # Get the roots
@@ -76,7 +76,7 @@ def denormalize_action(self, action):
 
 def get_dataloaders(cfg : DictConfig):
     # Load dataset - splitting will be done with random splitter
-    dataset = Dataset(data_dir=cfg.data_dir)
+    dataset = StateDataset(data_dir=cfg.data_dir)
 
     train_dset_size = int(len(dataset) * cfg.train_dset_split)
     test_dset_size = len(dataset) - train_dset_size
@@ -98,7 +98,7 @@ def get_dataloaders(cfg : DictConfig):
 
 if __name__ == '__main__':
     cfg = OmegaConf.load('/home/irmak/Workspace/DAWGE/contrastive_learning/configs/train.yaml')
-    dset = Dataset(
+    dset = StateDataset(
         data_dir = cfg.data_dir
     )
 

contrastive_learning/datasets/dataset.py → contrastive_learning/datasets/visual_dataset.py

@@ -16,7 +16,7 @@
 from contrastive_learning.datasets.preprocess import dump_video_to_images, create_pos_pairs
 import contrastive_learning.utils.utils as utils
 
-class Dataset:
+class VisualDataset:
     def __init__(self, data_dir: str,
                  frame_interval: int,
                  video_type: str = 'color') -> None:
@@ -87,7 +87,7 @@ def calculate_action_mean_std(self):
 
 def get_dataloaders(cfg : DictConfig):
     # Load dataset - splitting will be done with random splitter
-    dataset = Dataset(data_dir=cfg.data_dir, frame_interval=cfg.frame_interval, video_type=cfg.video_type)
+    dataset = VisualDataset(data_dir=cfg.data_dir, frame_interval=cfg.frame_interval, video_type=cfg.video_type)
 
     train_dset_size = int(len(dataset) * cfg.train_dset_split)
     test_dset_size = len(dataset) - train_dset_size
@@ -96,7 +96,6 @@ def get_dataloaders(cfg : DictConfig):
     train_dset, test_dset = data.random_split(dataset, 
                                              [train_dset_size, test_dset_size],
                                              generator=torch.Generator().manual_seed(cfg.seed))
-    print('len(train_dset): {}'.format(len(train_dset)))
     train_sampler = data.DistributedSampler(train_dset, drop_last=True, shuffle=True) if cfg.distributed else None
     test_sampler = data.DistributedSampler(test_dset, drop_last=True, shuffle=False) if cfg.distributed else None # val will not be shuffled
 
@@ -105,12 +104,12 @@ def get_dataloaders(cfg : DictConfig):
     test_loader = data.DataLoader(test_dset, batch_size=cfg.batch_size, shuffle=test_sampler is None,
                                     num_workers=cfg.num_workers, sampler=test_sampler)
 
-    return train_loader, test_loader, train_dset, test_dset
+    return train_loader, test_loader, dataset
 
 def plot_data(data_dir:str, frame_interval:int, num_images:int = 16) -> None:
 
     # Get the data loaders
-    dataset = Dataset(data_dir=data_dir, frame_interval=8)
+    dataset = VisualDataset(data_dir=data_dir, frame_interval=8)
     train_dset_size = int(len(dataset) * 0.8)
     val_dset_size = len(dataset) - train_dset_size
     # Random split the train and validation datasets

contrastive_learning/models/agents/agent_inits.py

@@ -0,0 +1,54 @@
+import hydra
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+
+# Scripts to initialize different agents - used in training scripts
+def init_pli(cfg, device, rank):
+    # Initialize the model
+    model = hydra.utils.instantiate(cfg.model,
+                                    input_dim=cfg.pos_dim*2, # For dog and box
+                                    action_dim=cfg.action_dim,
+                                    hidden_dim=cfg.hidden_dim).to(device)
+    model = DDP(model, device_ids=[rank], output_device=rank, broadcast_buffers=False)
+
+
+    # Initialize the optimizer
+    # parameters = list(encoder.parameters()) + list(trans.parameters())
+    optimizer = hydra.utils.instantiate(cfg.optimizer,
+                                        params = model.parameters(),
+                                        lr = cfg.lr,
+                                        weight_decay = cfg.weight_decay)
+
+    # Initialize the total agent
+    agent = hydra.utils.instantiate(cfg.agent,
+                                    model=model,
+                                    optimizer=optimizer)
+
+    agent.to(device)
+
+    return agent
+
+def init_cpn(cfg, device, rank):
+    # Initialize the encoder and the trans
+    encoder = hydra.utils.instantiate(cfg.encoder).to(device)
+    trans = hydra.utils.instantiate(cfg.trans,
+                                    z_dim=cfg.z_dim,
+                                    action_dim=cfg.action_dim).to(device)
+    encoder = DDP(encoder, device_ids=[rank], output_device=rank, broadcast_buffers=False) # To fix the inplace error https://github.com/pytorch/pytorch/issues/22095 
+    trans = DDP(trans, device_ids=[rank], output_device=rank, broadcast_buffers=False)
+
+    # Initialize the optimizer
+    parameters = list(encoder.parameters()) + list(trans.parameters())
+    optimizer = hydra.utils.instantiate(cfg.optimizer,
+                                        params = parameters,
+                                        lr = cfg.lr,
+                                        weight_decay = cfg.weight_decay)
+
+    # Initialize the total agent
+    agent = hydra.utils.instantiate(cfg.agent,
+                                    encoder=encoder,
+                                    trans=trans,
+                                    optimizer=optimizer)
+    agent.to(device)
+
+    return agent 

contrastive_learning/models/agents/cpn.py

@@ -20,7 +20,6 @@ def __init__(self,
                  loss_fn: str # will be a string to indicate the loss function to be used
                  ) -> None:
 
-        # print(f'encoder: {encoder}, trans: {trans}, optimizer: {optimizer}, loss_fn: {loss_fn}')
         self.encoder = encoder 
         self.trans = trans 
         self.optimizer = optimizer 
@@ -61,7 +60,9 @@ def train_epoch(self, train_loader):
             obs, obs_next, action = [b.to(self.device) for b in batch]
 
             # Get the loss - NOTE these parameters will need modificationafterwards
-            loss = self.loss_fn(self.encoder, self.trans, obs, obs_next, action)
+            z, z_next = self.encoder(obs), self.encoder(obs_next) # b x z_dim 
+            z_next_predict = self.trans(z, action)  # b x z_dim
+            loss = self.loss_fn(z, z_next, z_next_predict) # TODO: infonce was changed so you should check this
             train_loss += loss.item()
 
             # Back prop
@@ -82,7 +83,9 @@ def test_epoch(self, test_loader):
         for batch in test_loader:
             obs, obs_next, action = [b.to(self.device) for b in batch]
             with torch.no_grad():
-                loss = self.loss_fn(self.encoder, self.trans, obs, obs_next, action)
+                z, z_next = self.encoder(obs), self.encoder(obs_next) # b x z_dim 
+                z_next_predict = self.trans(z, action)  # b x z_dim
+                loss = self.loss_fn(z, z_next, z_next_predict) # TODO: infonce was changed so you should check this
                 test_loss += loss.item()
 
         return test_loss / len(test_loader)

contrastive_learning/models/custom_models.py

@@ -37,11 +37,11 @@ class LinearInverse(nn.Module):
     def __init__(self, input_dim, action_dim, hidden_dim):
         super().__init__()
         self.model = nn.Sequential(
-            nn.Linear(input_dim*2, hidden_dim),
+            nn.Linear(input_dim*2, hidden_dim), # input_dim*2: For current and next (so in total 32)
             nn.ReLU(),
-            nn.Linear(hidden_dim, hidden_dim),
+            nn.Linear(hidden_dim, int(hidden_dim/4)),
             nn.ReLU(),
-            nn.Linear(hidden_dim, action_dim)
+            nn.Linear(int(hidden_dim/4), action_dim)
         )        
 
     def forward(self, curr_pos, next_pos):