FPS set to 50 for data collection and deployment

aloha_scripts/constants.py

@@ -1,4 +1,5 @@
 ### Task parameters
+import pathlib
 import os
 
 DATA_DIR = os.path.expanduser('~/data')
@@ -45,22 +46,10 @@
         'episode_len': 1300,
         'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
     },
-    'aloha_mobile_wipe_wine_new':{
-        'dataset_dir': DATA_DIR + '/aloha_mobile_wipe_wine_new',
-        'num_episodes': 50,
-        'episode_len': 1300,
-        'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
-    },
-    'aloha_mobile_wipe_wine_new_30':{
-        'dataset_dir': DATA_DIR + '/aloha_mobile_wipe_wine_new_30',
-        'num_episodes': 50,
-        'episode_len': 1300,
-        'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
-    },
     'aloha_mobile_wash_pan':{
         'dataset_dir': DATA_DIR + '/aloha_mobile_wash_pan',
         'num_episodes': 50,
-        'episode_len': 1200,
+        'episode_len': 1100,
         'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
     },
     'aloha_mobile_dummy':{
@@ -79,9 +68,12 @@
 
 ### ALOHA fixed constants
 DT = 0.02
+FPS = 50
 JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"]
 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]
 
+XML_DIR = str(pathlib.Path(__file__).parent.resolve()) + '/assets/' # note: absolute path
+
 # Left finger position limits (qpos[7]), right_finger = -1 * left_finger
 MASTER_GRIPPER_POSITION_OPEN = 0.02417
 MASTER_GRIPPER_POSITION_CLOSE = 0.01244

aloha_scripts/real_env.py

@@ -196,7 +196,7 @@ def step(self, action, base_action=None, get_tracer_vel=False):
             # base_action_angular = np.clip(base_action[1], -angular_vel_limit, angular_vel_limit)
             base_action_linear, base_action_angular = base_action
             self.tracer.SetMotionCommand(linear_vel=base_action_linear, angular_vel=base_action_angular)
-        time.sleep(DT)
+        # time.sleep(DT)
         return dm_env.TimeStep(
             step_type=dm_env.StepType.MID,
             reward=self.get_reward(),

aloha_scripts/record_episodes.py

@@ -7,7 +7,7 @@
 from tqdm import tqdm
 import cv2
 
-from constants import DT, START_ARM_POSE, TASK_CONFIGS
+from constants import DT, START_ARM_POSE, TASK_CONFIGS, FPS
 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
@@ -92,6 +92,8 @@ def capture_one_episode(dt, max_timesteps, camera_names, dataset_dir, dataset_na
     timesteps = [ts]
     actions = []
     actual_dt_history = []
+    time0 = time.time()
+    DT = 1 / FPS
     for t in tqdm(range(max_timesteps)):
         t0 = time.time() #
         action = get_action(master_bot_left, master_bot_right)
@@ -101,6 +103,8 @@ def capture_one_episode(dt, max_timesteps, camera_names, dataset_dir, dataset_na
         timesteps.append(ts)
         actions.append(action)
         actual_dt_history.append([t0, t1, t2])
+        time.sleep(DT - (time.time() - t0))
+    print(f'Avg fps: {max_timesteps / (time.time() - time0)}')
 
     # Torque on both master bots
     torque_on(master_bot_left)

aloha_scripts/replay_and_record_episodes.py

@@ -5,7 +5,8 @@
 import argparse
 import matplotlib.pyplot as plt
 from real_env import make_real_env
-from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN
+from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN, fps
+import time
 
 import IPython
 e = IPython.embed
@@ -66,12 +67,15 @@ def main(args):
     apply_actions = actions
     apply_base_actions = base_actions[offset:] * scale
 
+    DT = 1 / fps
     for action, base_action in zip(apply_actions, apply_base_actions):
+        time1 = time.time()
         # base_action = calibrate_linear_vel(base_action, c=0.19)
         # base_action = postprocess_base_action(base_action)
         ts = env.step(action, base_action, get_tracer_vel=True)
         obs_wheels.append(ts.observation['base_vel'])
         obs_tracer.append(ts.observation['tracer_vel'])
+        time.sleep(max(0, DT - (time.time() - time1)))
     obs_wheels = np.array(obs_wheels)
     obs_tracer = np.array(obs_tracer)
 

aloha_scripts/replay_episodes.py

@@ -5,7 +5,8 @@
 import argparse
 import matplotlib.pyplot as plt
 from real_env import make_real_env
-from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN
+from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN, FPS
+import time
 
 import IPython
 e = IPython.embed
@@ -32,7 +33,7 @@ def main(args):
         base_actions = root['/base_action'][()]
 
     if use_actuator_net:
-        from act.train_actuator_network import ActuatorNetwork
+        from train_actuator_network import ActuatorNetwork
         import torch
         import pickle
         actuator_network = ActuatorNetwork(prediction_len)
@@ -67,6 +68,14 @@ def main(args):
 
         processed_base_actions = np.array(processed_base_actions)
         assert processed_base_actions.shape == base_actions.shape
+
+        plt.plot(base_actions[:, 0], label='action_linear')
+        plt.plot(processed_base_actions[:, 0], '--', label='processed_action_linear')
+        plt.plot(base_actions[:, 1], label='action_angular')
+        plt.plot(processed_base_actions[:, 1], '--', label='processed_action_angular')
+        plt.plot()
+        plt.legend()
+        plt.show()
     else:
         # processed_base_actions = smooth_base_action(base_actions)
         processed_base_actions = base_actions
@@ -76,25 +85,31 @@ def main(args):
     obs_wheels = []
     obs_tracer = []
 
+    time0 = time.time()
+    DT = 1 / FPS
     for action, base_action in zip(actions, processed_base_actions):
+        time1 = time.time()
         # base_action = calibrate_linear_vel(base_action, c=0.19)
         # base_action = postprocess_base_action(base_action)
         ts = env.step(action, base_action, get_tracer_vel=True)
         obs_wheels.append(ts.observation['base_vel'])
         obs_tracer.append(ts.observation['tracer_vel'])
+        time.sleep(max(0, DT - (time.time() - time1)))
+    print(f'Avg fps: {len(actions) / (time.time() - time0)}')
     obs_wheels = np.array(obs_wheels)
     obs_tracer = np.array(obs_tracer)
 
     plt.plot(base_actions[:, 0], label='action_linear')
-    plt.plot(processed_base_actions[:, 0], '-.-', label='processed_action_linear')
+    plt.plot(processed_base_actions[:, 0], '--', label='processed_action_linear')
     plt.plot(obs_wheels[:, 0], '--', label='obs_wheels_linear')
     plt.plot(obs_tracer[:, 0], '-.', label='obs_tracer_linear')
     plt.plot()
     plt.legend()
     plt.savefig('replay_episodes_linear_vel.png', dpi=300)
-
+
+    plt.clf()
     plt.plot(base_actions[:, 1], label='action_angular')
-    plt.plot(processed_base_actions[:, 1], '-.-', label='processed_action_angular')
+    plt.plot(processed_base_actions[:, 1], '--', label='processed_action_angular')
     plt.plot(obs_wheels[:, 1], '--', label='obs_wheels_angular')
     plt.plot(obs_tracer[:, 1], '-.', label='obs_tracer_angular')
     plt.plot()