当前位置：首页 > news >正文

MuJoCo 画出机械臂末端轨迹进行可视化（附代码）

news 来源：原创 2025/5/10 0:27:20

视频讲解：

上期视频我们使用matplotlib画出了路径规划及轨迹优化后的末端轨迹，有小伙伴留言说能不能在MuJoCo画出这个轨迹，粗看了下，官方文档中提供了这个方法，我们在上期视频基础上来添加这个功能。

仓库地址：GitHub - LitchiCheng/mujoco-learning

Python - MuJoCo Documentation

这个示例画了一堆球，颜色不一样，我们需要改一下

可以使用initGeom或者conector的方式

关键代码如下，将提取出来的末端tf赋值给pos，姿态的话就不显示，意义不大

self.handle.user_scn.ngeom = 0
i = 0
for i, tform in enumerate(self.tforms):position = tform.translationrotation_matrix = tform.rotationmujoco.mjv_initGeom(self.handle.user_scn.geoms[i],type=mujoco.mjtGeom.mjGEOM_SPHERE,size=[0.002, 0, 0],pos=np.array([tform.translation[0], tform.translation[1], tform.translation[2]]),mat=np.eye(3).flatten(),rgba=np.array([1, 0, 0, 1]))i += 1
self.handle.user_scn.ngeom = i

完整代码如下：

import mujoco_viewer
import mujoco,time,threading
import numpy as np
import pinocchio
import matplotlib
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import itertoolsfrom pyroboplan.core.utils import (get_random_collision_free_state,extract_cartesian_poses,
)
from pyroboplan.models.panda import (load_models,add_self_collisions,add_object_collisions,
)
from pyroboplan.planning.rrt import RRTPlanner, RRTPlannerOptions
from pyroboplan.trajectory.trajectory_optimization import (CubicTrajectoryOptimization,CubicTrajectoryOptimizationOptions,
)class Test(mujoco_viewer.CustomViewer):def __init__(self, path):super().__init__(path, 3, azimuth=180, elevation=-30)self.path = pathdef runBefore(self):# Create models and dataself.model_roboplan, self.collision_model, visual_model = load_models(use_sphere_collisions=True)add_self_collisions(self.model_roboplan, self.collision_model)add_object_collisions(self.model_roboplan, self.collision_model, visual_model, inflation_radius=0.1)data = self.model_roboplan.createData()collision_data = self.collision_model.createData()self.target_frame = "panda_hand"ignore_joint_indices = [self.model_roboplan.getJointId("panda_finger_joint1") - 1,self.model_roboplan.getJointId("panda_finger_joint2") - 1,]np.set_printoptions(precision=3)self.distance_padding = 0.001self.init_state = self.data.qpos.copy()while True:            q_start = self.random_valid_state()q_goal = self.random_valid_state()# Search for a pathoptions = RRTPlannerOptions(max_step_size=0.05,max_connection_dist=5.0,rrt_connect=False,bidirectional_rrt=True,rrt_star=True,max_rewire_dist=5.0,max_planning_time=20.0,fast_return=True,goal_biasing_probability=0.15,collision_distance_padding=0.01,)print("")print(f"Planning a path...")planner = RRTPlanner(self.model_roboplan, self.collision_model, options=options)q_path = planner.plan(q_start, q_goal)if len(q_path) > 0:print(f"Got a path with {len(q_path)} waypoints")else:print("Failed to plan.")# Perform trajectory optimization.dt = 0.025options = CubicTrajectoryOptimizationOptions(num_waypoints=len(q_path),samples_per_segment=7,min_segment_time=0.5,max_segment_time=10.0,min_vel=-1.5,max_vel=1.5,min_accel=-0.75,max_accel=0.75,min_jerk=-1.0,max_jerk=1.0,max_planning_time=30.0,check_collisions=True,min_collision_dist=self.distance_padding,collision_influence_dist=0.05,collision_avoidance_cost_weight=0.0,collision_link_list=["obstacle_box_1","obstacle_box_2","obstacle_sphere_1","obstacle_sphere_2","ground_plane","panda_hand",],)print("Optimizing the path...")optimizer = CubicTrajectoryOptimization(self.model_roboplan, self.collision_model, options)traj = optimizer.plan([q_path[0], q_path[-1]], init_path=q_path)if traj is None:print("Retrying with all the RRT waypoints...")traj = optimizer.plan(q_path, init_path=q_path)if traj is not None:print("Trajectory optimization successful")traj_gen = traj.generate(dt)self.q_vec = traj_gen[1]print(f"path has {self.q_vec.shape[1]} points")self.tforms = extract_cartesian_poses(self.model_roboplan, "panda_hand", self.q_vec.T)# 提取位置信息positions = []print(self.tforms[0].translation)print(self.tforms[0].rotation)self.handle.user_scn.ngeom = 0i = 0for i, tform in enumerate(self.tforms):position = tform.translationrotation_matrix = tform.rotationmujoco.mjv_initGeom(self.handle.user_scn.geoms[i],type=mujoco.mjtGeom.mjGEOM_SPHERE,size=[0.002, 0, 0],pos=np.array([tform.translation[0], tform.translation[1], tform.translation[2]]),mat=np.eye(3).flatten(),rgba=np.array([1, 0, 0, 1]))i += 1self.handle.user_scn.ngeom = iprint(f"Added {i} spheres to the scene.")for tform in self.tforms:position = tform.translationpositions.append(position)positions = np.array(positions)# 创建 3D 图形fig = plt.figure()ax = fig.add_subplot(111, projection='3d')# 绘制位置轨迹ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], marker='o')# 绘制姿态for i, tform in enumerate(self.tforms):position = tform.translationrotation_matrix = tform.rotation# 提取坐标轴方向的向量x_axis = rotation_matrix[:, 0]y_axis = rotation_matrix[:, 1]z_axis = rotation_matrix[:, 2]# 绘制坐标轴向量ax.quiver(position[0], position[1], position[2],x_axis[0], x_axis[1], x_axis[2], color='r', length=0.01)ax.quiver(position[0], position[1], position[2],y_axis[0], y_axis[1], y_axis[2], color='g', length=0.01)ax.quiver(position[0], position[1], position[2],z_axis[0], z_axis[1], z_axis[2], color='b', length=0.01)# 设置坐标轴标签ax.set_xlabel('X')ax.set_ylabel('Y')ax.set_zlabel('Z')# 显示图形plt.show(block=False)plt.pause(0.001)breakself.index = 0def random_valid_state(self):return get_random_collision_free_state(self.model_roboplan, self.collision_model, distance_padding=0.01)def runFunc(self):self.data.qpos[:7] = self.q_vec[:7, self.index]self.index += 1if self.index >= self.q_vec.shape[1]:self.index = 0time.sleep(0.01)if __name__ == "__main__":test = Test("/home/dar/MuJoCoBin/mujoco_menagerie/franka_emika_panda/scene.xml")test.run_loop()