add more tasks to constants, tweek sleep arm pose

agilexrobotics · Nov 1, 2023 · cdd65c4 · cdd65c4
1 parent 4832ecd
commit cdd65c4
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 4 deletions.
diff --git a/aloha_scripts/constants.py b/aloha_scripts/constants.py
@@ -21,6 +21,24 @@
         'episode_len': 500,
         'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
     },
+    'aloha_mobile_object_to_cabinet':{
+        'dataset_dir': DATA_DIR + '/aloha_mobile_object_to_cabinet',
+        'num_episodes': 50,
+        'episode_len': 1700,
+        'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
+    },
+    'aloha_mobile_glass_to_cabinet':{
+        'dataset_dir': DATA_DIR + '/aloha_mobile_glass_to_cabinet',
+        'num_episodes': 50,
+        'episode_len': 1500,
+        'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
+    },
+    'aloha_mobile_wipe_wine':{
+        'dataset_dir': DATA_DIR + '/aloha_mobile_wipe_wine',
+        'num_episodes': 50,
+        'episode_len': 1300,
+        'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
+    },
 }
 
 ### ALOHA fixed constants

diff --git a/aloha_scripts/replay_episodes.py b/aloha_scripts/replay_episodes.py
@@ -1,6 +1,6 @@
 import os
 import h5py
-from robot_utils import move_grippers
+from robot_utils import move_grippers, calibrate_linear_vel, smooth_base_action
 import argparse
 from real_env import make_real_env
 from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN
@@ -26,7 +26,9 @@ def main(args):
 
     env = make_real_env(init_node=True, setup_base=True)
     env.reset()
+    # base_actions = smooth_base_action(base_actions)
     for action, base_action in zip(actions, base_actions):
+        # base_action = calibrate_linear_vel(base_action, c=0)
         env.step(action, base_action)
 
     move_grippers([env.puppet_bot_left, env.puppet_bot_right], [PUPPET_GRIPPER_JOINT_OPEN] * 2, move_time=0.5)  # open

diff --git a/aloha_scripts/robot_utils.py b/aloha_scripts/robot_utils.py
@@ -185,3 +185,17 @@ def torque_off(bot):
 def torque_on(bot):
     bot.dxl.robot_torque_enable("group", "arm", True)
     bot.dxl.robot_torque_enable("single", "gripper", True)
+
+def calibrate_linear_vel(base_action, c=None):
+    if c is None:
+        c = 0.19
+    v = base_action[..., 0]
+    w = base_action[..., 1]
+    base_action = base_action.copy()
+    base_action[..., 0] = v - c * w
+    return base_action
+
+def smooth_base_action(base_action):
+    return np.stack([
+        np.convolve(base_action[:, i], np.ones(20)/20, mode='same') for i in range(base_action.shape[1])
+    ], axis=-1).astype(np.float32)
diff --git a/aloha_scripts/sleep.py b/aloha_scripts/sleep.py
@@ -17,15 +17,15 @@ def main():
     for bot in all_bots:
         torque_on(bot)
 
-    puppet_sleep_position = (0, -1.7, 1.55, 0.12, 0.65, 0)
+    puppet_sleep_position = (0, -1.7, 1.55, 0, 0.65, 0)
     master_sleep_left_position = (-0.61, 0., 0.43, 0., 1.04, -0.65)
     master_sleep_right_position = (0.61, 0., 0.43, 0., 1.04, 0.65)
     all_positions = [puppet_sleep_position] * 2 + [master_sleep_left_position, master_sleep_right_position] if args.all else [puppet_sleep_position] * 2
     move_arms(all_bots, all_positions, move_time=2)
 
     if args.all:
-        master_sleep_left_position_2 = (0., 0.66, -0.27, -0.055, 1.1, 0)
-        master_sleep_right_position_2 = (0., 0.66, -0.27, -0.055, 1.1, 0)
+        master_sleep_left_position_2 = (0., 0.66, -0.27, -0.0, 1.1, 0)
+        master_sleep_right_position_2 = (0., 0.66, -0.27, -0.0, 1.1, 0)
         move_arms(master_bots, [master_sleep_left_position_2, master_sleep_right_position_2], move_time=1)