feat: add spec.set_skrobot_model_state

example/attach_box_ik_and_plan.py

@@ -11,7 +11,7 @@
 from plainmp.problem import Problem
 from plainmp.psdf import GroundSDF, UnionSDF
 from plainmp.robot_spec import PandaSpec
-from plainmp.utils import box_to_grid_poitns, primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import box_to_grid_poitns, primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -77,7 +77,7 @@
     if args.visualize:
         # visualization
         robot = spec.get_robot_model(with_mesh=True)
-        set_robot_state(robot, spec.control_joint_names, ik_res.q)
+        spec.set_skrobot_model_state(robot, ik_res.q)
         ground_vis = Box([2.0, 2.0, 0.03])
         v = PyrenderViewer()
         v.add(robot)
@@ -90,7 +90,7 @@
         robot.__dict__["panda_hand"].assoc(cardboard_box)
 
         for q in result.traj.resample(30):
-            set_robot_state(robot, spec.control_joint_names, q)
+            spec.set_skrobot_model_state(robot, q)
             v.redraw()
             time.sleep(0.1)
         time.sleep(1000)

example/beta/humanoid_lift_and_place_box.py

@@ -8,14 +8,13 @@
 
 from plainmp.constraint import EqCompositeCst, IneqCompositeCst
 from plainmp.ik import solve_ik
-from plainmp.kinematics import BaseType
 from plainmp.manifold_rrt.manifold_rrt_solver import ManiRRTConfig, ManiRRTConnectSolver
 from plainmp.nlp_solver import SQPBasedSolver, SQPBasedSolverConfig
 from plainmp.parallel import ParallelSolver
 from plainmp.problem import Problem
 from plainmp.psdf import UnionSDF
 from plainmp.robot_spec import JaxonSpec, RotType
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -113,7 +112,7 @@
         v = PyrenderViewer()
         ground = Box([2, 2, 0.01])
         robot = jspec.get_robot_model(with_mesh=True)
-        set_robot_state(robot, jspec.control_joint_names, ik_ret1.q, BaseType.FLOATING)
+        jspec.set_skrobot_model_state(robot, ik_ret1.q)
         robot.rarm_end_coords.assoc(box)
 
         v.add(ground)
@@ -124,7 +123,7 @@
         time.sleep(2)
 
         for q in ret.traj:
-            set_robot_state(robot, jspec.control_joint_names, q, BaseType.FLOATING)
+            jspec.set_skrobot_model_state(robot, q)
             v.redraw()
             time.sleep(0.1)
         time.sleep(10)

example/beta/humanoid_planning.py

@@ -8,13 +8,12 @@
 
 from plainmp.constraint import EqCompositeCst, IneqCompositeCst
 from plainmp.ik import solve_ik
-from plainmp.kinematics import BaseType
 from plainmp.manifold_rrt.manifold_rrt_solver import ManiRRTConfig, ManiRRTConnectSolver
 from plainmp.nlp_solver import SQPBasedSolver, SQPBasedSolverConfig
 from plainmp.parallel import ParallelSolver
 from plainmp.problem import Problem
 from plainmp.robot_spec import JaxonSpec, RotType
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -80,15 +79,15 @@
     if args.visualize:
         v = PyrenderViewer()
         robot = jspec.get_robot_model(with_mesh=True)
-        set_robot_state(robot, jspec.control_joint_names, ik_ret.q, BaseType.FLOATING)
+        jspec.set_skrobot_model_state(robot, ik_ret.q)
         ground = Box([2, 2, 0.01])
         v.add(ground)
         v.add(table)
         v.add(robot)
         v.show()
         time.sleep(2)
         for q in ret.traj:
-            set_robot_state(robot, jspec.control_joint_names, q, BaseType.FLOATING)
+            jspec.set_skrobot_model_state(robot, q)
             v.redraw()
             time.sleep(0.15)
         time.sleep(20)

example/fetch_ik.py

@@ -9,7 +9,7 @@
 from plainmp.ik import solve_ik
 from plainmp.psdf import CloudSDF
 from plainmp.robot_spec import FetchSpec
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -45,7 +45,7 @@
 
     if args.visualize:
         fetch = Fetch()
-        set_robot_state(fetch, fs.control_joint_names, ret.q)
+        fs.set_skrobot_model_state(fetch, ret.q)
         v = PyrenderViewer()
         v.add(fetch)
         v.add(table)

example/fetch_plan.py

@@ -9,7 +9,7 @@
 from plainmp.problem import Problem
 from plainmp.psdf import CloudSDF, GroundSDF, UnionSDF
 from plainmp.robot_spec import FetchSpec
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -48,7 +48,7 @@
         robot_model = fs.get_robot_model(with_mesh=True)
 
         # get goal point
-        set_robot_state(robot_model, fs.control_joint_names, q_goal)
+        fs.set_skrobot_model_state(robot_model, q_goal)
         co = robot_model.gripper_link.copy_worldcoords()
 
         v = PyrenderViewer()
@@ -59,7 +59,7 @@
         v.show()
         input("Press Enter to replay the path")
         for q in ret.traj.resample(50):
-            set_robot_state(robot_model, fs.control_joint_names, q)
+            fs.set_skrobot_model_state(robot_model, q)
             v.redraw()
             time.sleep(0.2)
         time.sleep(1000)

example/misc/fetch_visualize_coll_spheres.py

@@ -6,15 +6,14 @@
 from skrobot.viewers import PyrenderViewer
 
 from plainmp.robot_spec import FetchSpec
-from plainmp.utils import set_robot_state
 
 fs = FetchSpec()
 cst = fs.create_collision_const()
 q = np.zeros(8)
 cst.update_kintree(q, True)
 
 fetch = Fetch()
-set_robot_state(fetch, fs.control_joint_names, q)
+fs.set_skrobot_model_state(fetch, q)
 v = PyrenderViewer()
 v.add(fetch)
 for center, r in cst.get_group_spheres():

example/misc/panda_visualize_coll_spheres.py

@@ -6,7 +6,6 @@
 from skrobot.viewers import PyrenderViewer
 
 from plainmp.robot_spec import PandaSpec
-from plainmp.utils import set_robot_state
 
 s = PandaSpec()
 cst = s.create_collision_const()
@@ -16,7 +15,7 @@
 q = np.random.uniform(lb, ub)
 cst.update_kintree(q, True)
 model = RobotModelFromURDF(urdf_file=str(s.urdf_path))
-set_robot_state(model, s.control_joint_names, q)
+s.set_skrobot_model_state(model, q)
 
 sk_all_spheres = []
 for center, r in cst.get_all_spheres():

example/misc/pr2_visualize_coll_spheres.py

@@ -4,17 +4,18 @@
 from skrobot.model.primitives import Sphere
 from skrobot.viewers import PyrenderViewer
 
-from plainmp.robot_spec import PR2RarmSpec
-from plainmp.utils import set_robot_state
+from plainmp.robot_spec import PR2BaseOnlySpec
 
-s = PR2RarmSpec()
+# s = PR2RarmSpec()
+s = PR2BaseOnlySpec()
 cst = s.create_collision_const()
 
 lb, ub = s.angle_bounds()
 q = np.random.uniform(lb, ub)
+q = np.zeros(3)
 cst.update_kintree(q, True)
 model = s.get_robot_model(with_mesh=True)
-set_robot_state(model, s.control_joint_names, q)
+s.set_skrobot_model_state(model, q)
 
 sk_all_spheres = []
 for center, r in cst.get_all_spheres():

example/panda_plan.py

@@ -9,7 +9,7 @@
 from plainmp.problem import Problem
 from plainmp.psdf import UnionSDF
 from plainmp.robot_spec import PandaSpec
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -52,7 +52,8 @@
     if args.visualize:
         v = PyrenderViewer()
         robot_model = spec.get_robot_model(with_mesh=True)
-        set_robot_state(robot_model, spec.control_joint_names, q1)
+        spec.set_skrobot_model_state(robot_model, q1)
+
         co = robot_model.panda_hand.copy_worldcoords()
         v.add(Sphere(0.05, color=[255, 0, 0]).translate(co.worldpos()))
 
@@ -62,7 +63,7 @@
         v.show()
         input("Press Enter to replay the path")
         for q in ret.traj.resample(50):
-            set_robot_state(robot_model, spec.control_joint_names, q)
+            spec.set_skrobot_model_state(robot_model, q)
             v.redraw()
             time.sleep(0.2)
         time.sleep(1000)

example/pr2_dualarm_ik_and_plan.py

@@ -10,7 +10,7 @@
 from plainmp.ompl_solver import OMPLSolver, OMPLSolverConfig
 from plainmp.problem import Problem
 from plainmp.robot_spec import PR2DualarmSpec
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -63,15 +63,15 @@
         robot_model = spec.get_robot_model(with_mesh=True)
         for name, angle in default_joint_positions.items():
             robot_model.__dict__[name].joint_angle(angle)
-        set_robot_state(robot_model, spec.control_joint_names, ik_result.q)
+        spec.set_skrobot_model_state(robot_model, ik_result.q)
         viewer.add(Axis.from_coords(target_coords_rarm))
         viewer.add(Axis.from_coords(target_coords_larm))
         viewer.add(robot_model)
         viewer.add(table)
         viewer.show()
         input("Press Enter to show the planned path")
         for q in mp_result.traj.resample(50):
-            set_robot_state(robot_model, spec.control_joint_names, q)
+            spec.set_skrobot_model_state(robot_model, q)
             viewer.redraw()
             time.sleep(0.15)
         time.sleep(1000)

example/pr2_rarm_ik_and_plan.py

@@ -10,7 +10,7 @@
 from plainmp.ompl_solver import OMPLSolver, OMPLSolverConfig
 from plainmp.problem import Problem
 from plainmp.robot_spec import PR2RarmSpec
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -58,14 +58,14 @@
         robot_model = spec.get_robot_model(with_mesh=True)
         for name, angle in default_joint_positions.items():
             robot_model.__dict__[name].joint_angle(angle)
-        set_robot_state(robot_model, spec.control_joint_names, ik_result.q)
+        spec.set_skrobot_model_state(robot_model, ik_result.q)
         viewer.add(Axis.from_coords(target_coords))
         viewer.add(robot_model)
         viewer.add(table)
         viewer.show()
         input("Press Enter to show the planned path")
         for q in mp_result.traj.resample(50):
-            set_robot_state(robot_model, spec.control_joint_names, q)
+            spec.set_skrobot_model_state(robot_model, q)
             viewer.redraw()
             time.sleep(0.15)
         time.sleep(1000)

example/pr2_with_mobile_base.py

@@ -10,7 +10,7 @@
 from plainmp.ompl_solver import OMPLSolver, OMPLSolverConfig
 from plainmp.problem import Problem
 from plainmp.robot_spec import PR2RarmSpec
-from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
+from plainmp.utils import primitive_to_plainmp_sdf
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -50,15 +50,15 @@
         robot_model = spec.get_robot_model(with_mesh=True)
         for name, angle in default_joint_positions.items():
             robot_model.__dict__[name].joint_angle(angle)
-        set_robot_state(robot_model, spec.control_joint_names, q_start, base_type=BaseType.PLANAR)
+        spec.set_skrobot_model_state(robot_model, q_start)
         co = Coordinates([0.9, -0.2, 0.9])
         viewer.add(Axis.from_coords(co))
         viewer.add(robot_model)
         viewer.add(table)
         viewer.show()
         input("Press Enter to show the planned path")
         for q in mp_result.traj.resample(50):
-            set_robot_state(robot_model, spec.control_joint_names, q, base_type=BaseType.PLANAR)
+            spec.set_skrobot_model_state(robot_model, q)
             viewer.redraw()
             time.sleep(0.15)
         time.sleep(1000)

python/plainmp/robot_spec.py

@@ -34,7 +34,7 @@
 )
 from plainmp.kinematics import BaseType, KinematicModel
 from plainmp.psdf import BoxSDF, Pose, UnionSDF
-from plainmp.utils import primitive_to_plainmp_sdf
+from plainmp.utils import primitive_to_plainmp_sdf, set_robot_state
 
 _loaded_urdf_models: Dict[str, URDF] = {}
 _loaded_yamls: Dict[str, Dict] = {}  # loading yaml is quite slow
@@ -157,6 +157,9 @@ def reflect_kin_to_skrobot_model(self, robot_model: RobotModel) -> None:
         co = Coordinates(position, xyzw2wxyz(quat))
         robot_model.newcoords(co)
 
+    def set_skrobot_model_state(self, robot_model: RobotModel, q: np.ndarray) -> None:
+        set_robot_state(robot_model, self.control_joint_names, q, base_type=self.base_type)
+
     def extract_skrobot_model_q(
         self, robot_model: RobotModel, base_type: BaseType = BaseType.FIXED
     ) -> np.ndarray: