auto record mode and small reorg

.gitignore

@@ -1,4 +1,4 @@
 .idea
 .DS_store
 /aloha.egg-info/
-/scripts/.idea/
+/aloha_scripts/.idea/

README.md

@@ -7,7 +7,7 @@ This codebase contains implementation for teleoperation and data collection with
 ### Repo Structure
 - ``config``: a config for each robot, designating the port they should bind to, more details in quick start guide.
 - ``launch``: a ROS launch file for all 4 cameras and all 4 robots.
-- ``scripts``: python code for teleop and data collection
+- ``aloha_scripts``: python code for teleop and data collection
 
 ## Quick start guide
 
@@ -139,12 +139,12 @@ All robots will rise to a height that is easy for teleoperation.
 
     # Right hand terminal
     conda activate aloha
-    cd ~/interbotix_ws/src/aloha/scripts
+    cd ~/interbotix_ws/src/aloha/aloha_scripts
     python3 one_side_teleop.py right
 
     # Left hand terminal
     conda activate aloha
-    cd ~/interbotix_ws/src/aloha/scripts
+    cd ~/interbotix_ws/src/aloha/aloha_scripts
     python3 one_side_teleop.py left
 
 The teleoperation will start when the master side gripper is closed.
@@ -155,7 +155,7 @@ The teleoperation will start when the master side gripper is closed.
 To set up a new terminal, run:
 
     conda activate aloha
-    cd ~/interbotix_ws/src/aloha/scripts
+    cd ~/interbotix_ws/src/aloha/aloha_scripts
 
 
 The ``one_side_teleop.py`` we ran is for testing teleoperation and has no data collection. To collect data for an episode, run:

aloha_scripts/auto_record.sh

@@ -0,0 +1,15 @@
+if [ "$2" -lt 0 ]; then
+  echo "# of episodes not valid"
+  exit
+fi
+
+echo "Task: $1"
+for (( i=0; i<$2; i++ ))
+do
+  echo "Starting episode $i"
+  python3 record_episodes.py --task "$1"
+  if [ $? -ne 0 ]; then
+    echo "Failed to execute command. Returning"
+    exit
+  fi
+done

scripts/constants.py → aloha_scripts/constants.py

@@ -1,14 +1,19 @@
-
-### Parameters that changes across tasks
-EPISODE_LEN = 600
+### Task parameters
+
+DATA_DIR = '<put your data dir here>'
+TASK_CONFIGS = {
+    'aloha_wear_shoe':{
+        'dataset_dir': DATA_DIR + '/aloha_wear_shoe',
+        'num_episodes': 50,
+        'episode_len': 400,
+        'camera_names': ['cam_high', 'cam_low', 'cam_left_wrist', 'cam_right_wrist']
+    },
+}
 
 ### ALOHA fixed constants
 DT = 0.02
 JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"]
-CAMERA_NAMES = ['cam_high', 'cam_low', 'cam_left_wrist', 'cam_right_wrist'] # defines the number and ordering of cameras
-BOX_INIT_POSE = [0.2, 0.5, 0.05, 1, 0, 0, 0]
 START_ARM_POSE = [0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239,  0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239]
-SIM_CAMERA_NAMES = ['main']
 
 # Left finger position limits (qpos[7]), right_finger = -1 * left_finger
 MASTER_GRIPPER_POSITION_OPEN = 0.02417

scripts/one_side_teleop.py → aloha_scripts/one_side_teleop.py

@@ -6,7 +6,7 @@
 from interbotix_xs_modules.arm import InterbotixManipulatorXS
 from interbotix_xs_msgs.msg import JointSingleCommand
 from constants import MASTER2PUPPET_JOINT_FN, DT, START_ARM_POSE, MASTER_GRIPPER_JOINT_MID, PUPPET_GRIPPER_JOINT_CLOSE
-from utils import torque_on, torque_off, move_arms, move_grippers, get_arm_gripper_positions
+from robot_utils import torque_on, torque_off, move_arms, move_grippers, get_arm_gripper_positions
 
 def prep_robots(master_bot, puppet_bot):
     # reboot gripper motors, and set operating modes for all motors

scripts/real_env.py → aloha_scripts/real_env.py

@@ -7,8 +7,8 @@
 from constants import DT, START_ARM_POSE, MASTER_GRIPPER_JOINT_NORMALIZE_FN, PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN
 from constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN, PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN
 from constants import PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE
-from utils import Recorder, ImageRecorder
-from utils import setup_master_bot, setup_puppet_bot, move_arms, move_grippers
+from robot_utils import Recorder, ImageRecorder
+from robot_utils import setup_master_bot, setup_puppet_bot, move_arms, move_grippers
 from interbotix_xs_modules.arm import InterbotixManipulatorXS
 from interbotix_xs_msgs.msg import JointSingleCommand
 
@@ -142,10 +142,12 @@ def get_action(master_bot_left, master_bot_right):
 
     return action
 
+
 def make_real_env(init_node, setup_robots=True):
     env = RealEnv(init_node, setup_robots)
     return env
 
+
 def test_real_teleop():
     """
     Test bimanual teleoperation and show image observations onscreen.

scripts/record_episodes.py → aloha_scripts/record_episodes.py

@@ -1,19 +1,15 @@
-# !/usr/bin/env python
 import os
-import rospy
-import cv2
 import time
-import pickle
-import numpy as np
-from tqdm import tqdm
 import h5py
-import h5py_cache
 import argparse
+import h5py_cache
+import numpy as np
+from tqdm import tqdm
 
-from constants import DT, CAMERA_NAMES, EPISODE_LEN, START_ARM_POSE, MASTER_GRIPPER_JOINT_MID, PUPPET_GRIPPER_JOINT_CLOSE
-from constants import PUPPET_GRIPPER_JOINT_OPEN
-from utils import Recorder, ImageRecorder, get_arm_joint_positions, get_arm_gripper_positions
-from utils import setup_master_bot, setup_puppet_bot, move_arms, torque_on, torque_off, move_grippers
+from constants import DT, START_ARM_POSE, TASK_CONFIGS
+from constants import MASTER_GRIPPER_JOINT_MID, PUPPET_GRIPPER_JOINT_CLOSE, PUPPET_GRIPPER_JOINT_OPEN
+from robot_utils import Recorder, ImageRecorder, get_arm_gripper_positions
+from robot_utils import move_arms, torque_on, torque_off, move_grippers
 from real_env import make_real_env, get_action
 
 from interbotix_xs_modules.arm import InterbotixManipulatorXS
@@ -69,7 +65,7 @@ def opening_ceremony(master_bot_left, master_bot_right, puppet_bot_left, puppet_
     print(f'Started!')
 
 
-def capture_one_episode(dt, max_timesteps, dataset_dir, dataset_name, overwrite):
+def capture_one_episode(dt, max_timesteps, camera_names, dataset_dir, dataset_name, overwrite):
     print(f'Dataset name: {dataset_name}')
 
     # source of data
@@ -136,7 +132,7 @@ def capture_one_episode(dt, max_timesteps, dataset_dir, dataset_name, overwrite)
         '/observations/qvel': [],
         '/action': [],
     }
-    for cam_name in CAMERA_NAMES:
+    for cam_name in camera_names:
         data_dict[f'/observations/images/{cam_name}'] = []
 
     # len(action): max_timesteps, len(time_steps): max_timesteps + 1
@@ -146,7 +142,7 @@ def capture_one_episode(dt, max_timesteps, dataset_dir, dataset_name, overwrite)
         data_dict['/observations/qpos'].append(ts.observation['qpos'])
         data_dict['/observations/qvel'].append(ts.observation['qvel'])
         data_dict['/action'].append(action)
-        for cam_name in CAMERA_NAMES:
+        for cam_name in camera_names:
             data_dict[f'/observations/images/{cam_name}'].append(ts.observation['images'][cam_name])
 
     # HDF5
@@ -157,18 +153,11 @@ def capture_one_episode(dt, max_timesteps, dataset_dir, dataset_name, overwrite)
         root.attrs['sim'] = False
         obs = root.create_group('observations')
         image = obs.create_group('images')
-        cam_high = image.create_dataset('cam_high', (max_timesteps, 480, 640, 3), dtype='uint8', chunks=(1, 480, 640, 3),)
-                                        # compression='gzip',compression_opts=2,)
-                                        # compression=32001, compression_opts=(0, 0, 0, 0, 9, 1, 1), shuffle=False)
-        cam_low = image.create_dataset('cam_low', (max_timesteps, 480, 640, 3), dtype='uint8', chunks=(1, 480, 640, 3),)
-                                       # compression='gzip',compression_opts=2,)
-                                       # compression=32001, compression_opts=(0, 0, 0, 0, 9, 1, 1), shuffle=False)
-        cam_left_wrist = image.create_dataset('cam_left_wrist', (max_timesteps, 480, 640, 3), dtype='uint8', chunks=(1, 480, 640, 3),)
-                                              # compression='gzip',compression_opts=2,)
-                                              # compression=32001, compression_opts=(0, 0, 0, 0, 9, 1, 1), shuffle=False)
-        cam_right_wrist = image.create_dataset('cam_right_wrist', (max_timesteps, 480, 640, 3), dtype='uint8', chunks=(1, 480, 640, 3),)
-                                               # compression='gzip',compression_opts=2,)
-                                               # compression=32001, compression_opts=(0, 0, 0, 0, 9, 1, 1), shuffle=False)
+        for cam_name in camera_names:
+            _ = image.create_dataset(cam_name, (max_timesteps, 480, 640, 3), dtype='uint8',
+                                     chunks=(1, 480, 640, 3), )
+            # compression='gzip',compression_opts=2,)
+            # compression=32001, compression_opts=(0, 0, 0, 0, 9, 1, 1), shuffle=False)
         qpos = obs.create_dataset('qpos', (max_timesteps, 14))
         qvel = obs.create_dataset('qvel', (max_timesteps, 14))
         action = root.create_dataset('action', (max_timesteps, 14))
@@ -181,19 +170,35 @@ def capture_one_episode(dt, max_timesteps, dataset_dir, dataset_name, overwrite)
 
 
 def main(args):
-    dataset_dir = args['dataset_dir']
-    episode_idx = args['episode_idx']
-    max_timesteps = EPISODE_LEN
-    dataset_name_prefix = ''
+    task_config = TASK_CONFIGS[args['task_name']]
+    dataset_dir = task_config['dataset_dir']
+    max_timesteps = task_config['episode_len']
+    camera_names = task_config['camera_names']
+
+    if args['episode_idx'] is not None:
+        episode_idx = args['episode_idx']
+    else:
+        episode_idx = get_auto_index(dataset_dir)
     overwrite = True
 
-    dataset_name = f'{dataset_name_prefix}episode_{episode_idx}'
+    dataset_name = f'episode_{episode_idx}'
     print(dataset_name + '\n')
     while True:
-        is_success = capture_one_episode(DT, max_timesteps, dataset_dir, dataset_name, overwrite)
-        if is_success:
+        is_healthy = capture_one_episode(DT, max_timesteps, camera_names, dataset_dir, dataset_name, overwrite)
+        if is_healthy:
             break
 
+
+def get_auto_index(dataset_dir, dataset_name_prefix = '', data_suffix = 'hdf5'):
+    max_idx = 1000
+    if not os.path.isdir(dataset_dir):
+        os.makedirs(dataset_dir)
+    for i in range(max_idx+1):
+        if not os.path.isfile(os.path.join(dataset_dir, f'{dataset_name_prefix}episode_{i}.{data_suffix}')):
+            return i
+    raise Exception(f"Error getting auto index, or more than {max_idx} episodes")
+
+
 def print_dt_diagnosis(actual_dt_history):
     actual_dt_history = np.array(actual_dt_history)
     get_action_time = actual_dt_history[:, 1] - actual_dt_history[:, 0]
@@ -217,8 +222,8 @@ def debug():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset_dir', action='store', type=str, help='Dataset dir.', required=True)
-    parser.add_argument('--episode_idx', action='store', type=int, help='Episode index.', required=False)
+    parser.add_argument('--task_name', action='store', type=str, help='Task name.', required=True)
+    parser.add_argument('--episode_idx', action='store', type=int, help='Episode index.', default=None, required=False)
     main(vars(parser.parse_args()))
     # debug()
 

scripts/replay_episodes.py → aloha_scripts/replay_episodes.py

@@ -1,6 +1,6 @@
 import os
 import h5py
-from utils import move_grippers
+from robot_utils import move_grippers
 import argparse
 from real_env import make_real_env
 from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN