diff --git a/src/comodo/visualizers/__init__.py b/src/comodo/visualizers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/comodo/visualizers/iDynTreeVisualizer.py b/src/comodo/visualizers/iDynTreeVisualizer.py new file mode 100644 index 0000000..9a556d9 --- /dev/null +++ b/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 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() \ No newline at end of file