Adding visualizers

src/comodo/visualizers/iDynTreeVisualizer.py

@@ -0,0 +1,60 @@
+import idyntree.bindings as iDynTree
+import numpy as np
+import time 
+
+class iDynTreeVisualizer(): 
+    def __init__(self, model_name) -> None:
+        self.model_name= model_name
+        pass  
+
+    def prepare_visualization(self):
+        self.viz = iDynTree.Visualizer()  
+        vizOpt = iDynTree.VisualizerOptions()
+        vizOpt.winWidth = 1500
+        vizOpt.winHeight = 1500 
+        self.viz.init(vizOpt)
+
+        self.env = self.viz.enviroment()
+        self.env.setElementVisibility('floor_grid',True)
+        self.env.setElementVisibility('world_frame',False)
+        self.viz.setColorPalette("meshcat")
+        self.env.setElementVisibility('world_frame',False)
+        self.frames = self.viz.frames()  
+        cam = self.viz.camera()
+        cam.setPosition(iDynTree.Position(0,3,1.2))
+        self.viz.camera().animator().enableMouseControl(True)
+
+    def add_model(self,robot_model, urdf_path):
+        mdlLoader = iDynTree.ModelLoader()
+        mdlLoader.loadReducedModelFromFile(urdf_path,robot_model.joint_name_list, robot_model.base_link)
+        self.viz.addModel(mdlLoader.model(),self.model_name)
+
+    def update_model(self,s, H_b): 
+        s_idyn = self.get_idyntree_joint_pos(s)
+        H_b_idyn = self.get_idyntree_transf(H_b)
+        self.viz.modelViz(self.model_name).setPositions(H_b_idyn,s_idyn)
+
+    def get_idyntree_transf(self,H): 
+        pos_iDynTree = iDynTree.Position()
+        R_iDynTree = iDynTree.Rotation()
+        R_iDynTree.FromPython(H[:3,:3])
+        for i in range(3):
+            pos_iDynTree.setVal(i,float(H[i,3]))
+            for j in range(3):
+                R_iDynTree.setVal(j,i,float(H[j,i]))
+        T = iDynTree.Transform()
+        T.setPosition(pos_iDynTree)
+        T.setRotation(R_iDynTree)
+        T.setPosition(pos_iDynTree)
+        return T
+
+    def visualize(self, time_step): 
+        time_now = time.time()
+        while(time.time()-time_now<time_step and self.viz.run()): 
+            self.viz.draw()
+    def get_idyntree_joint_pos(self,s):
+        N_DoF = len(s)
+        s_idyntree =iDynTree.VectorDynSize(N_DoF)
+        for i in range(N_DoF):
+            s_idyntree.setVal(i,s[i])
+        return s_idyntree

src/comodo/visualizers/mujocoVisualizer.py

@@ -0,0 +1,52 @@
+import mujoco
+import math
+import numpy as np
+import mujoco_viewer
+import scipy 
+
+class MujocoVisualizer():
+
+    def __init__(self) -> None:
+        pass 
+
+    def load_model(self, robot_model, s, H_b):
+        self.robot_model = robot_model
+        mujoco_xml = robot_model.get_mujoco_model()
+        self.model = mujoco.MjModel.from_xml_string(mujoco_xml)
+        self.data = mujoco.MjData(self.model)
+        self.set_joint_vector_in_mujoco(s)
+        xyz, quat = self.from_H_b_to_xyz_quat(H_b)
+        self.set_base_pose_in_mujoco(xyz, quat)
+        mujoco.mj_forward(self.model, self.data)
+        self.viewer = mujoco_viewer.MujocoViewer(self.model, self.data)
+
+    def set_base_pose_in_mujoco(self, xyz, quat):
+        indexes_joint = self.model.jnt_qposadr[1:]
+        # Extract quaternion components
+        self.data.qpos[3 : indexes_joint[0]] = quat[:]
+        self.data.qpos[:3] = xyz[:]
+
+    def set_joint_vector_in_mujoco(self, pos):
+        indexes_joint = self.model.jnt_qposadr[1:]
+        for i in range(self.robot_model.NDoF):
+            self.data.qpos[indexes_joint[i]] = pos[i]
+
+    def update_vis(self, s, H_b): 
+        xyz, quat = self.from_H_b_to_xyz_quat(H_b)
+        self.set_base_pose_in_mujoco(xyz, quat)
+        self.set_joint_vector_in_mujoco(s)
+        mujoco.mj_forward(self.model, self.data)
+        self.viewer.render()
+
+    def from_H_b_to_xyz_quat(self, H_b):
+        xyz = H_b[:3,3] 
+        rot_mat = scipy.spatial.transform.Rotation.from_matrix(H_b[:3,:3])
+        quat_xyzw = rot_mat.as_quat()
+        # note that for mujoco the ordering is w,x,y,z
+        quat_wxyz = np.asarray([quat_xyzw[3],quat_xyzw[0],quat_xyzw[1],quat_xyzw[2]])
+        return xyz, quat_wxyz
+    def close(self):
+            self.viewer.close()
+
+    def visualize_robot(self):
+        self.viewer.render()