Fix tcdm 1 (#66)

* FIX: when robot_vel reference is not present, set to zero

* FIX:  returns both magnitude and axis, for the distance we just need the magnitude

---------

Co-authored-by: Vittorio Caggiano <caggiano@fb.com>

robohive/envs/tcdm/reference_motion.py

@@ -43,6 +43,8 @@ def __init__(self, reference:Union[str, dict],
         # load reference
         if isinstance(reference, str):
             self.reference = self.load(reference)
+            if 'robot_vel' not in self.reference.keys():
+                self.reference['robot_vel'] = np.zeros(self.reference['robot'].shape)
         elif isinstance(reference, dict):
             self.reference = reference
         else:
@@ -94,6 +96,7 @@ def check_format(self, reference):
         """
         assert 'time' in reference.keys(), "Missing keys in reference"
         assert 'robot' in reference.keys(), "Missing keys in reference"
+        assert 'robot_vel' in reference.keys(), "Missing keys in reference"
         assert 'object' in reference.keys(), "Missing keys in reference"
         assert reference['robot'].ndim == 2, "Check robot reference, must be shape(N, n_robot_jnt)"
         assert reference['object'].ndim == 2, "Check object reference, must be shape(N, n_object_jnt)"

robohive/envs/tcdm/track.py

@@ -143,7 +143,8 @@ def rotation_distance(self, q1, q2, euler=True):
         if euler:
             q1 = euler2quat(q1)
             q2 = euler2quat(q2)
-        return np.abs(quatDiff2Vel(q2,q1,1))
+
+        return np.abs(quatDiff2Vel(q2,q1,1)[0])
 
 
     def update_reference_insim(self, curr_ref):
@@ -290,4 +291,4 @@ def check_termination(self, obs_dict):
         else:
             qpos_term = False
 
-        return obj_term or qpos_term or base_term # combining termination for object + posture
+        return obj_term or qpos_term or base_term # combining termination for object + posture