|
| 1 | +import logging |
| 2 | +import pathlib |
| 3 | +from typing import Any, cast |
| 4 | + |
| 5 | +import gymnasium as gym |
| 6 | +import mujoco |
| 7 | +import numpy as np |
| 8 | +from rcs._core.common import Pose |
| 9 | +from rcs.envs.base import ControlMode, GripperWrapper, RobotEnv |
| 10 | +from rcs.envs.creators import FR3SimplePickUpSimEnvCreator |
| 11 | +from rcs.ompl.mj_ompl import MjOMPL |
| 12 | + |
| 13 | +import rcs |
| 14 | + |
| 15 | +logger = logging.getLogger(__name__) |
| 16 | +logger.setLevel(logging.INFO) |
| 17 | + |
| 18 | +RECORD_PATH = "/home/sbien/Documents/Development/RCS/internal_rcs/test_data" |
| 19 | +SCENE_ROOT = pathlib.Path("~", "Documents", "Development", "RCS", "models", "scenes", "viral").expanduser() |
| 20 | +SCENE_FILE = pathlib.Path(SCENE_ROOT, "scene.xml") |
| 21 | +CAMERAS = [ |
| 22 | + "wrist", |
| 23 | + "arro", |
| 24 | + "bird_eye", |
| 25 | + "full_side", |
| 26 | + "front_up", |
| 27 | + "front_down", |
| 28 | + "side_right", |
| 29 | + "side_left_up", |
| 30 | + "back_right_up", |
| 31 | + "back_center", |
| 32 | + "back_left", |
| 33 | + "side_opposite", |
| 34 | + "side_center", |
| 35 | +] |
| 36 | + |
| 37 | + |
| 38 | +class OmplTrajectoryDemo: |
| 39 | + def __init__(self, env: gym.Env, planner: MjOMPL): |
| 40 | + self.env = env |
| 41 | + self.unwrapped: RobotEnv = cast(RobotEnv, self.env.unwrapped) |
| 42 | + self.home_pose: Pose = self.unwrapped.robot.get_cartesian_position() |
| 43 | + self.home_qpos: np.ndarray = self.unwrapped.robot.get_joint_position() |
| 44 | + self.sol_path = None |
| 45 | + self.planner = planner |
| 46 | + |
| 47 | + def _action(self, pose: Pose, gripper: float) -> dict[str, Any]: |
| 48 | + return {"xyzrpy": pose.xyzrpy(), "gripper": gripper} |
| 49 | + |
| 50 | + def _jaction(self, joints: np.ndarray, gripper: float) -> dict[str, Any]: |
| 51 | + return {"joints": joints, "gripper": gripper} |
| 52 | + |
| 53 | + def get_object_pose(self, geom_name) -> Pose: |
| 54 | + model = self.env.get_wrapper_attr("sim").model |
| 55 | + data = self.env.get_wrapper_attr("sim").data |
| 56 | + |
| 57 | + geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, geom_name) |
| 58 | + obj_pose_world_coordinates = Pose( |
| 59 | + translation=data.geom_xpos[geom_id], rotation=data.geom_xmat[geom_id].reshape(3, 3) |
| 60 | + ) * Pose( |
| 61 | + rpy_vector=np.array([0, 0, np.pi]), translation=[0, 0, 0] # type: ignore |
| 62 | + ) |
| 63 | + return self.unwrapped.robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates) |
| 64 | + |
| 65 | + def plan_path_to_object(self, obj_name: str, delta_up): |
| 66 | + self.move_home() |
| 67 | + obj_pose_og = self.get_object_pose(geom_name=obj_name) |
| 68 | + obj_pose = obj_pose_og * Pose( |
| 69 | + translation=np.array([0, 0, delta_up]), quaternion=np.array([1, 0, 0, 0]) # type: ignore |
| 70 | + ) # type: ignore |
| 71 | + goal_q = self.planner.ik(obj_pose, None) |
| 72 | + if goal_q is None: |
| 73 | + return False |
| 74 | + sol, path = self.planner.plan(goal=goal_q, start=self.home_qpos) |
| 75 | + if not sol: |
| 76 | + return False |
| 77 | + for waypoint in path: |
| 78 | + self.step(self._jaction(waypoint, GripperWrapper.BINARY_GRIPPER_OPEN)) |
| 79 | + return True |
| 80 | + |
| 81 | + def approach_and_grasp(self, obj_name: str, delta_up: float = 0.2): |
| 82 | + obj_pose_og = self.get_object_pose(geom_name=obj_name) |
| 83 | + |
| 84 | + obj_pose_grasp = obj_pose_og * Pose(translation=np.array([0, 0, delta_up]), quaternion=np.array([1, 0, 0, 0])) # type: ignore |
| 85 | + waypoints = self.generate_waypoints( |
| 86 | + start_pose=self.unwrapped.robot.get_cartesian_position(), end_pose=obj_pose_grasp, num_waypoints=5 |
| 87 | + ) |
| 88 | + for waypoint in waypoints: |
| 89 | + self.step(self._jaction(waypoint, GripperWrapper.BINARY_GRIPPER_OPEN)) # type: ignore |
| 90 | + for _ in range(3): # Required for simulation stability |
| 91 | + self.step(self._jaction(waypoints[-1], GripperWrapper.BINARY_GRIPPER_CLOSED)) # type: ignore |
| 92 | + |
| 93 | + def generate_waypoints(self, start_pose: Pose, end_pose: Pose, num_waypoints: int) -> list[Pose]: |
| 94 | + start_q = self.planner.ik(start_pose, None) |
| 95 | + start_q = start_q[:7] # ignore finger joints |
| 96 | + waypoints = [start_q] |
| 97 | + for i in range(num_waypoints + 1): |
| 98 | + t = i / (num_waypoints) |
| 99 | + waypoints.append(self.planner.ik(start_pose.interpolate(end_pose, t), None)[:7]) # ignore finger joints |
| 100 | + return waypoints |
| 101 | + |
| 102 | + def step(self, action: dict) -> dict: |
| 103 | + return self.env.step(action)[0] |
| 104 | + |
| 105 | + def execute_motion(self, waypoints: list[Pose], gripper: float = GripperWrapper.BINARY_GRIPPER_OPEN) -> dict: |
| 106 | + for i in range(len(waypoints)): |
| 107 | + obs = self.step(self._jaction(waypoints[i], gripper)) # type: ignore |
| 108 | + return obs |
| 109 | + |
| 110 | + def move_home(self): |
| 111 | + end_eff_pose = self.unwrapped.robot.get_cartesian_position() |
| 112 | + waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=15) |
| 113 | + self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED) |
| 114 | + |
| 115 | + |
| 116 | +def main(): |
| 117 | + # Make an environment as usual |
| 118 | + env = FR3SimplePickUpSimEnvCreator()(render_mode="human", delta_actions=False, control_mode=ControlMode.JOINTS) |
| 119 | + robot_tcp = rcs.common.Pose( |
| 120 | + translation=np.array([0.0, 0.0, 0.1034]), # type: ignore |
| 121 | + rotation=np.array([[0.707, 0.707, 0], [-0.707, 0.707, 0], [0, 0, 1]]), # type: ignore |
| 122 | + ) |
| 123 | + planner = "RRTConnect" # one of ["PRM","RRT","RRTConnect","RRTstar","EST","FMT","BITstar"] |
| 124 | + |
| 125 | + for _ in range(50): |
| 126 | + env.reset() |
| 127 | + |
| 128 | + # For each episode, we create a new planner instance to reset its space information. |
| 129 | + ompl_planner = MjOMPL( |
| 130 | + robot_name="fr3", # Name is only for logging purposes. |
| 131 | + robot_env=env, # Pass the environment created above to the planner. |
| 132 | + njoints=7, # Specify the number of joints |
| 133 | + robot_tcp=robot_tcp, # Specify the TCP of the robot to be planned around. |
| 134 | + robot_xml_name=rcs.scenes[ |
| 135 | + "fr3_simple_pick_up" |
| 136 | + ].mjcf_robot, # Path to the robot xml file (NOT the scene.xml) |
| 137 | + robot_root_name="base_0", # Name of the robot root body in the xml |
| 138 | + robot_joint_name="fr3_joint#_0", # Joint name pattern of the robot in the xml, where # is [1~njoints] |
| 139 | + robot_actuator_name="fr3_joint#_0", # Actuator name pattern, used to ensure that we only plan for actuated joints. |
| 140 | + obstacle_body_names=[], # What MuJoCo bodies to consider as obstacles |
| 141 | + obstacle_geom_names=["floor", "box_geom"], # What MuJoCo geoms to consider as obstacles |
| 142 | + interpolate_num=50, # Number of interpolation points between start and goal configuration |
| 143 | + ) |
| 144 | + ompl_planner.set_planner(planner) |
| 145 | + controller = OmplTrajectoryDemo(env, ompl_planner) |
| 146 | + |
| 147 | + # Since the target object is part of the collision to be checked, we need sufficient clearance (delta_up=0.1) |
| 148 | + # for the planner to find a solution. |
| 149 | + success = controller.plan_path_to_object("box_geom", delta_up=0.05) |
| 150 | + if success: |
| 151 | + logger.info(f"Successfully planned with {planner}.") |
| 152 | + # Once successful, we approach the object and grasp it. |
| 153 | + controller.approach_and_grasp("box_geom", delta_up=0.01) |
| 154 | + # And then return to home position. |
| 155 | + controller.move_home() |
| 156 | + else: |
| 157 | + logger.warning(f"No solution found. Failed planning with {planner}.") |
| 158 | + continue |
| 159 | + # Since we are creating the planner object at each loop to reset its space information, we delete it here. |
| 160 | + # This also frees the memory and prevents a potential crash. |
| 161 | + ompl_planner = None |
| 162 | + logger.info("Finished all episodes.") |
| 163 | + env.close() |
| 164 | + |
| 165 | + |
| 166 | +if __name__ == "__main__": |
| 167 | + main() |
0 commit comments