gkuppa/compress imgs

.gitignore

@@ -137,4 +137,6 @@ celerybeat.pid
 dmypy.json
 
 # Pyre type checker
-.pyre/
+.pyre/
+
+training_outputs/

README.md

@@ -29,20 +29,7 @@ You can find all scripted/human demo for simulated environments [here](https://d
 
     conda create -n aloha python=3.8.10
     conda activate aloha
-    pip install torchvision
-    pip install torch
-    pip install pyquaternion
-    pip install pyyaml
-    pip install rospkg
-    pip install pexpect
-    pip install mujoco==2.3.7
-    pip install dm_control==1.0.14
-    pip install opencv-python
-    pip install matplotlib
-    pip install einops
-    pip install packaging
-    pip install h5py
-    pip install ipython
+    pip install -r requirements.txt
     cd act/detr && pip install -e .
 
 ### Example Usages

detr/experiment_configs/real_env_constants.py

@@ -0,0 +1,79 @@
+### Task parameters
+
+import os
+
+DATA_DIR = ''
+TASK_CONFIGS = {
+    # this config is only inserting in the middle
+    'insert_redness_relief_slanted':{
+        'dataset_dir': os.environ['DATA'] + '/closed_loop_demos/insert_redness_relief_slanted',
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+    },
+
+    # this config is only inserting anywhere
+    'insert_redness_relief_slanted_anywhere':{
+        'dataset_dir': [
+                        os.environ['DATA'] + '/closed_loop_demos/insert_redness_relief_slanted_anywhere',
+                        ],
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+    },
+
+    # this config is only inserting anywhere
+    'pick_vial':{
+        'dataset_dir': [
+                        os.environ['DATA'] + '/closed_loop_demos/pick_place_phenylephrine/',
+                        ],
+        'stage_key': '1_pick',
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+    },
+
+    'place_vial':{
+        'dataset_dir': [
+                        os.environ['DATA'] + '/closed_loop_demos/pick_place_phenylephrine/',
+                        ],
+        'stage_key': '2_place',
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+    },
+
+    'insert_ibuprofen':{
+        'dataset_dir': [
+                        os.environ['DATA'] + '/closed_loop_demos/insert_ibuprofen',
+                        ],
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+    },
+
+    'insert_ibuprofen_zeroqpos':{
+        'dataset_dir': [
+                        os.environ['DATA'] + '/closed_loop_demos/insert_ibuprofen',
+                        ],
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+        'zero_qpos': True
+    },
+
+    'insert_ibuprofen_rel':{
+        'dataset_dir': [
+                        os.environ['DATA'] + '/closed_loop_demos/insert_ibuprofen_relactions',
+                        ],
+        'num_episodes': -1,
+        'episode_len': 80,
+        'camera_names': ['color_image'],
+        'actiondim':7,
+        'relative_actions': True
+    },
+}
+
+### ALOHA fixed constants
+DT = 0.02
+

detr/main.py

@@ -30,6 +30,11 @@ def get_args_parser():
                         help="Type of positional embedding to use on top of the image features")
     parser.add_argument('--camera_names', default=[], type=list, # will be overridden
                         help="A list of camera names")
+
+    parser.add_argument('--statedim', type=int, default=9,# will be overridden
+                        help="statedim")
+    parser.add_argument('--actiondim', type=int, default=10, # will be overridden
+                        help="actiondim")
 
     # * Transformer
     parser.add_argument('--enc_layers', default=4, type=int, # will be overridden
@@ -56,13 +61,17 @@ def get_args_parser():
     parser.add_argument('--eval', action='store_true')
     parser.add_argument('--onscreen_render', action='store_true')
     parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', required=True)
+    parser.add_argument('--ckpt_name', action='store', type=str, help='ckpt_name', required=False)
     parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize', required=True)
     parser.add_argument('--task_name', action='store', type=str, help='task_name', required=True)
     parser.add_argument('--seed', action='store', type=int, help='seed', required=True)
     parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', required=True)
     parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', required=False)
     parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size', required=False)
     parser.add_argument('--temporal_agg', action='store_true')
+    parser.add_argument('--run_name', action='store', type=str, help='run_name', required=True)
+
+    parser.add_argument('--wandb_mode', type=str, default='online', help='wandb mode') # online, disabled
 
     return parser
 

detr/models/backbone.py

@@ -11,7 +11,8 @@
 from torchvision.models._utils import IntermediateLayerGetter
 from typing import Dict, List
 
-from util.misc import NestedTensor, is_main_process
+# from util.misc import NestedTensor, is_main_process
+from detr.util.misc import NestedTensor, is_main_process
 
 from .position_encoding import build_position_encoding
 

detr/models/detr_vae.py

@@ -31,9 +31,13 @@ def get_position_angle_vec(position):
     return torch.FloatTensor(sinusoid_table).unsqueeze(0)
 
 
+# STATE_DIM = 9  # 3 for xyz + 6 for r6 rotations
+# ACTION_DIM = 10
+
+
 class DETRVAE(nn.Module):
     """ This is the DETR module that performs object detection """
-    def __init__(self, backbones, transformer, encoder, state_dim, num_queries, camera_names):
+    def __init__(self, backbones, transformer, encoder, num_queries, camera_names, statedim, actiondim):
         """ Initializes the model.
         Parameters:
             backbones: torch module of the backbone to be used. See backbone.py
@@ -49,25 +53,25 @@ def __init__(self, backbones, transformer, encoder, state_dim, num_queries, came
         self.transformer = transformer
         self.encoder = encoder
         hidden_dim = transformer.d_model
-        self.action_head = nn.Linear(hidden_dim, state_dim)
+        self.action_head = nn.Linear(hidden_dim, actiondim)
         self.is_pad_head = nn.Linear(hidden_dim, 1)
         self.query_embed = nn.Embedding(num_queries, hidden_dim)
         if backbones is not None:
             self.input_proj = nn.Conv2d(backbones[0].num_channels, hidden_dim, kernel_size=1)
             self.backbones = nn.ModuleList(backbones)
-            self.input_proj_robot_state = nn.Linear(14, hidden_dim)
+            self.input_proj_robot_state = nn.Linear(statedim, hidden_dim)
         else:
             # input_dim = 14 + 7 # robot_state + env_state
-            self.input_proj_robot_state = nn.Linear(14, hidden_dim)
+            self.input_proj_robot_state = nn.Linear(statedim, hidden_dim)
             self.input_proj_env_state = nn.Linear(7, hidden_dim)
             self.pos = torch.nn.Embedding(2, hidden_dim)
             self.backbones = None
 
         # encoder extra parameters
         self.latent_dim = 32 # final size of latent z # TODO tune
         self.cls_embed = nn.Embedding(1, hidden_dim) # extra cls token embedding
-        self.encoder_action_proj = nn.Linear(14, hidden_dim) # project action to embedding
-        self.encoder_joint_proj = nn.Linear(14, hidden_dim)  # project qpos to embedding
+        self.encoder_action_proj = nn.Linear(actiondim, hidden_dim) # project action to embedding
+        self.encoder_joint_proj = nn.Linear(statedim, hidden_dim)  # project qpos to embedding
         self.latent_proj = nn.Linear(hidden_dim, self.latent_dim*2) # project hidden state to latent std, var
         self.register_buffer('pos_table', get_sinusoid_encoding_table(1+1+num_queries, hidden_dim)) # [CLS], qpos, a_seq
 
@@ -82,6 +86,7 @@ def forward(self, qpos, image, env_state, actions=None, is_pad=None):
         env_state: None
         actions: batch, seq, action_dim
         """
+
         is_training = actions is not None # train or val
         bs, _ = qpos.shape
         ### Obtain latent z from action sequence
@@ -166,8 +171,8 @@ def __init__(self, backbones, state_dim, camera_names):
                 backbone_down_projs.append(down_proj)
             self.backbone_down_projs = nn.ModuleList(backbone_down_projs)
 
-            mlp_in_dim = 768 * len(backbones) + 14
-            self.mlp = mlp(input_dim=mlp_in_dim, hidden_dim=1024, output_dim=14, hidden_depth=2)
+            mlp_in_dim = 768 * len(backbones) + STATE_DIM
+            self.mlp = mlp(input_dim=mlp_in_dim, hidden_dim=1024, output_dim=ACTION_DIM, hidden_depth=2)
         else:
             raise NotImplementedError
 
@@ -192,7 +197,7 @@ def forward(self, qpos, image, env_state, actions=None):
         for cam_feature in all_cam_features:
             flattened_features.append(cam_feature.reshape([bs, -1]))
         flattened_features = torch.cat(flattened_features, axis=1) # 768 each
-        features = torch.cat([flattened_features, qpos], axis=1) # qpos: 14
+        features = torch.cat([flattened_features, qpos], axis=1)
         a_hat = self.mlp(features)
         return a_hat
 
@@ -227,7 +232,6 @@ def build_encoder(args):
 
 
 def build(args):
-    state_dim = 14 # TODO hardcode
 
     # From state
     # backbone = None # from state for now, no need for conv nets
@@ -244,7 +248,8 @@ def build(args):
         backbones,
         transformer,
         encoder,
-        state_dim=state_dim,
+        statedim=args.statedim,
+        actiondim=args.actiondim,
         num_queries=args.num_queries,
         camera_names=args.camera_names,
     )

detr/models/position_encoding.py

@@ -6,7 +6,8 @@
 import torch
 from torch import nn
 
-from util.misc import NestedTensor
+# from util.misc import NestedTensor
+from detr.util.misc import NestedTensor
 
 import IPython
 e = IPython.embed

detr/util/misc.py

@@ -17,12 +17,29 @@
 import torch.distributed as dist
 from torch import Tensor
 
+import wandb 
+
 # needed due to empty tensor bug in pytorch and torchvision 0.5
 import torchvision
 if version.parse(torchvision.__version__) < version.parse('0.7'):
     from torchvision.ops import _new_empty_tensor
     from torchvision.ops.misc import _output_size
 
+def wandb_setup(params, save_dir, mode, run_id=None):
+    os.environ['WANDB_API_KEY'] = '0d98919f1dac74f61fae02938822ece99e40d095'
+    os.environ['WANDB_USER_EMAIL'] = 'febert@emancro.ai'
+    os.environ['WANDB_USERNAME'] ='febert1'
+    wandb_run = wandb.init(
+                        project='closed_loop_policy',
+                        config=params,
+                        name=params['run_name'],
+                        id=run_id,
+                        resume=run_id is not None,
+                        dir=save_dir,
+                        mode=mode
+                  )
+    return wandb_run.id
+
 
 class SmoothedValue(object):
     """Track a series of values and provide access to smoothed values over a

docker/README.md

@@ -0,0 +1,2 @@
+# copy entire folder to the cloud:
+gcloud compute scp --recurse act febert@a100-2:/home/febert/code --zone=us-central1-c

docker/access_docker.sh

@@ -0,0 +1,5 @@
+docker run --gpus all -it --rm --shm-size 8G \
+    --volume $CODE/act:/home/app/act \
+    --volume $DATA:/home/app/data \
+    -e "DATA=/home/app/data/" \
+    emancro/act:latest /bin/bash -c "cd /home/app/act/detr && pip install -e . && cd /home/app/act && /bin/bash "

docker/build.sh

@@ -0,0 +1,2 @@
+cd ..
+docker build -f docker/dockerfile -t emancro/act:latest .

docker/dockerfile

@@ -0,0 +1,20 @@
+FROM continuumio/miniconda3
+RUN conda create -n aloha python=3.10
+
+# Set the working directory
+WORKDIR /app
+# Copy the requirements file from the local folder to the container
+COPY requirements.txt /app/requirements.txt
+
+ENV DATA /workspace
+
+
+RUN /bin/bash -c "source activate aloha && \
+            pip install -r /app/requirements.txt"
+
+WORKDIR /workspace/act/detr
+
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
+
+
+CMD ["bash"]

download_files.py

@@ -0,0 +1,54 @@
+import os
+
+
+GRAB_ITERATIONS = [
+    'policy_epoch_12000_seed_0.ckpt',
+    'policy_epoch_18000_seed_0.ckpt',
+    'policy_epoch_5000_seed_0.ckpt'    
+]
+
+OTHER_FOLDERS_AND_FILES = [
+    'dataset_stats.pkl'
+]
+
+
+
+def download_from_gcs(args):
+
+    destination_path = f'/home/user/data/act_training_runs/{args.run_name}'
+    source_path = f'/mnt/disks/extra_data/data/act_training/training_outputs/{args.run_name}'
+
+    if not args.dry_run:
+        os.makedirs(destination_path, exist_ok=True)
+    else:
+        print('would create', destination_path)
+
+    for file in OTHER_FOLDERS_AND_FILES:
+        # cmd = 'scp -r ' + args.machine + ':' + source_path + '/' + file + ' ' + destination_path
+        cmd = 'gcloud compute scp --recurse ' + args.machine + ':' + source_path + '/' + file + ' ' + destination_path + ' --zone=us-central1-c'
+        if not args.dry_run:
+            os.system(cmd)
+        else:
+            print('would execute', cmd)
+
+    for iteration in GRAB_ITERATIONS:
+        # cmd = 'scp -r ' +  args.machine + ':' + source_path + '/' + iteration + ' ' + destination_path
+        cmd = 'gcloud compute scp --recurse ' + args.machine + ':' + source_path + '/' + iteration + ' ' + destination_path + ' --zone=us-central1-c'
+        if not args.dry_run:
+            os.system(cmd)
+        else:
+            print('would execute', cmd)
+
+
+
+if __name__ == "__main__":
+    # example
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--machine', default='a100-2', type=str)
+    parser.add_argument('--run_name', default='experiment_20240610_205136', type=str, required=True)   
+    parser.add_argument('--dry_run', action='store_true')
+
+    args = parser.parse_args()
+
+    download_from_gcs(args)

ee_sim_env.py

@@ -2,11 +2,11 @@
 import collections
 import os
 
-from constants import DT, XML_DIR, START_ARM_POSE
-from constants import PUPPET_GRIPPER_POSITION_CLOSE
-from constants import PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN
-from constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN
-from constants import PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN
+from detr.experiment_configs.sim_constants import DT, XML_DIR, START_ARM_POSE
+from detr.experiment_configs.sim_constants import PUPPET_GRIPPER_POSITION_CLOSE
+from detr.experiment_configs.sim_constants import PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN
+from detr.experiment_configs.sim_constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN
+from detr.experiment_configs.sim_constants import PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN
 
 from utils import sample_box_pose, sample_insertion_pose
 from dm_control import mujoco