debug linear speed, tune open-loop replay_episodes

aloha_scripts/real_env.py

@@ -3,6 +3,7 @@
 import collections
 import matplotlib.pyplot as plt
 import dm_env
+from pyquaternion import Quaternion
 
 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
@@ -104,7 +105,14 @@ def get_base_vel(self):
         frames = self.pipeline.wait_for_frames()
         pose_frame = frames.get_pose_frame()
         pose = pose_frame.get_pose_data()
-        base_linear_vel = pose.velocity.z
+
+        q1 = Quaternion(w=pose.rotation.w, x=pose.rotation.x, y=pose.rotation.y, z=pose.rotation.z)
+        rotation = -np.array(q1.yaw_pitch_roll)[0]
+        rotation_vec = np.array([np.cos(rotation), np.sin(rotation)])
+        linear_vel_vec = np.array([pose.velocity.z, pose.velocity.x])
+        is_forward = rotation_vec.dot(linear_vel_vec) > 0
+
+        base_linear_vel = np.sqrt(pose.velocity.z ** 2 + pose.velocity.x ** 2) * (1 if is_forward else -1)
         base_angular_vel = pose.angular_velocity.y
         return np.array([base_linear_vel, base_angular_vel])
 
@@ -171,7 +179,6 @@ def step(self, action, base_action=None):
             discount=None,
             observation=self.get_observation())
 
-
 def get_action(master_bot_left, master_bot_right):
     action = np.zeros(14) # 6 joint + 1 gripper, for two arms
     # Arm actions

aloha_scripts/replay_episodes.py

@@ -1,6 +1,6 @@
 import os
 import h5py
-from robot_utils import move_grippers, calibrate_linear_vel, smooth_base_action
+from robot_utils import move_grippers, calibrate_linear_vel, smooth_base_action, postprocess_base_action
 import argparse
 from real_env import make_real_env
 from constants import JOINT_NAMES, PUPPET_GRIPPER_JOINT_OPEN
@@ -26,9 +26,10 @@ def main(args):
 
     env = make_real_env(init_node=True, setup_base=True)
     env.reset()
-    # base_actions = smooth_base_action(base_actions)
+    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)
+        base_action = calibrate_linear_vel(base_action, c=0.19)
+        # base_action = postprocess_base_action(base_action)
         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

aloha_scripts/robot_utils.py

@@ -188,7 +188,7 @@ def torque_on(bot):
 
 def calibrate_linear_vel(base_action, c=None):
     if c is None:
-        c = 0.19
+        c = 0.
     v = base_action[..., 0]
     w = base_action[..., 1]
     base_action = base_action.copy()
@@ -197,5 +197,10 @@ def calibrate_linear_vel(base_action, c=None):
 
 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])
+        np.convolve(base_action[:, i], np.ones(5)/5, mode='same') for i in range(base_action.shape[1])
     ], axis=-1).astype(np.float32)
+
+def postprocess_base_action(base_action):
+    linear_vel, angular_vel = base_action
+    angular_vel *= 0.9
+    return np.array([linear_vel, angular_vel])