Add support for datasets recorded from robot; set up as pip installable library

.gitignore

@@ -523,4 +523,10 @@ test_*
 *.out
 
 # Hardware gitignore
-.DS_Store
+.DS_Store
+
+# Vim stuff
+*.swp
+*.swo
+*.swn
+*.swm

imitation-in-homes/dobbe/configs/dataset/debug.yaml

@@ -0,0 +1,57 @@
+train:
+  _target_: dobbe.dataloaders.decord_dataset.DecordDataset
+  config:
+    _target_: dobbe.dataloaders.utils.DataLoaderConfig
+    dataset_root: ${env_vars.debug_data_root}
+    trajectory_roots:
+      _target_: dobbe.dataloaders.utils.create_trajectory_roots
+      trajectory_root_path: ${dataset.train.config.dataset_root}
+      original_root: ${env_vars.debug_data_original_root}
+      new_root: ${dataset.train.config.dataset_root}
+    include_tasks: 
+      - ${include_task}
+    exclude_homes: []
+    include_envs: 
+      - ${include_env}
+    sequence_length: 1
+    trajectory_subsample_fraction: 1.0
+    control_timeskip: 7
+    fps_subsample: 1
+    relative_gripper: False
+    use_depth: ${use_depth}
+    binarize_gripper: false
+    binarize_gripper_threshold: 0.8
+    binarize_gripper_upper_value: 1.0
+    binarize_gripper_lower_value: 0.35
+  max_videoreaders_in_memory: 128
+  shuffle_mode: RANDOM_ORDER
+  device: cpu
+
+
+test:
+  _target_: dobbe.dataloaders.decord_dataset.DecordDataset
+  config:
+    _target_: dobbe.dataloaders.utils.DataLoaderConfig
+    dataset_root: ${env_vars.debug_data_root}
+    trajectory_roots:
+      _target_: dobbe.dataloaders.utils.create_trajectory_roots
+      trajectory_root_path: ${dataset.test.config.dataset_root}
+      original_root: ${env_vars.debug_data_original_root}
+      new_root: ${dataset.test.config.dataset_root}
+    include_tasks: 
+      - ${include_task}
+    exclude_homes: []
+    include_envs: 
+      - validation
+    sequence_length: ${dataset.train.config.sequence_length}
+    control_timeskip: ${dataset.train.config.control_timeskip}
+    fps_subsample: ${dataset.train.config.fps_subsample}
+    relative_gripper: ${dataset.train.config.relative_gripper}
+    use_depth: ${use_depth}
+    binarize_gripper: ${dataset.train.config.binarize_gripper}
+    binarize_gripper_threshold: ${dataset.train.config.binarize_gripper_threshold}
+    binarize_gripper_upper_value: ${dataset.train.config.binarize_gripper_upper_value}
+    binarize_gripper_lower_value: ${dataset.train.config.binarize_gripper_lower_value}
+  max_videoreaders_in_memory: 128
+  shuffle_mode: SEQUENTIAL
+  device: cpu

imitation-in-homes/configs/dataset/finetune_task_data.yaml → imitation-in-homes/dobbe/configs/dataset/finetune_task_data.yaml

@@ -1,10 +1,10 @@
 train:
-  _target_: dataloaders.decord_dataset.DecordDataset
+  _target_: dobbe.dataloaders.decord_dataset.DecordDataset
   config:
-    _target_: dataloaders.utils.DataLoaderConfig
+    _target_: dobbe.dataloaders.utils.DataLoaderConfig
     dataset_root: ${env_vars.finetune_task_data_root}
     trajectory_roots:
-      _target_: dataloaders.utils.load_trajectory_roots
+      _target_: dobbe.dataloaders.utils.load_trajectory_roots
       trajectory_root_path: ${dataset.train.config.dataset_root}/r3d_files.txt
       original_root: ${env_vars.finetune_task_data_original_root}
       new_root: ${dataset.train.config.dataset_root}
@@ -29,12 +29,12 @@ train:
 
 
 test:
-  _target_: dataloaders.decord_dataset.DecordDataset
+  _target_: dobbe.dataloaders.decord_dataset.DecordDataset
   config:
-    _target_: dataloaders.utils.DataLoaderConfig
+    _target_: dobbe.dataloaders.utils.DataLoaderConfig
     dataset_root: ${env_vars.finetune_task_data_root}
     trajectory_roots:
-      _target_: dataloaders.utils.load_trajectory_roots
+      _target_: dobbe.dataloaders.utils.load_trajectory_roots
       trajectory_root_path: ${dataset.test.config.dataset_root}/r3d_files.txt
       original_root: ${env_vars.finetune_task_data_original_root}
       new_root: ${dataset.test.config.dataset_root}

imitation-in-homes/dobbe/configs/debug.yaml

@@ -0,0 +1,53 @@
+defaults:
+  - env_vars: env_vars
+  - dataset:  debug
+  - model: image_bc_depth_timm
+  - loss_fn: bc_loss
+  - _self_
+
+num_epochs: 50
+eval_every: 1
+device: cuda
+batch_size: 256
+use_depth: true
+
+include_task: cactus
+include_env: default_env
+
+sampler: null
+
+batch_sampler:
+  _target_: dobbe.dataloaders.decord_dataset.DecordBatchSampler
+  batch_size: ${batch_size}
+
+dataloader:
+  _target_: dobbe.dataloaders.decord_dataloader.DecordVideoLoader
+  device: cpu
+  num_workers: 8
+
+
+optimizer:
+  _target_: torch.optim.Adam
+  lr: 0.00003
+  weight_decay: 1.5e-6
+
+scheduler:
+  _target_: dobbe.models.ssl.moco.optimizer.MoCoScheduler
+  warmup_epochs: 10
+  start_lr: 2e-6
+  num_epochs: ${num_epochs}
+
+
+model_path: checkpoints/${now:%Y-%m-%d}/${include_task}-${include_env}-moco-${now:%H-%M-%S}
+resume_training: false
+save_every: 10
+
+wandb:
+  entity: ${env_vars.wandb.entity}
+  project: dobbe-${include_task}
+  name: ${include_task}-${include_env}-moco-4x4
+  # id starts out null, hoawever, if you resume training, it will be set to the original id.
+  id: null
+  save_code: true
+
+seed: 42

imitation-in-homes/dobbe/configs/env_vars/env_vars.yaml.sample

@@ -0,0 +1,9 @@
+home_ssl_data_root: /home/shared/data/iphone_data  # root to home ssl data folder
+finetune_task_data_root: /home/cpaxton/data/finetune_directory  # root to the exported finetuning data folder
+
+home_ssl_data_original_root: /path/to/directory  # path included in the r3d_files.txt file of dataset
+finetune_task_data_original_root: /home/cpaxton/finetune_directory  # path included in r3d_files.txt of dataset (may be the same as finetune_task_data_root if data has not been moved nor transferred across machines)
+project_root: /home/shared/code/imitation-in-homes  # path to this repo
+
+wandb:
+    entity: null  # wandb username

imitation-in-homes/configs/finetune.yaml → imitation-in-homes/dobbe/configs/finetune.yaml

@@ -17,11 +17,11 @@ include_env: Env1
 sampler: null
 
 batch_sampler:
-  _target_: dataloaders.decord_dataset.DecordBatchSampler
+  _target_: dobbe.dataloaders.decord_dataset.DecordBatchSampler
   batch_size: ${batch_size}
 
 dataloader:
-  _target_: dataloaders.decord_dataloader.DecordVideoLoader
+  _target_: dobbe.dataloaders.decord_dataloader.DecordVideoLoader
   device: cpu
   num_workers: 8
 
@@ -32,7 +32,7 @@ optimizer:
   weight_decay: 1.5e-6
 
 scheduler:
-  _target_: models.ssl.moco.optimizer.MoCoScheduler
+  _target_: dobbe.models.ssl.moco.optimizer.MoCoScheduler
   warmup_epochs: 10
   start_lr: 2e-6
   num_epochs: ${num_epochs}
@@ -50,4 +50,4 @@ wandb:
   id: null
   save_code: true
 
-seed: 42
+seed: 42

imitation-in-homes/dobbe/configs/loss_fn/bc_loss.yaml

@@ -0,0 +1 @@
+_target_: dobbe.loss_fns.bc_loss_fn.BCPolicyLossFn

imitation-in-homes/configs/loss_fn/mae_loss.yaml → imitation-in-homes/dobbe/configs/loss_fn/mae_loss.yaml

@@ -1,4 +1,4 @@
-_target_: loss_fns.ssl.mae_loss_fn.MAELossFn
+_target_: dobbe.loss_fns.ssl.mae_loss_fn.MAELossFn
 decoder_dim: 512
 masking_ratio: 0.75
 decoder_depth: 6

imitation-in-homes/dobbe/configs/loss_fn/moco_loss.yaml

@@ -0,0 +1 @@
+_target_: dobbe.loss_fns.ssl.moco_loss_fn.MocoLossFn

imitation-in-homes/configs/model/image_bc_depth.yaml → imitation-in-homes/dobbe/configs/model/image_bc_depth.yaml

@@ -1,6 +1,6 @@
-_target_: models.policies.bc.BCDepthModel
+_target_: dobbe.models.policies.bc.BCDepthModel
 encoder:
-  _target_: models.encoders.torch_encoders.TorchSSL
+  _target_: dobbe.models.encoders.torch_encoders.TorchSSL
   model_name: resnet34
   override_aug_kwargs:
     hflip: 0.0

imitation-in-homes/configs/model/image_bc_depth_timm.yaml → imitation-in-homes/dobbe/configs/model/image_bc_depth_timm.yaml

@@ -1,6 +1,6 @@
-_target_: models.policies.bc.BCDepthModel
+_target_: dobbe.models.policies.bc.BCDepthModel
 encoder:
-  _target_: models.encoders.timm_encoders.TimmSSL
+  _target_: dobbe.models.encoders.timm_encoders.TimmSSL
   model_name: hf-hub:notmahi/dobb-e
   override_aug_kwargs:
     hflip: 0.0

imitation-in-homes/dobbe/configs/model/moco.yaml

@@ -0,0 +1,3 @@
+_target_: dobbe.models.encoders.timm_encoders.TimmSSLByolMoco
+model_name: resnet18
+pretrained: True

imitation-in-homes/configs/run.yaml → imitation-in-homes/dobbe/configs/run.yaml

@@ -1,22 +1,22 @@
 defaults:
   - env_vars: env_vars
-  - model: image_bc_depth
-  - dataset: finetune_task_data
+  - model: image_bc_depth_timm
+  - dataset: debug
   - _self_
 
 image_save_dir: ${env_vars.project_root}/robot_images
 image_buffer_size: 1
 device: cpu
 
-include_task: Random
-include_env: Env1
+include_task: cactus
+include_env: default_env
 
-model_weight_pth: TODO_FINETUNED_MODEL_PATH
+model_weight_pth: ${env_vars.project_root}/checkpoints/2024-05-07/cactus-default_env-moco-13-27-02/checkpoint_40.pt
 
 experiment: bc_with_hpr_encoder
 batch_size: 32
 sampler: null
-run_offline: false
+run_offline: true
 use_depth: true
 
 robot_params:

imitation-in-homes/configs/train_moco.yaml → imitation-in-homes/dobbe/configs/train_moco.yaml

@@ -1,23 +1,28 @@
 defaults:
   - env_vars: env_vars
-  - dataset: hony
+  - dataset: debug
   - model: moco
   - loss_fn: moco_loss
   - _self_
 
 num_epochs: 60
 eval_every: 5
 device: cuda
-batch_size: 800
+# batch_size: 800
+batch_size: 256
+use_depth: true
+
+include_task: cactus
+include_env: default_env
 
 sampler: null
 
 batch_sampler:
-  _target_: dataloaders.decord_dataset.DecordBatchSampler
+  _target_: dobbe.dataloaders.decord_dataset.DecordBatchSampler
   batch_size: ${batch_size}
 
 dataloader:
-  _target_: dataloaders.decord_dataloader.DecordVideoLoader
+  _target_: dobbe.dataloaders.decord_dataloader.DecordVideoLoader
   device: cpu
   num_workers: 10
 
@@ -28,7 +33,7 @@ optimizer:
   weight_decay: 1.0e-6
 
 scheduler:
-  _target_: models.ssl.moco.optimizer.MoCoScheduler
+  _target_: dobbe.models.ssl.moco.optimizer.MoCoScheduler
   warmup_epochs: 70
   num_epochs: ${num_epochs}
 
@@ -44,4 +49,4 @@ wandb:
   id: null
   save_code: true
 
-seed: 42
+seed: 42

imitation-in-homes/dataloaders/abstract_dataloader.py → imitation-in-homes/dobbe/dataloaders/abstract_dataloader.py

@@ -4,9 +4,9 @@
 
 from torch.utils.data import Dataset
 
-from dataloaders.depth_data import DepthDataLoader
-from dataloaders.pose_data import PoseDataLoader
-from dataloaders.utils import DataLoaderConfig, calculate_start_end_indices
+from dobbe.dataloaders.depth_data import DepthDataLoader
+from dobbe.dataloaders.pose_data import PoseDataLoader
+from dobbe.dataloaders.utils import DataLoaderConfig, calculate_start_end_indices
 
 
 class AbstractVideoDataset(ABC, Dataset):

imitation-in-homes/dataloaders/decord_dataloader.py → imitation-in-homes/dobbe/dataloaders/decord_dataloader.py

@@ -5,7 +5,7 @@
 import decord
 from torch.utils.data import DataLoader, Dataset, Sampler
 
-from dataloaders.decord_dataset import DecordDataset
+from dobbe.dataloaders.decord_dataset import DecordDataset
 
 
 class DecordVideoLoader(DataLoader):
@@ -25,6 +25,9 @@ def __init__(
         # Check if the dataset is a DecordDataset and the device is GPU
         if isinstance(dataset, DecordDataset):
             self._is_decord_dataset = True
+        else:
+            print(dataset)
+            self._is_decord_dataset = False
         if num_workers > 0 and self._is_decord_dataset:
             if device == "gpu":
                 warnings.warn(

imitation-in-homes/dataloaders/decord_dataset.py → imitation-in-homes/dobbe/dataloaders/decord_dataset.py

@@ -20,8 +20,8 @@
 import torch
 import tqdm
 
-from dataloaders.abstract_dataloader import AbstractVideoDataset
-from dataloaders.utils import DataLoaderConfig, TrajectorySlice, flatten_nested_lists
+from dobbe.dataloaders.abstract_dataloader import AbstractVideoDataset
+from dobbe.dataloaders.utils import DataLoaderConfig, TrajectorySlice, flatten_nested_lists
 
 
 # Create a dataclass for shuffle modes.