feat: reflect base state of skrobot robot model <=> kin

python/plainmp/robot_spec.py

@@ -15,6 +15,7 @@
     rpy_angle,
     rpy_matrix,
     wxyz2xyzw,
+    xyzw2wxyz,
 )
 from skrobot.model.primitives import Box, Cylinder, Sphere
 from skrobot.model.robot_model import RobotModel
@@ -140,6 +141,8 @@ def reflect_skrobot_model_to_kin(self, robot_model: RobotModel) -> None:
             angle = robot_model.__dict__[jn].joint_angle()
             table[jn] = angle
         self.reflect_joint_positions(table)
+        pose_vec = self.coordinates_to_pose_vec(robot_model.worldcoords(), RotType.XYZW)
+        self.get_kin().set_base_pose(pose_vec)
 
     def reflect_kin_to_skrobot_model(self, robot_model: RobotModel) -> None:
         # inverse of reflect_skrobot_model_to_kin
@@ -149,6 +152,11 @@ def reflect_kin_to_skrobot_model(self, robot_model: RobotModel) -> None:
         angles = kin.get_joint_positions(joint_ids)
         robot_model.angle_vector(angles)
 
+        pose_vec = kin.get_base_pose()
+        position, quat = pose_vec[:3], pose_vec[3:]
+        co = Coordinates(position, xyzw2wxyz(quat))
+        robot_model.newcoords(co)
+
     def extract_skrobot_model_q(
         self, robot_model: RobotModel, base_type: BaseType = BaseType.FIXED
     ) -> np.ndarray:
@@ -358,16 +366,7 @@ def create_pose_const_from_coords(
     ) -> LinkPoseCst:
         pose_list = []
         for co, rt in zip(link_poses, rot_types):
-            pos = co.worldpos()
-            if rt == RotType.RPY:
-                ypr = rpy_angle(co.rotation)[0]
-                rpy = [ypr[2], ypr[1], ypr[0]]
-                pose = np.hstack([pos, rpy])
-            elif rt == RotType.XYZW:
-                quat_wxyz = matrix2quaternion(co.rotation)
-                pose = np.hstack([pos, wxyz2xyzw(quat_wxyz)])
-            else:
-                pose = pos
+            pose = self.coordinates_to_pose_vec(co, rt)
             pose_list.append(pose)
         return self.create_pose_const(link_names, pose_list)
 
@@ -425,6 +424,26 @@ def create_attached_box_collision_const(
         )
         return cst
 
+    @staticmethod
+    def coordinates_to_pose_vec(co: Coordinates, rot_type: RotType) -> np.ndarray:
+        """convert skrobot coordinates to the pose vector
+        Args:
+            co: skrobot coordinates
+        Returns:
+            pose vector: [x, y, z, orientation...]
+        where orientation depends on the rot_type
+        """
+        pos = co.worldpos()
+        if rot_type == RotType.RPY:
+            ypr = rpy_angle(co.rotation)[0]
+            rpy = [ypr[2], ypr[1], ypr[0]]
+            return np.hstack([pos, rpy])
+        elif rot_type == RotType.XYZW:
+            quat_wxyz = matrix2quaternion(co.rotation)
+            return np.hstack([pos, wxyz2xyzw(quat_wxyz)])
+        else:
+            return pos
+
 
 class FetchSpec(RobotSpec):
     def __init__(self, base_type: BaseType = BaseType.FIXED):
@@ -547,18 +566,14 @@ def create_default_stand_pose_const(self) -> LinkPoseCst:
         # set reset manip pose
         for jn, angle in zip(self.control_joint_names, self.reset_manip_pose_q):
             robot_model.__dict__[jn].joint_angle(angle)
-
-        def skcoords_to_xyzrpy(co):
-            pos = co.worldpos()
-            ypr = rpy_angle(co.rotation)[0]
-            rpy = [ypr[2], ypr[1], ypr[0]]
-            return np.hstack([pos, rpy])
-
         rleg = robot_model.rleg_end_coords.copy_worldcoords()
         lleg = robot_model.lleg_end_coords.copy_worldcoords()
         return self.create_pose_const(
             ["rleg_end_coords", "lleg_end_coords"],
-            [skcoords_to_xyzrpy(rleg), skcoords_to_xyzrpy(lleg)],
+            [
+                self.coordinates_to_pose_vec(rleg, RotType.RPY),
+                self.coordinates_to_pose_vec(lleg, RotType.RPY),
+            ],
         )
 
     def create_default_com_const(

python/tests/test_kinematics.py

@@ -1,6 +1,11 @@
 import numpy as np
 import pytest
-from skrobot.coordinates.math import quaternion2matrix, xyzw2wxyz
+from skrobot.coordinates import Coordinates
+from skrobot.coordinates.math import (
+    quaternion2matrix,
+    rotation_matrix_from_quat,
+    xyzw2wxyz,
+)
 
 from plainmp.robot_spec import (
     FetchSpec,
@@ -60,3 +65,26 @@ def _test_get_link_pose(spec: RobotSpec):
 )
 def test_get_link_pose(spec):
     _test_get_link_pose(spec)
+
+
+def test_skrobot_kin_base_mirror():
+    spec = FetchSpec()
+    model = spec.get_robot_model()
+    kin = spec.get_kin()
+
+    # skrobot => kin
+    co = Coordinates([1, 2, 3], [0.1, 0.2, 0.3])
+    model.newcoords(co)
+    spec.reflect_skrobot_model_to_kin(model)
+    pose_vec = kin.get_base_pose()
+    pos, quat = pose_vec[:3], pose_vec[3:]
+    mat = rotation_matrix_from_quat(xyzw2wxyz(quat))
+    np.testing.assert_allclose(pos, co.worldpos(), atol=1e-3)
+    np.testing.assert_allclose(mat, co.worldrot(), atol=1e-3)
+
+    # kin => skrobot
+    kin.set_base_pose([4, 5, 6, 0, 0, 0, 1.0])  # whatever
+    spec.reflect_kin_to_skrobot_model(model)
+    pos, rotmat = model.worldpos(), model.worldrot()
+    np.testing.assert_allclose(pos, [4, 5, 6], atol=1e-3)
+    np.testing.assert_allclose(rotmat, np.eye(3), atol=1e-3)