Add wing joint transformation.

flybody/tasks/task_utils.py

@@ -25,6 +25,37 @@ def observable_indices_in_tensor(
     return sorted_obs_dict
 
 
+def wing_qpos_to_conventional(model_wing_qpos: np.ndarray,
+                              body_pitch_angle: float = 47.5,
+                             ) -> np.ndarray:
+    """Transform model wing qpos to conventional wing kinematics definition.
+    
+    Args:
+        model_wing_qpos: Wing MjData.qpos in radians, shape (B, 6).
+            Order of angles: yaw_left, roll_left, pitch_left,
+                             yaw_right, roll_right, pitch_right.
+        body_pitch_angle: Body pitch angle for initial flight pose, relative to
+            ground, degrees. 0: horizontal body position. Default value from
+            https://doi.org/10.1126/science.1248955
+
+    Returns:
+        Wing angles transformed to conventional representation.
+    """
+    if not isinstance(model_wing_qpos, np.ndarray):
+        model_wing_qpos = np.array(model_wing_qpos)
+    conventional = np.zeros_like(model_wing_qpos)
+    body_pitch_angle = np.deg2rad(body_pitch_angle)
+    for i in [0, 3]:
+        # Yaw, doesn't require transformation.
+        conventional[..., i] = model_wing_qpos[..., i].copy()
+        # Roll.
+        conventional[..., i+1] = - model_wing_qpos[..., i+1]
+        # Pitch.
+        conventional[..., i+2] = (np.pi / 2 - body_pitch_angle -
+                                  model_wing_qpos[..., i+2])
+    return conventional
+
+
 def get_random_policy(action_spec: 'dm_env.specs.BoundedArray',
                       minimum: float = -0.2,
                       maximum: float = 0.2) -> Callable[[Any], np.ndarray]: