456 lines
16 KiB
Python
456 lines
16 KiB
Python
import pybullet as p
|
||
import pybullet_data
|
||
import numpy as np
|
||
import os
|
||
import time
|
||
from PytorchBoot.runners.runner import Runner
|
||
import PytorchBoot.stereotype as stereotype
|
||
from PytorchBoot.config import ConfigManager
|
||
from utils.control import ControlUtil
|
||
|
||
|
||
@stereotype.runner("simulator")
|
||
class Simulator(Runner):
|
||
CREATE: str = "create"
|
||
SIMULATE: str = "simulate"
|
||
INIT_GRIPPER_POSE:np.ndarray = np.asarray(
|
||
[[0.41869126 ,0.87596275 , 0.23951774 , 0.36005292],
|
||
[ 0.70787907 ,-0.4800251 , 0.51813998 ,-0.40499909],
|
||
[ 0.56884584, -0.04739109 ,-0.82107382 ,0.76881103],
|
||
[ 0. , 0. , 0. , 1. ]])
|
||
TURNTABLE_WORLD_TO_PYBULLET_WORLD:np.ndarray = np.asarray(
|
||
[[1, 0, 0, 0.8],
|
||
[0, 1, 0, 0],
|
||
[0, 0, 1, 0.5],
|
||
[0, 0, 0, 1]])
|
||
|
||
debug_pose = np.asarray([
|
||
[
|
||
0.992167055606842,
|
||
-0.10552699863910675,
|
||
0.06684812903404236,
|
||
-0.07388903945684433
|
||
],
|
||
[
|
||
0.10134342312812805,
|
||
0.3670985698699951,
|
||
-0.9246448874473572,
|
||
-0.41582486033439636
|
||
],
|
||
[
|
||
0.07303514331579208,
|
||
0.9241767525672913,
|
||
0.37491756677627563,
|
||
1.0754833221435547
|
||
],
|
||
[
|
||
0.0,
|
||
0.0,
|
||
0.0,
|
||
1.0
|
||
]])
|
||
|
||
def __init__(self, config_path):
|
||
super().__init__(config_path)
|
||
self.config_path = config_path
|
||
self.robot_id = None
|
||
self.turntable_id = None
|
||
self.target_id = None
|
||
camera_config = ConfigManager.get("simulation", "camera")
|
||
self.camera_params = {
|
||
'width': camera_config["width"],
|
||
'height': camera_config["height"],
|
||
'fov': camera_config["fov"],
|
||
'near': camera_config["near"],
|
||
'far': camera_config["far"]
|
||
}
|
||
self.sim_config = ConfigManager.get("simulation")
|
||
|
||
def run(self, cmd):
|
||
print(f"Simulator run {cmd}")
|
||
if cmd == self.CREATE:
|
||
self.prepare_env()
|
||
self.create_env()
|
||
elif cmd == self.SIMULATE:
|
||
self.simulate()
|
||
|
||
def simulate(self):
|
||
self.reset()
|
||
self.init()
|
||
debug_pose = Simulator.debug_pose
|
||
offset = np.asarray([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
|
||
debug_pose = debug_pose @ offset
|
||
for _ in range(10000):
|
||
debug_pose_2 = np.eye(4)
|
||
debug_pose_2[0,0] = -1
|
||
debug_pose_2[2,3] = 0.5
|
||
self.move_to(debug_pose_2)
|
||
# Wait for the system to stabilize
|
||
for _ in range(20): # Simulate 20 steps to ensure stability
|
||
p.stepSimulation()
|
||
time.sleep(0.001) # Add small delay to ensure physics simulation
|
||
|
||
depth_img, segm_img = self.take_picture()
|
||
p.stepSimulation()
|
||
|
||
def prepare_env(self):
|
||
p.connect(p.GUI)
|
||
p.setAdditionalSearchPath(pybullet_data.getDataPath())
|
||
p.setGravity(0, 0, 0)
|
||
p.loadURDF("plane.urdf")
|
||
|
||
def create_env(self):
|
||
print(self.config)
|
||
robot_config = self.sim_config["robot"]
|
||
turntable_config = self.sim_config["turntable"]
|
||
target_config = self.sim_config["target"]
|
||
|
||
self.robot_id = p.loadURDF(
|
||
robot_config["urdf_path"],
|
||
robot_config["initial_position"],
|
||
p.getQuaternionFromEuler(robot_config["initial_orientation"]),
|
||
useFixedBase=True
|
||
)
|
||
|
||
p.changeDynamics(
|
||
self.robot_id,
|
||
linkIndex=-1,
|
||
mass=0,
|
||
linearDamping=0,
|
||
angularDamping=0,
|
||
lateralFriction=0
|
||
)
|
||
|
||
visual_shape_id = p.createVisualShape(
|
||
shapeType=p.GEOM_CYLINDER,
|
||
radius=turntable_config["radius"],
|
||
length=turntable_config["height"],
|
||
rgbaColor=[0.7, 0.7, 0.7, 1]
|
||
)
|
||
collision_shape_id = p.createCollisionShape(
|
||
shapeType=p.GEOM_CYLINDER,
|
||
radius=turntable_config["radius"],
|
||
height=turntable_config["height"]
|
||
)
|
||
self.turntable_id = p.createMultiBody(
|
||
baseMass=0, # 设置质量为0使其成为静态物体
|
||
baseCollisionShapeIndex=collision_shape_id,
|
||
baseVisualShapeIndex=visual_shape_id,
|
||
basePosition=turntable_config["center_position"]
|
||
)
|
||
|
||
# 禁用转盘的动力学
|
||
p.changeDynamics(
|
||
self.turntable_id,
|
||
-1, # -1 表示基座
|
||
mass=0,
|
||
linearDamping=0,
|
||
angularDamping=0,
|
||
lateralFriction=0
|
||
)
|
||
|
||
|
||
obj_path = os.path.join(target_config["obj_dir"], target_config["obj_name"], "mesh.obj")
|
||
|
||
assert os.path.exists(obj_path), f"Error: File not found at {obj_path}"
|
||
|
||
# 加载OBJ文件作为目标物体
|
||
target_visual = p.createVisualShape(
|
||
shapeType=p.GEOM_MESH,
|
||
fileName=obj_path,
|
||
rgbaColor=target_config["rgba_color"],
|
||
specularColor=[0.4, 0.4, 0.4],
|
||
meshScale=[target_config["scale"]] * 3
|
||
)
|
||
|
||
# 使用简化的碰撞形状
|
||
target_collision = p.createCollisionShape(
|
||
shapeType=p.GEOM_MESH,
|
||
fileName=obj_path,
|
||
meshScale=[target_config["scale"]] * 3,
|
||
flags=p.GEOM_FORCE_CONCAVE_TRIMESH # 尝试使用凹面网格
|
||
)
|
||
|
||
|
||
# 创建目标物体
|
||
self.target_id = p.createMultiBody(
|
||
baseMass=0, # 设置质量为0使其成为静态物体
|
||
baseCollisionShapeIndex=target_collision,
|
||
baseVisualShapeIndex=target_visual,
|
||
basePosition=[
|
||
turntable_config["center_position"][0],
|
||
turntable_config["center_position"][1],
|
||
turntable_config["height"] + turntable_config["center_position"][2]
|
||
],
|
||
baseOrientation=p.getQuaternionFromEuler([np.pi/2, 0, 0])
|
||
)
|
||
|
||
# 禁用目标物体的动力学
|
||
p.changeDynamics(
|
||
self.target_id,
|
||
-1, # -1 表示基座
|
||
mass=0,
|
||
linearDamping=0,
|
||
angularDamping=0,
|
||
lateralFriction=0
|
||
)
|
||
|
||
# 创建固定约束，将目标物体固定在转盘上
|
||
cid = p.createConstraint(
|
||
parentBodyUniqueId=self.turntable_id,
|
||
parentLinkIndex=-1, # -1 表示基座
|
||
childBodyUniqueId=self.target_id,
|
||
childLinkIndex=-1, # -1 表示基座
|
||
jointType=p.JOINT_FIXED,
|
||
jointAxis=[0, 0, 0],
|
||
parentFramePosition=[0, 0, 0], # 相对于转盘中心的偏移
|
||
childFramePosition=[0, 0, 0] # 相对于物体中心的偏移
|
||
)
|
||
|
||
# 设置约束参数
|
||
p.changeConstraint(cid, maxForce=100) # 设置最大力，确保约束稳定
|
||
|
||
def move_robot_to_pose(self, target_matrix):
|
||
# 从4x4齐次矩阵中提取位置（前3个元素）
|
||
position = target_matrix[:3, 3]
|
||
|
||
# 从3x3旋转矩阵中提取方向四元数
|
||
R = target_matrix[:3, :3]
|
||
|
||
# 计算四元数的w分量
|
||
w = np.sqrt(max(0, 1 + R[0,0] + R[1,1] + R[2,2])) / 2
|
||
|
||
# 避免除零错误，同时处理不同情况
|
||
if abs(w) < 1e-8:
|
||
# 当w接近0时的特殊情况
|
||
x = np.sqrt(max(0, 1 + R[0,0] - R[1,1] - R[2,2])) / 2
|
||
y = np.sqrt(max(0, 1 - R[0,0] + R[1,1] - R[2,2])) / 2
|
||
z = np.sqrt(max(0, 1 - R[0,0] - R[1,1] + R[2,2])) / 2
|
||
|
||
# 确定符号
|
||
if R[2,1] - R[1,2] < 0: x = -x
|
||
if R[0,2] - R[2,0] < 0: y = -y
|
||
if R[1,0] - R[0,1] < 0: z = -z
|
||
else:
|
||
# 正常情况
|
||
x = (R[2,1] - R[1,2]) / (4 * w)
|
||
y = (R[0,2] - R[2,0]) / (4 * w)
|
||
z = (R[1,0] - R[0,1]) / (4 * w)
|
||
|
||
orientation = (x, y, z, w)
|
||
|
||
# 设置IK求解参数
|
||
num_joints = p.getNumJoints(self.robot_id)
|
||
lower_limits = []
|
||
upper_limits = []
|
||
joint_ranges = []
|
||
rest_poses = []
|
||
|
||
# 获取关节限制和默认姿态
|
||
for i in range(num_joints):
|
||
joint_info = p.getJointInfo(self.robot_id, i)
|
||
lower_limits.append(joint_info[8])
|
||
upper_limits.append(joint_info[9])
|
||
joint_ranges.append(joint_info[9] - joint_info[8])
|
||
rest_poses.append(0) # 可以设置一个较好的默认姿态
|
||
|
||
# 使用增强版IK求解器，考虑碰撞避障
|
||
joint_poses = p.calculateInverseKinematics(
|
||
self.robot_id,
|
||
7, # end effector link index
|
||
position,
|
||
orientation,
|
||
lowerLimits=lower_limits,
|
||
upperLimits=upper_limits,
|
||
jointRanges=joint_ranges,
|
||
restPoses=rest_poses,
|
||
maxNumIterations=100,
|
||
residualThreshold=1e-4
|
||
)
|
||
|
||
# 分步移动到目标位置，同时检查碰撞
|
||
current_poses = [p.getJointState(self.robot_id, i)[0] for i in range(7)]
|
||
steps = 50 # 分50步移动
|
||
|
||
for step in range(steps):
|
||
# 线性插值计算中间位置
|
||
intermediate_poses = []
|
||
for current, target in zip(current_poses, joint_poses):
|
||
t = (step + 1) / steps
|
||
intermediate = current + (target - current) * t
|
||
intermediate_poses.append(intermediate)
|
||
|
||
# 设置关节位置
|
||
for i in range(7):
|
||
p.setJointMotorControl2(
|
||
self.robot_id,
|
||
i,
|
||
p.POSITION_CONTROL,
|
||
intermediate_poses[i]
|
||
)
|
||
|
||
# 执行一步模拟
|
||
p.stepSimulation()
|
||
|
||
# 检查碰撞
|
||
if p.getContactPoints(self.robot_id, self.turntable_id):
|
||
print("检测到潜在碰撞，停止移动")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
def rotate_turntable(self, angle_degrees):
|
||
# 旋转转盘
|
||
current_pos, current_orn = p.getBasePositionAndOrientation(self.turntable_id)
|
||
current_orn = p.getEulerFromQuaternion(current_orn)
|
||
|
||
new_orn = list(current_orn)
|
||
new_orn[2] += np.radians(angle_degrees)
|
||
new_orn_quat = p.getQuaternionFromEuler(new_orn)
|
||
|
||
p.resetBasePositionAndOrientation(
|
||
self.turntable_id,
|
||
current_pos,
|
||
new_orn_quat
|
||
)
|
||
|
||
# 同时旋转目标物体
|
||
target_pos, target_orn = p.getBasePositionAndOrientation(self.target_id)
|
||
target_orn = p.getEulerFromQuaternion(target_orn)
|
||
|
||
# 更新目标物体的方向
|
||
target_orn = list(target_orn)
|
||
target_orn[2] += np.radians(angle_degrees)
|
||
target_orn_quat = p.getQuaternionFromEuler(target_orn)
|
||
|
||
# 计算物体新的位置（绕转盘中心旋转）
|
||
turntable_center = current_pos
|
||
relative_pos = np.array(target_pos) - np.array(turntable_center)
|
||
|
||
# 创建旋转矩阵
|
||
theta = np.radians(angle_degrees)
|
||
rotation_matrix = np.array([
|
||
[np.cos(theta), -np.sin(theta), 0],
|
||
[np.sin(theta), np.cos(theta), 0],
|
||
[0, 0, 1]
|
||
])
|
||
|
||
# 计算新的相对位置
|
||
new_relative_pos = rotation_matrix.dot(relative_pos)
|
||
new_pos = np.array(turntable_center) + new_relative_pos
|
||
|
||
# 更新目标物体的位置和方向
|
||
p.resetBasePositionAndOrientation(
|
||
self.target_id,
|
||
new_pos,
|
||
target_orn_quat
|
||
)
|
||
|
||
def get_camera_pose(self):
|
||
end_effector_link = 7 # Franka末端执行器的链接索引
|
||
state = p.getLinkState(self.robot_id, end_effector_link)
|
||
ee_pos = state[0] # 世界坐标系中的位置
|
||
camera_orn = state[1] # 世界坐标系中的朝向（四元数）
|
||
|
||
# 计算相机的视角矩阵
|
||
rot_matrix = p.getMatrixFromQuaternion(camera_orn)
|
||
rot_matrix = np.array(rot_matrix).reshape(3, 3)
|
||
|
||
# 相机的前向向量（与末端执行器的x轴对齐）
|
||
camera_forward = rot_matrix.dot(np.array([0, 0, 1])) # x轴方向
|
||
|
||
# 将相机位置向前偏移0.1米
|
||
offset = 0.12
|
||
camera_pos = np.array(ee_pos) + camera_forward * offset
|
||
camera_target = camera_pos + camera_forward
|
||
|
||
# 相机的上向量（与末端执行器的z轴对齐）
|
||
camera_up = rot_matrix.dot(np.array([1, 0, 0])) # z轴方向
|
||
|
||
return camera_pos, camera_target, camera_up
|
||
|
||
def take_picture(self):
|
||
camera_pos, camera_target, camera_up = self.get_camera_pose()
|
||
|
||
view_matrix = p.computeViewMatrix(
|
||
cameraEyePosition=camera_pos,
|
||
cameraTargetPosition=camera_target,
|
||
cameraUpVector=camera_up
|
||
)
|
||
|
||
projection_matrix = p.computeProjectionMatrixFOV(
|
||
fov=self.camera_params['fov'],
|
||
aspect=self.camera_params['width'] / self.camera_params['height'],
|
||
nearVal=self.camera_params['near'],
|
||
farVal=self.camera_params['far']
|
||
)
|
||
|
||
_,_,rgb_img,depth_img,segm_img = p.getCameraImage(
|
||
width=self.camera_params['width'],
|
||
height=self.camera_params['height'],
|
||
viewMatrix=view_matrix,
|
||
projectionMatrix=projection_matrix,
|
||
renderer=p.ER_BULLET_HARDWARE_OPENGL
|
||
)
|
||
|
||
depth_img = self.camera_params['far'] * self.camera_params['near'] / (
|
||
self.camera_params['far'] - (self.camera_params['far'] - self.camera_params['near']) * depth_img)
|
||
|
||
depth_img = np.array(depth_img)
|
||
segm_img = np.array(segm_img)
|
||
|
||
return depth_img, segm_img
|
||
|
||
def reset(self):
|
||
target_pos = [0.5, 0, 1]
|
||
target_orn = p.getQuaternionFromEuler([np.pi, 0, 0])
|
||
target_matrix = np.eye(4)
|
||
target_matrix[:3, 3] = target_pos
|
||
target_matrix[:3, :3] = np.asarray(p.getMatrixFromQuaternion(target_orn)).reshape(3,3)
|
||
self.move_robot_to_pose(target_matrix)
|
||
|
||
def init(self):
|
||
self.move_to(Simulator.INIT_GRIPPER_POSE)
|
||
|
||
def move_to(self, pose: np.ndarray):
|
||
#delta_degree, min_new_cam_to_world = ControlUtil.solve_display_table_rot_and_cam_to_world(pose)
|
||
#print(delta_degree)
|
||
min_new_cam_to_pybullet_world = Simulator.TURNTABLE_WORLD_TO_PYBULLET_WORLD@pose
|
||
self.move_to_cam_pose(min_new_cam_to_pybullet_world)
|
||
#self.rotate_turntable(delta_degree)
|
||
|
||
|
||
|
||
def __del__(self):
|
||
p.disconnect()
|
||
|
||
def create_experiment(self, backup_name=None):
|
||
return super().create_experiment(backup_name)
|
||
|
||
def load_experiment(self, backup_name=None):
|
||
super().load_experiment(backup_name)
|
||
|
||
def move_to_cam_pose(self, camera_pose: np.ndarray):
|
||
# 从相机位姿矩阵中提取位置和旋转矩阵
|
||
camera_pos = camera_pose[:3, 3]
|
||
R_camera = camera_pose[:3, :3]
|
||
|
||
# 相机的朝向向量（z轴）
|
||
forward = R_camera[:, 2]
|
||
|
||
# 由于相机与末端执行器之间有固定偏移，需要计算末端执行器位置
|
||
# 相机在末端执行器前方0.12米
|
||
gripper_pos = camera_pos - forward * 0.12
|
||
|
||
# 末端执行器的旋转矩阵需要考虑与相机坐标系的固定变换
|
||
# 假设相机的forward对应gripper的z轴，相机的x轴对应gripper的x轴
|
||
R_gripper = R_camera
|
||
|
||
# 构建4x4齐次变换矩阵
|
||
gripper_pose = np.eye(4)
|
||
gripper_pose[:3, :3] = R_gripper
|
||
gripper_pose[:3, 3] = gripper_pos
|
||
print(gripper_pose)
|
||
# 移动机器人到计算出的位姿
|
||
return self.move_robot_to_pose(gripper_pose) |