├── .gitignore ├── README.md ├── homestri_ur5e_rl ├── __init__.py ├── controllers │ ├── __init__.py │ ├── cartesian_motion_controller.py │ ├── force_torque_sensor_controller.py │ ├── joint_controller.py │ ├── joint_effort_controller.py │ ├── joint_position_controller.py │ ├── joint_velocity_controller.py │ └── operational_space_controller.py ├── envs │ ├── __init__.py │ ├── assets │ │ └── base_robot │ │ │ ├── assets │ │ │ ├── rethink_pedestal │ │ │ │ └── pedestal.stl │ │ │ ├── robotiq_2f85 │ │ │ │ ├── base.stl │ │ │ │ ├── base_mount.stl │ │ │ │ ├── coupler.stl │ │ │ │ ├── driver.stl │ │ │ │ ├── follower.stl │ │ │ │ ├── pad.stl │ │ │ │ ├── silicone_pad.stl │ │ │ │ └── spring_link.stl │ │ │ └── universal_robots_ur5e │ │ │ │ ├── base_0.obj │ │ │ │ ├── base_1.obj │ │ │ │ ├── forearm_0.obj │ │ │ │ ├── forearm_1.obj │ │ │ │ ├── forearm_2.obj │ │ │ │ ├── forearm_3.obj │ │ │ │ ├── shoulder_0.obj │ │ │ │ ├── shoulder_1.obj │ │ │ │ ├── shoulder_2.obj │ │ │ │ ├── upperarm_0.obj │ │ │ │ ├── upperarm_1.obj │ │ │ │ ├── upperarm_2.obj │ │ │ │ ├── upperarm_3.obj │ │ │ │ ├── wrist1_0.obj │ │ │ │ ├── wrist1_1.obj │ │ │ │ ├── wrist1_2.obj │ │ │ │ ├── wrist2_0.obj │ │ │ │ ├── wrist2_1.obj │ │ │ │ ├── wrist2_2.obj │ │ │ │ └── wrist3.obj │ │ │ ├── scene.xml │ │ │ ├── table.xml │ │ │ └── wire.xml │ ├── base_robot │ │ ├── __init__.py │ │ └── base_robot_env.py │ └── mujoco │ │ ├── mujoco_env.py │ │ └── mujoco_rendering.py ├── input_devices │ └── keyboard_input.py └── utils │ ├── controller_utils.py │ ├── mujoco_utils.py │ ├── pid_controller_utils.py │ ├── solver_utils.py │ └── transform_utils.py ├── media └── wire_manipulation.gif ├── scripts ├── circle_points.py └── test_base_robot_env.py └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # homestri-ur5e-rl 2 | 3 | Welcome to the **homestri-ur5e-rl** repository! This repository provides a Mujoco Gymnasium Environment designed for manipulating flexible objects using a UR5e robot arm and a Robotiq 2F-85 gripper. 4 | 5 | ![Robot Manipulation](media/wire_manipulation.gif) 6 | 7 | ## Overview 8 | 9 | The key features of this repository are: 10 | 11 | - **Mujoco Simulation**: The repository contains an example Mujoco XML configuration featuring a UR5e robot arm with an attached Robotiq 2F-85 gripper. This setup allows the robot to interact with a flexible cable in a simulated environment. 12 | 13 | - **Gymnasium Integration**: The environment is seamlessly integrated into the Gymnasium framework and comes with a graphical user interface (GUI) that visualizes the Mujoco simulation. This GUI provides an interactive way to observe the robot's actions and interactions. 14 | 15 | - **Controller Variety**: The UR5e robot arm offers multiple controller options, enabling various modes of control, including: 16 | - Joint effort controller 17 | - Joint velocity controller 18 | - Joint position controller 19 | - Operational space controller 20 | 21 | - **Teleoperation**: For manual control, keyboard input is supported, allowing teleoperation of the robot in task space. This can be particularly useful for specific tasks that require direct human intervention. 22 | 23 | ## Repository Information 24 | 25 | Please note the following details about this repository: 26 | 27 | - **Maintenance and Status**: This repository is not actively maintained and should be considered experimental. It was created during the initial stages of learning Mujoco. 28 | 29 | - **References**: The development of this repository drew inspiration from the following repositories: 30 | - [ARISE-Initiative/robosuite.git](https://github.com/ARISE-Initiative/robosuite.git) 31 | - [ir-lab/irl_control.git](https://github.com/ir-lab/irl_control.git) 32 | - [abr/abr_control.git](https://github.com/abr/abr_control.git) 33 | 34 | ## Installation 35 | 36 | Follow these steps to set up the environment: 37 | 38 | 1. Clone the repository: 39 | ``` 40 | git clone https://github.com/ian-chuang/homestri-ur5e-rl.git 41 | ``` 42 | 43 | 2. Navigate to the repository directory: 44 | ``` 45 | cd homestri-ur5e-rl 46 | ``` 47 | 48 | 3. Install the required packages using pip: 49 | ``` 50 | pip install -e . 51 | ``` 52 | 53 | ## Testing the Environment 54 | 55 | To test the environment, you can execute the following command: 56 | ``` 57 | python scripts/test_base_robot_env.py 58 | ``` 59 | 60 | Feel free to explore and experiment with the capabilities of the **homestri-ur5e-rl** environment. Keep in mind that this repository is an experimental project and might not receive regular updates. If you encounter issues or have questions, refer to the repository or seek support from the developer community. -------------------------------------------------------------------------------- /homestri_ur5e_rl/__init__.py: -------------------------------------------------------------------------------- 1 | from gymnasium.envs.registration import register 2 | 3 | register( 4 | id="BaseRobot-v0", 5 | entry_point="homestri_ur5e_rl.envs.base_robot:BaseRobot", 6 | ) -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/controllers/__init__.py -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/cartesian_motion_controller.py: -------------------------------------------------------------------------------- 1 | from homestri_ur5e_rl.utils.mujoco_utils import ( 2 | get_site_xpos_by_id, 3 | get_site_xmat_by_id, 4 | get_site_xvelp_by_id, 5 | get_site_xvelr_by_id, 6 | get_site_jacp_by_id, 7 | get_site_jacr_by_id, 8 | get_obs_by_id, 9 | get_obs_by_name 10 | ) 11 | from homestri_ur5e_rl.utils.transform_utils import ( 12 | compute_position_error, 13 | set_goal_position, 14 | set_goal_orientation 15 | ) 16 | from homestri_ur5e_rl.utils.pd_controller_utils import ( 17 | SpatialPDController 18 | ) 19 | from homestri_ur5e_rl.utils.solver_utils import ( 20 | JacobianTransposeSolver 21 | ) 22 | import numpy as np 23 | 24 | 25 | class OperationalSpaceController(): 26 | def __init__( 27 | self, 28 | model, 29 | data, 30 | model_names, 31 | eef_name, 32 | joint_names, 33 | actuator_names, 34 | p_values=None, 35 | d_values=None 36 | ): 37 | self.model = model 38 | self.data = data 39 | self.model_names = model_names 40 | self.eef_name = eef_name 41 | self.joint_names = joint_names 42 | self.actuator_names = actuator_names 43 | 44 | self.eef_id = self.model_names.site_name2id[self.eef_name] 45 | self.joint_ids = [self.model_names.joint_name2id[name] for name in self.joint_names] 46 | self.jnt_qpos_ids = [self.model.jnt_qposadr[id] for id in self.joint_ids] 47 | self.jnt_dof_ids = [self.model.jnt_dofadr[id] for id in self.joint_ids] 48 | self.actuator_ids = [self.model_names.actuator_name2id[name] for name in actuator_names] 49 | 50 | self.goal_pos = None 51 | self.goal_ori_mat = None 52 | self.control_period = self.model.opt.timestep 53 | self.internal_period = 0.02 54 | 55 | self.spatial_controller = SpatialPDController(p_values, d_values) 56 | self.solver = JacobianTransposeSolver(len(self.joint_ids)) 57 | 58 | # set goal to current end effector position and orientation 59 | self.reset_controller() 60 | 61 | def set_goal(self, goal_pos, goal_ori_mat): 62 | self.goal_pos = goal_pos 63 | self.goal_ori_mat = goal_ori_mat 64 | 65 | def reset_controller(self): 66 | self.goal_pos = get_site_xpos_by_id(self.model, self.data, self.eef_id) 67 | self.goal_ori_mat = get_site_xmat_by_id(self.model, self.data, self.eef_id) 68 | 69 | self.goal_twist = np.zeros(6) 70 | 71 | self.solver.sync_jnt_pos(self.get_joint_pos()) 72 | 73 | self._compute_joint_control_cmds(np.zeros(6), self.internal_period) 74 | 75 | def run_controller(self): 76 | 77 | 78 | error = self._compute_error() 79 | 80 | joint_pos_cmds, joint_vel_cmds = self._compute_joint_control_cmds(error, self.control_period) 81 | 82 | return joint_pos_cmds, joint_vel_cmds 83 | 84 | def get_eef_pos(self): 85 | ee_pos = get_site_xpos_by_id(self.model, self.data, self.eef_id) 86 | ee_ori_mat = get_site_xmat_by_id(self.model, self.data, self.eef_id) 87 | return ee_pos, ee_ori_mat 88 | 89 | def get_joint_pos(self): 90 | return self.data.qpos[self.jnt_qpos_ids] 91 | 92 | def get_actuator_ids(self): 93 | return self.actuator_ids 94 | 95 | def _compute_error(self): 96 | # store end effector position and orientation using mujoco_utils 97 | ee_pos = get_site_xpos_by_id(self.model, self.data, self.eef_id) 98 | ee_ori_mat = get_site_xmat_by_id(self.model, self.data, self.eef_id) 99 | ee_pos_vel = get_site_xvelp_by_id(self.model, self.data, self.eef_id) 100 | ee_ori_vel = get_site_xvelr_by_id(self.model, self.data, self.eef_id) 101 | 102 | # new_ee_pos = set_goal_position(self.goal_twist[:3], ee_pos) 103 | # new_ee_ori_mat = set_goal_orientation(self.goal_twist[3:], ee_ori_mat) 104 | 105 | # compute error 106 | pos_error = compute_position_error(self.goal_pos, self.goal_ori_mat, ee_pos, ee_ori_mat) 107 | 108 | 109 | 110 | return pos_error 111 | 112 | def _compute_joint_control_cmds(self, error, period): 113 | 114 | # might want to do iterations thus have to update kinematic model 115 | # use this as resource https://github.com/deepmind/mujoco/issues/411 116 | 117 | # store jacobian matrices using mujoco_utils 118 | J_pos = get_site_jacp_by_id(self.model, self.data, self.eef_id) 119 | J_ori = get_site_jacr_by_id(self.model, self.data, self.eef_id) 120 | J_full = np.vstack([J_pos, J_ori]) 121 | J = J_full[:, self.jnt_dof_ids] 122 | 123 | # compute control output 124 | cartesian_input = self.spatial_controller(error, self.internal_period) 125 | 126 | joint_pos_cmds, joint_vel_cmds = self.solver.compute_jnt_ctrl(J, self.internal_period, cartesian_input) 127 | 128 | return joint_pos_cmds, joint_vel_cmds -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/force_torque_sensor_controller.py: -------------------------------------------------------------------------------- 1 | # Import necessary modules and classes 2 | from mujoco._structs import MjModel 3 | from mujoco._structs import MjData 4 | import mujoco 5 | from homestri_ur5e_rl.utils.mujoco_utils import MujocoModelNames 6 | import numpy as np 7 | 8 | class ForceTorqueSensorController: 9 | def __init__( 10 | self, 11 | model: MjModel, 12 | data: MjData, 13 | model_names: MujocoModelNames, 14 | ft_sensor_site: str, 15 | force_sensor_name: str, 16 | torque_sensor_name: str, 17 | subtree_body_name: str, 18 | smoothing_factor: float = 0.1, 19 | ) -> None: 20 | """ 21 | Base controller class to define a common interface for robot controllers. 22 | 23 | Parameters: 24 | model (mujoco._structs.MjModel): Mujoco model representing the robot. 25 | data (mujoco._structs.MjData): Mujoco data associated with the model. 26 | model_names (homestri_ur5e_rl.utils.mujoco_utils.MujocoModelNames): Names of different components in the Mujoco model. 27 | ft_sensor_site (str): Name of the force-torque sensor site in the Mujoco model. 28 | force_sensor_name (str): Name of the force sensor in the Mujoco model. 29 | torque_sensor_name (str): Name of the torque sensor in the Mujoco model. 30 | subtree_body_name (str): Name of the body on which the force-torque sensor is mounted. 31 | """ 32 | self.model = model 33 | self.data = data 34 | self.model_names = model_names 35 | self.ft_sensor_site = ft_sensor_site 36 | self.force_sensor_name = force_sensor_name 37 | self.torque_sensor_name = torque_sensor_name 38 | self.subtree_body_name = subtree_body_name 39 | self.smoothing_factor = smoothing_factor 40 | 41 | # Get IDs for the site, force sensor, torque sensor, and subtree body 42 | self.ft_site_id = self.model_names.site_name2id[self.ft_sensor_site] 43 | self.force_id = self.model_names.sensor_name2id[self.force_sensor_name] 44 | self.torque_id = self.model_names.sensor_name2id[self.torque_sensor_name] 45 | self.subtree_body_id = self.model_names.body_name2id[self.subtree_body_name] 46 | 47 | # Make sure the specified sensors are of the correct type 48 | assert model.sensor_type[self.force_id] == mujoco.mjtSensor.mjSENS_FORCE 49 | assert model.sensor_type[self.torque_id] == mujoco.mjtSensor.mjSENS_TORQUE 50 | 51 | # Get addresses for force and torque sensor data 52 | self.force_adr = model.sensor_adr[self.force_id] 53 | self.torque_adr = model.sensor_adr[self.torque_id] 54 | 55 | # Number of dimensions for force and torque vectors (typically 3 for each) 56 | self.force_ndim = 3 57 | self.torque_ndim = 3 58 | 59 | # Initialize zero signal offset as a 6-dimensional zero array 60 | self.zero_signal_offset = np.zeros(6) 61 | self.filtered_wrench = None 62 | 63 | def _gravity_compensation(self) -> np.ndarray: 64 | """ 65 | Calculate the gravity compensation for the subtree. 66 | 67 | Returns: 68 | np.ndarray: Wrench containing the gravity compensation forces and torques in world frame (6-dimensional array). 69 | """ 70 | 71 | # Get the center of mass of the specified body (CoM in world coordinates) 72 | subtree_com = self.data.subtree_com[self.subtree_body_id] 73 | 74 | # Get the mass of the subtree of the specified body 75 | subtree_mass = self.model.body_subtreemass[self.subtree_body_id] 76 | 77 | # Get the gravity vector from the physics model 78 | gravity = self.model.opt.gravity 79 | 80 | # Get the position of the FT sensor site in world coordinates 81 | ft_pos = self.data.site_xpos[self.ft_site_id] 82 | 83 | # Calculate the relative position vector from the FT sensor to the subtree's CoM 84 | relative_pos = subtree_com - ft_pos 85 | 86 | # The relative gravity acting on the subtree's CoM (typically just the gravity vector) 87 | relative_gravity = gravity 88 | 89 | # Calculate the gravitational force 90 | F_gravity = subtree_mass * (relative_gravity) 91 | 92 | # Calculate the gravitational torque 93 | T_gravity = np.cross(relative_pos, F_gravity) 94 | 95 | return np.concatenate([F_gravity, T_gravity]) 96 | 97 | def zero_signal(self) -> None: 98 | """ 99 | Capture the current wrench as the zero signal offset. 100 | """ 101 | self.zero_signal_offset = self.get_wrench() 102 | self.filtered_wrench = None 103 | 104 | def get_wrench(self) -> np.ndarray: 105 | """ 106 | Get the current wrench (force and torque) applied at the force-torque sensor. 107 | 108 | Returns: 109 | np.ndarray: Wrench containing the forces and torques applied at the force-torque sensor in world frame (6-dimensional array). 110 | """ 111 | 112 | # Get the orientation matrix of the force-torque (FT) sensor 113 | ft_ori_mat = self.data.site_xmat[self.ft_site_id].reshape(3, 3) 114 | 115 | # Get the raw force and torque measurements from the sensor 116 | force = self.data.sensordata[self.force_adr : self.force_adr + self.force_ndim] 117 | torque = self.data.sensordata[self.torque_adr : self.torque_adr + self.torque_ndim] 118 | 119 | # Transform the force and torque from the sensor frame to the world frame 120 | force = ft_ori_mat @ force 121 | torque = ft_ori_mat @ torque 122 | 123 | # Concatenate the force and torque vectors to obtain the resultant wrench 124 | wrench = np.concatenate([force, torque]) 125 | 126 | # Add the gravity compensation term to the wrench 127 | # wrench += self._gravity_compensation() 128 | 129 | # Subtract the zero signal offset from the wrench 130 | # wrench -= self.zero_signal_offset 131 | 132 | # Apply the smoothing filter 133 | # if self.filtered_wrench is None: 134 | # self.filtered_wrench = wrench 135 | # else: 136 | # self.filtered_wrench = (1 - self.smoothing_factor) * self.filtered_wrench + self.smoothing_factor * wrench 137 | 138 | # return self.filtered_wrench 139 | 140 | return wrench 141 | 142 | 143 | 144 | 145 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/joint_controller.py: -------------------------------------------------------------------------------- 1 | # Import necessary modules and classes 2 | from mujoco._structs import MjModel 3 | from mujoco._structs import MjData 4 | from homestri_ur5e_rl.utils.mujoco_utils import MujocoModelNames 5 | from typing import List 6 | import numpy as np 7 | 8 | class JointController: 9 | def __init__( 10 | self, 11 | model: MjModel, 12 | data: MjData, 13 | model_names: MujocoModelNames, 14 | eef_name: str, 15 | joint_names: List[str], 16 | actuator_names: List[str] 17 | ) -> None: 18 | """ 19 | Base controller class to define a common interface for robot controllers. 20 | 21 | Parameters: 22 | model (mujoco._structs.MjModel): Mujoco model representing the robot. 23 | data (mujoco._structs.MjData): Mujoco data associated with the model. 24 | model_names (homestri_ur5e_rl.utils.mujoco_utils.MujocoModelNames): Names of different components in the Mujoco model. 25 | eef_name (str): Name of the end-effector in the Mujoco model. 26 | joint_names (List[str]): List of joint names for the robot. 27 | actuator_names (List[str]): List of actuator names for the robot. 28 | """ 29 | self.model = model 30 | self.data = data 31 | self.model_names = model_names 32 | self.eef_name = eef_name 33 | self.joint_names = joint_names 34 | self.actuator_names = actuator_names 35 | 36 | self.n_joints = len(self.joint_names) 37 | # Get the site ID of the end-effector from its name in the Mujoco model 38 | self.eef_id = self.model_names.site_name2id[self.eef_name] 39 | # Get the joint IDs corresponding to the joint names 40 | self.joint_ids = [self.model_names.joint_name2id[name] for name in self.joint_names] 41 | # Get the qpos (position) IDs corresponding to the joint IDs 42 | self.jnt_qpos_ids = [self.model.jnt_qposadr[id] for id in self.joint_ids] 43 | # Get the dof (degree of freedom) IDs corresponding to the joint IDs 44 | self.jnt_dof_ids = [self.model.jnt_dofadr[id] for id in self.joint_ids] 45 | # Get the actuator IDs corresponding to the actuator names 46 | self.actuator_ids = [self.model_names.actuator_name2id[name] for name in actuator_names] 47 | 48 | def run(self, target: np.ndarray, ctrl: np.ndarray) -> None: 49 | """ 50 | Run the robot controller. 51 | 52 | Parameters: 53 | target (numpy.ndarray): The desired target joint positions or states for the robot. 54 | The size of `target` should be (n_joints,) where n_joints is the number of robot joints. 55 | ctrl (numpy.ndarray): Control signals for the robot actuators from `mujoco._structs.MjData.ctrl` of size (nu,). 56 | """ 57 | # Set the control signals for the actuators to the desired target joint positions or states 58 | ctrl[self.actuator_ids] = target 59 | 60 | def reset(self) -> None: 61 | """ 62 | Reset the controller's internal state to its initial configuration. 63 | 64 | Notes: 65 | This method is intended to reset any internal variables or states of the controller. 66 | """ 67 | # Raise a NotImplementedError to indicate that the method needs to be implemented in the derived classes 68 | raise NotImplementedError("The 'reset' method must be implemented in the derived class.") 69 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/joint_effort_controller.py: -------------------------------------------------------------------------------- 1 | # Import necessary modules and classes 2 | from homestri_ur5e_rl.controllers.joint_controller import JointController 3 | from mujoco._structs import MjModel 4 | from mujoco._structs import MjData 5 | from homestri_ur5e_rl.utils.mujoco_utils import MujocoModelNames 6 | from typing import List 7 | import numpy as np 8 | 9 | class JointEffortController(JointController): 10 | def __init__( 11 | self, 12 | model: MjModel, 13 | data: MjData, 14 | model_names: MujocoModelNames, 15 | eef_name: str, 16 | joint_names: List[str], 17 | actuator_names: List[str], 18 | min_effort: List[float], 19 | max_effort: List[float], 20 | ) -> None: 21 | """ 22 | Joint Effort Controller class to control the robot's joints by specifying joint efforts (controller joint torques). 23 | 24 | Parameters: 25 | model (mujoco._structs.MjModel): Mujoco model representing the robot. 26 | data (mujoco._structs.MjData): Mujoco data associated with the model. 27 | model_names (homestri_ur5e_rl.utils.mujoco_utils.MujocoModelNames): Names of different components in the Mujoco model. 28 | eef_name (str): Name of the end-effector in the Mujoco model. 29 | joint_names (List[str]): List of joint names for the robot. 30 | actuator_names (List[str]): List of actuator names for the robot. 31 | min_effort (List[float]): List of minimum allowable effort (torque) for each joint. 32 | max_effort (List[float]): List of maximum allowable effort (torque) for each joint. 33 | """ 34 | super().__init__(model, data, model_names, eef_name, joint_names, actuator_names) 35 | 36 | # Convert the minimum and maximum effort lists to NumPy arrays for easier manipulation 37 | self.min_effort = np.array(min_effort) 38 | self.max_effort = np.array(max_effort) 39 | 40 | def run(self, target: np.ndarray, ctrl: np.ndarray) -> None: 41 | """ 42 | Run the joint effort controller to control the robot's joints. 43 | 44 | Parameters: 45 | target (numpy.ndarray): The desired target efforts (controller joint torques) for the actuators. 46 | It should have the same length as the number of actuators in the robot. 47 | ctrl (numpy.ndarray): Control signals for the robot actuators. 48 | It should have the same length as the number of actuators in the robot. 49 | 50 | Notes: 51 | - The controller clamps the target efforts to ensure they are within the allowable effort range. 52 | If a target effort exceeds the specified maximum or minimum effort for a joint, it will be clipped to the corresponding bound. 53 | """ 54 | 55 | # Clip the target efforts to ensure they are within the allowable effort range 56 | target_effort = np.clip(target, self.min_effort, self.max_effort) 57 | 58 | # Get the gravity compensation for each joint 59 | gravity_compensation = self.data.qfrc_bias[self.jnt_dof_ids] 60 | 61 | # Add gravity compensation to the target effort 62 | target_effort += gravity_compensation 63 | 64 | # Call the base controller's `run_controller` method with the clipped target efforts 65 | super().run(target_effort, ctrl) 66 | 67 | def reset(self) -> None: 68 | """ 69 | Reset the controller's internal state to its initial configuration. 70 | 71 | Notes: 72 | This method does not perform any actions for resetting the controller's state. 73 | It is intended to be overridden in the derived classes if any specific reset behavior is required. 74 | """ 75 | pass # The reset method does not perform any actions for this controller 76 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/joint_position_controller.py: -------------------------------------------------------------------------------- 1 | # Import necessary modules and classes 2 | from homestri_ur5e_rl.controllers.joint_effort_controller import JointEffortController 3 | from homestri_ur5e_rl.utils.mujoco_utils import ( 4 | get_fullM, 5 | ) 6 | from mujoco._structs import MjModel 7 | from mujoco._structs import MjData 8 | from homestri_ur5e_rl.utils.mujoco_utils import MujocoModelNames 9 | from typing import List 10 | import numpy as np 11 | 12 | class JointPositionController(JointEffortController): 13 | def __init__( 14 | self, 15 | model: MjModel, 16 | data: MjData, 17 | model_names: MujocoModelNames, 18 | eef_name: str, 19 | joint_names: List[str], 20 | actuator_names: List[str], 21 | kp: List[float], 22 | kd: List[float], 23 | min_effort: List[float], 24 | max_effort: List[float], 25 | min_position: List[float], 26 | max_position: List[float], 27 | ) -> None: 28 | """ 29 | Joint Position Controller class to control the robot's joints using PD control for position control with gravity compensation. 30 | 31 | Parameters: 32 | model (mujoco._structs.MjModel): Mujoco model representing the robot. 33 | data (mujoco._structs.MjData): Mujoco data associated with the model. 34 | model_names (homestri_ur5e_rl.utils.mujoco_utils.MujocoModelNames): Names of different components in the Mujoco model. 35 | eef_name (str): Name of the end-effector in the Mujoco model. 36 | joint_names (List[str]): List of joint names for the robot. 37 | actuator_names (List[str]): List of actuator names for the robot. 38 | kp (List[float]): Proportional gains for the PD controller for each joint. 39 | kd (List[float]): Derivative gains for the PD controller for each joint. A good starting point value for kd is sqrt(kp) 40 | min_effort (List[float]): List of minimum allowable effort (torque) for each joint. 41 | max_effort (List[float]): List of maximum allowable effort (torque) for each joint. 42 | min_position (List[float]): List of minimum allowable position for each joint. 43 | max_position (List[float]): List of maximum allowable position for each joint. 44 | """ 45 | super().__init__(model, data, model_names, eef_name, joint_names, actuator_names, min_effort, max_effort) 46 | 47 | self.min_position = np.array(min_position) 48 | self.max_position = np.array(max_position) 49 | self.kp = np.array(kp) 50 | self.kd = np.array(kd) 51 | 52 | def run(self, target: np.ndarray, ctrl: np.ndarray) -> None: 53 | """ 54 | Run the joint position controller to control the robot's joints using PD control with gravity compensation. 55 | 56 | Parameters: 57 | target (numpy.ndarray): The desired target joint positions for the robot. It should have the same length as the number of controlled joints in the robot. 58 | ctrl (numpy.ndarray): Control signals for the robot actuators. It should have the same length as the number of actuators in the robot. 59 | 60 | Notes: 61 | The controller sets the control signals (efforts, i.e., controller joint torques) for the actuators specified by the 'actuator_ids' attribute based on the PD control with gravity compensation to achieve the desired joint positions specified in 'target'. 62 | """ 63 | # Get the joint space inertia matrix 64 | M_full = get_fullM(self.model, self.data) 65 | M = M_full[self.jnt_dof_ids, :][:, self.jnt_dof_ids] 66 | 67 | # Clip the target joint positions to ensure they are within the allowable position range 68 | target_positions = np.clip(target, self.min_position, self.max_position) 69 | 70 | # Get the current joint positions and velocities from Mujoco 71 | current_positions = self.data.qpos[self.jnt_qpos_ids] 72 | current_velocities = self.data.qvel[self.jnt_dof_ids] 73 | 74 | # Calculate the error term for PD control 75 | position_error = target_positions - current_positions 76 | velocity_error = -current_velocities 77 | 78 | # Calculate the control effort (u) using PD control 79 | u = self.kp * position_error + self.kd * velocity_error 80 | 81 | # Apply the joint space inertia matrix to obtain the desired joint effort (torques) 82 | target_effort = np.dot(M, u) 83 | 84 | # Pass the calculated joint torques to the parent's run_controller 85 | super().run(target_effort, ctrl) 86 | 87 | 88 | def reset(self) -> None: 89 | """ 90 | Reset the controller's internal state to its initial configuration. 91 | 92 | Notes: 93 | This method does not perform any actions for resetting the controller's state. 94 | It is intended to be overridden in the derived classes if any specific reset behavior is required. 95 | """ 96 | pass # The reset method does not perform any actions for this controller 97 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/joint_velocity_controller.py: -------------------------------------------------------------------------------- 1 | # Import necessary modules and classes 2 | from homestri_ur5e_rl.controllers.joint_effort_controller import JointEffortController 3 | from mujoco._structs import MjModel 4 | from mujoco._structs import MjData 5 | from homestri_ur5e_rl.utils.mujoco_utils import MujocoModelNames 6 | from typing import List 7 | import numpy as np 8 | 9 | class JointVelocityController(JointEffortController): 10 | def __init__( 11 | self, 12 | model: MjModel, 13 | data: MjData, 14 | model_names: MujocoModelNames, 15 | eef_name: str, 16 | joint_names: List[str], 17 | actuator_names: List[str], 18 | min_effort: List[float], 19 | max_effort: List[float], 20 | min_velocity: List[float], 21 | max_velocity: List[float], 22 | kp: List[float], 23 | ki: List[float], 24 | kd: List[float], 25 | antiwindup: bool = False, 26 | max_integral: float = 10.0, 27 | ) -> None: 28 | """ 29 | Joint Velocity Controller class to control the robot's joints using PID control for velocity control with gravity compensation. 30 | 31 | Parameters: 32 | model (mujoco._structs.MjModel): Mujoco model representing the robot. 33 | data (mujoco._structs.MjData): Mujoco data associated with the model. 34 | model_names (homestri_ur5e_rl.utils.mujoco_utils.MujocoModelNames): Names of different components in the Mujoco model. 35 | eef_name (str): Name of the end-effector in the Mujoco model. 36 | joint_names (List[str]): List of joint names for the robot. 37 | actuator_names (List[str]): List of actuator names for the robot. 38 | min_effort (List[float]): List of minimum allowable effort (torque) for each joint. 39 | max_effort (List[float]): List of maximum allowable effort (torque) for each joint. 40 | min_velocity (List[float]): List of minimum allowable velocity for each joint. 41 | max_velocity (List[float]): List of maximum allowable velocity for each joint. 42 | kp (List[float]): Proportional gains for the PID controller for each joint. 43 | kd (List[float]): Derivative gains for the PID controller for each joint. 44 | ki (List[float]): Integral gains for the PID controller for each joint. 45 | antiwindup (bool, optional): Enable anti-windup to limit the integral term. Defaults to False. 46 | max_integral (float, optional): Maximum allowed value for the integral term. Defaults to 10.0. 47 | """ 48 | super().__init__(model, data, model_names, eef_name, joint_names, actuator_names, min_effort, max_effort) 49 | 50 | self.min_velocity = np.array(min_velocity) 51 | self.max_velocity = np.array(max_velocity) 52 | self.kp = np.array(kp) 53 | self.kd = np.array(kd) 54 | self.ki = np.array(ki) 55 | self.antiwindup = antiwindup 56 | self.max_integral = max_integral 57 | 58 | # Initialize previous error and integral term arrays for each joint 59 | self.prev_error = np.zeros(self.n_joints) 60 | self.integral = np.zeros(self.n_joints) 61 | 62 | def run(self, target: np.ndarray, ctrl: np.ndarray) -> None: 63 | """ 64 | Run the joint velocity controller to control the robot's joints using PID control with gravity compensation. 65 | 66 | Parameters: 67 | target (numpy.ndarray): The desired target joint velocities for the robot. It should have the same length as the number of controlled joints in the robot. 68 | ctrl (numpy.ndarray): Control signals for the robot actuators. It should have the same length as the number of actuators in the robot. 69 | 70 | Notes: 71 | The controller sets the control signals (efforts, i.e., controller joint torques) for the actuators specified by the 'actuator_ids' attribute based on the PID control and gravity compensation to achieve the desired joint velocities specified in 'target'. 72 | """ 73 | # Clip the target joint velocities to ensure they are within the allowable velocity range 74 | target_velocities = np.clip(target, self.min_velocity, self.max_velocity) 75 | 76 | # Calculate the error term for PID control 77 | current_velocities = self.data.qvel[self.jnt_dof_ids] 78 | error = target_velocities - current_velocities 79 | 80 | # Integral term 81 | if self.antiwindup: 82 | self.integral += error 83 | # Anti-windup to limit the integral term 84 | self.integral = np.clip(self.integral, -self.max_integral, self.max_integral) 85 | else: 86 | self.integral += error 87 | 88 | # Derivative term 89 | derivative = error - self.prev_error 90 | 91 | # Calculate the desired joint effort (torques) using PID control 92 | target_effort = self.kp * error + self.kd * derivative + self.ki * self.integral 93 | 94 | # Save the current error for the next iteration 95 | self.prev_error = error 96 | 97 | # Pass the calculated joint torques to the parent's run_controller 98 | super().run(target_effort, ctrl) 99 | 100 | def reset(self) -> None: 101 | """ 102 | Reset the controller's internal state to its initial configuration. 103 | 104 | Notes: 105 | This method does not perform any actions for resetting the controller's state. 106 | It is intended to be overridden in the derived classes if any specific reset behavior is required. 107 | """ 108 | pass # The reset method does not perform any actions for this controller 109 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/controllers/operational_space_controller.py: -------------------------------------------------------------------------------- 1 | from homestri_ur5e_rl.controllers.joint_effort_controller import JointEffortController 2 | from homestri_ur5e_rl.controllers.joint_velocity_controller import JointVelocityController 3 | from homestri_ur5e_rl.controllers.force_torque_sensor_controller import ForceTorqueSensorController 4 | from homestri_ur5e_rl.utils.mujoco_utils import ( 5 | get_site_jac, 6 | get_fullM, 7 | ) 8 | from homestri_ur5e_rl.utils.transform_utils import ( 9 | mat2quat, 10 | ) 11 | from homestri_ur5e_rl.utils.controller_utils import ( 12 | task_space_inertia_matrix, 13 | pose_error, 14 | ) 15 | from homestri_ur5e_rl.utils.pid_controller_utils import ( 16 | SpatialPIDController, 17 | ) 18 | from mujoco._structs import MjModel 19 | from mujoco._structs import MjData 20 | from homestri_ur5e_rl.utils.mujoco_utils import MujocoModelNames 21 | from typing import List 22 | import numpy as np 23 | from enum import Enum 24 | 25 | class TargetType(Enum): 26 | POSE = 0 27 | TWIST = 1 28 | 29 | class OperationalSpaceController(JointEffortController): 30 | def __init__( 31 | self, 32 | model: MjModel, 33 | data: MjData, 34 | model_names: MujocoModelNames, 35 | eef_name: str, 36 | joint_names: List[str], 37 | actuator_names: List[str], 38 | min_effort: List[float], 39 | max_effort: List[float], 40 | target_type: TargetType, 41 | kp: float, 42 | ko: float, 43 | kv: float, 44 | vmax_xyz: float, 45 | vmax_abg: float, 46 | null_damp_kv: float, 47 | ) -> None: 48 | """ 49 | Operational Space Controller class to control the robot's joints using operational space control with gravity compensation. 50 | 51 | Parameters: 52 | model (mujoco._structs.MjModel): Mujoco model representing the robot. 53 | data (mujoco._structs.MjData): Mujoco data associated with the model. 54 | model_names (homestri_ur5e_rl.utils.mujoco_utils.MujocoModelNames): Names of different components in the Mujoco model. 55 | eef_name (str): Name of the end-effector in the Mujoco model. 56 | joint_names (List[str]): List of joint names for the robot. 57 | actuator_names (List[str]): List of actuator names for the robot. 58 | min_effort (List[float]): List of minimum allowable effort (torque) for each joint. 59 | max_effort (List[float]): List of maximum allowable effort (torque) for each joint. 60 | target_type (TargetType): The type of target input for the controller (POSE or TWIST). 61 | kp (float): Proportional gain for the PD controller in position space. 62 | ko (float): Proportional gain for the PD controller in orientation space. 63 | kv (float): Velocity gain for the PD controller. 64 | vmax_xyz (float): Maximum velocity for linear position control. 65 | vmax_abg (float): Maximum velocity for orientation control. 66 | ctrl_dof (List[bool]): Control degrees of freedom for each joint (True if controlled, False if uncontrolled). 67 | null_damp_kv (float): Damping gain for null space control. 68 | """ 69 | super().__init__(model, data, model_names, eef_name, joint_names, actuator_names, min_effort, max_effort) 70 | 71 | self.target_type = target_type 72 | self.kp = kp 73 | self.ko = ko 74 | self.kv = kv 75 | self.vmax_xyz = vmax_xyz 76 | self.vmax_abg = vmax_abg 77 | self.null_damp_kv = null_damp_kv 78 | 79 | self.task_space_gains = np.array([self.kp] * 3 + [self.ko] * 3) 80 | self.lamb = self.task_space_gains / self.kv 81 | self.sat_gain_xyz = vmax_xyz / self.kp * self.kv 82 | self.sat_gain_abg = vmax_abg / self.ko * self.kv 83 | self.scale_xyz = vmax_xyz / self.kp * self.kv 84 | self.scale_abg = vmax_abg / self.ko * self.kv 85 | 86 | def run(self, target: np.ndarray, ctrl: np.ndarray) -> None: 87 | """ 88 | Run the operational space controller to control the robot's joints using operational space control with gravity compensation. 89 | 90 | Parameters: 91 | target (numpy.ndarray): The desired target input for the controller. 92 | ctrl (numpy.ndarray): Control signals for the robot actuators. 93 | 94 | Notes: 95 | The controller sets the control signals (efforts, i.e., controller joint torques) for the actuators based on operational space control to achieve the desired target (either pose or twist). 96 | """ 97 | # Get the Jacobian matrix for the end-effector. 98 | J = get_site_jac(self.model, self.data, self.eef_id) 99 | J = J[:, self.jnt_dof_ids] 100 | 101 | # Get the mass matrix and its inverse for the controlled degrees of freedom (DOF) of the robot. 102 | M_full = get_fullM(self.model, self.data) 103 | M = M_full[self.jnt_dof_ids, :][:, self.jnt_dof_ids] 104 | Mx, M_inv = task_space_inertia_matrix(M, J) 105 | 106 | # Get the joint velocities for the controlled DOF. 107 | dq = self.data.qvel[self.jnt_dof_ids] 108 | 109 | # Get the end-effector position, orientation matrix, and twist (spatial velocity). 110 | ee_pos = self.data.site_xpos[self.eef_id] 111 | ee_quat = mat2quat(self.data.site_xmat[self.eef_id].reshape(3, 3)) 112 | ee_pose = np.concatenate([ee_pos, ee_quat]) 113 | ee_twist = J @ dq 114 | 115 | # Initialize the task space control signal (desired end-effector motion). 116 | u_task = np.zeros(6) 117 | 118 | if self.target_type == TargetType.POSE: 119 | # If the target type is pose, the target contains both position and orientation. 120 | target_pose = target 121 | target_twist = np.zeros(6) 122 | 123 | # Scale the task space control signal while ensuring it doesn't exceed the specified velocity limits. 124 | u_task = self._scale_signal_vel_limited(pose_error(ee_pose, target_pose)) 125 | 126 | elif self.target_type == TargetType.TWIST: 127 | # If the target type is twist, the target contains the desired spatial velocity. 128 | target_twist = target 129 | 130 | else: 131 | raise ValueError("Invalid target type: {}".format(self.target_type)) 132 | 133 | # Initialize the joint effort control signal (controller joint torques). 134 | u = np.zeros(self.n_joints) 135 | 136 | if np.all(target_twist == 0): 137 | # If the target twist is zero (no desired motion), apply damping to the controlled DOF. 138 | u -= self.kv * np.dot(M, dq) 139 | else: 140 | # If the target twist is not zero, calculate the task space control signal error. 141 | u_task += self.kv * (ee_twist - target_twist) 142 | 143 | # Compute the joint effort control signal based on the task space control signal. 144 | u -= np.dot(J.T, np.dot(Mx, u_task)) 145 | 146 | # Compute the null space control signal to minimize the joint efforts in the null space of the task. 147 | u_null = np.dot(M, -self.null_damp_kv * dq) 148 | Jbar = np.dot(M_inv, np.dot(J.T, Mx)) 149 | null_filter = np.eye(self.n_joints) - np.dot(J.T, Jbar.T) 150 | u += np.dot(null_filter, u_null) 151 | 152 | # Call the parent class's run method to apply the computed joint efforts to the robot actuators. 153 | super().run(u, ctrl) 154 | 155 | def _scale_signal_vel_limited(self, u_task: np.ndarray) -> np.ndarray: 156 | """ 157 | Scale the control signal such that the arm isn't driven to move faster in position or orientation than the specified vmax values. 158 | 159 | Parameters: 160 | u_task (numpy.ndarray): The task space control signal. 161 | 162 | Returns: 163 | numpy.ndarray: The scaled task space control signal. 164 | """ 165 | norm_xyz = np.linalg.norm(u_task[:3]) 166 | norm_abg = np.linalg.norm(u_task[3:]) 167 | scale = np.ones(6) 168 | if norm_xyz > self.sat_gain_xyz: 169 | scale[:3] *= self.scale_xyz / norm_xyz 170 | if norm_abg > self.sat_gain_abg: 171 | scale[3:] *= self.scale_abg / norm_abg 172 | 173 | return self.kv * scale * self.lamb * u_task 174 | 175 | class ImpedanceController(JointEffortController): 176 | def __init__( 177 | self, 178 | model: MjModel, 179 | data: MjData, 180 | model_names: MujocoModelNames, 181 | eef_name: str, 182 | joint_names: List[str], 183 | actuator_names: List[str], 184 | min_effort: List[float], 185 | max_effort: List[float], 186 | null_damp_kv: float, 187 | ) -> None: 188 | super().__init__( 189 | model, 190 | data, 191 | model_names, 192 | eef_name, 193 | joint_names, 194 | actuator_names, 195 | min_effort, 196 | max_effort, 197 | ) 198 | self.null_damp_kv = null_damp_kv 199 | 200 | def run( 201 | self, 202 | target_pose: np.ndarray, 203 | target_twist: np.ndarray, 204 | kp: np.ndarray, 205 | kd: np.ndarray, 206 | ctrl: np.ndarray, 207 | ) -> None: 208 | # Get the Jacobian matrix for the end-effector. 209 | J_full = get_site_jac(self.model, self.data, self.eef_id) 210 | J = J_full[:, self.jnt_dof_ids] 211 | # Get the joint velocities 212 | dq = self.data.qvel[self.jnt_dof_ids] 213 | # Get the mass matrix and its inverse for the controlled degrees of freedom (DOF) of the robot. 214 | M_full = get_fullM(self.model, self.data) 215 | M = M_full[self.jnt_dof_ids, :][:, self.jnt_dof_ids] 216 | Mx, M_inv = task_space_inertia_matrix(M, J) 217 | 218 | # Get the end-effector pose and twist (spatial velocity). 219 | ee_pos = self.data.site_xpos[self.eef_id] 220 | ee_quat = mat2quat(self.data.site_xmat[self.eef_id].reshape(3, 3)) 221 | ee_pose = np.concatenate([ee_pos, ee_quat]) 222 | ee_twist = J @ dq 223 | 224 | u_task_pos = kp * pose_error(target_pose, ee_pose) 225 | u_task_vel = kd * (target_twist -ee_twist) 226 | u_task = u_task_pos + u_task_vel 227 | 228 | # Compute the joint effort control signal based on the task space control signal. 229 | u = np.dot(J.T, np.dot(Mx, u_task)) 230 | 231 | # Compute the null space control signal to minimize the joint efforts in the null space of the task. 232 | u_null = np.dot(M, -self.null_damp_kv * dq) 233 | Jbar = np.dot(M_inv, np.dot(J.T, Mx)) 234 | null_filter = np.eye(self.n_joints) - np.dot(J.T, Jbar.T) 235 | u += np.dot(null_filter, u_null) 236 | 237 | # Call the parent class's run method to apply the computed joint efforts to the robot actuators. 238 | super().run(u, ctrl) 239 | 240 | class ComplianceController(JointVelocityController): 241 | def __init__( 242 | self, 243 | model: MjModel, 244 | data: MjData, 245 | model_names: MujocoModelNames, 246 | eef_name: str, 247 | joint_names: List[str], 248 | actuator_names: List[str], 249 | min_effort: List[float], 250 | max_effort: List[float], 251 | min_velocity: List[float], 252 | max_velocity: List[float], 253 | kp_jnt_vel: List[float], 254 | ki_jnt_vel: List[float], 255 | kd_jnt_vel: List[float], 256 | kp: np.ndarray, 257 | ki: np.ndarray, 258 | kd: np.ndarray, 259 | control_period: float, 260 | ft_sensor_site: str, 261 | force_sensor_name: str, 262 | torque_sensor_name: str, 263 | subtree_body_name: str, 264 | ft_smoothing_factor: float, 265 | ) -> None: 266 | super().__init__( 267 | model, 268 | data, 269 | model_names, 270 | eef_name, 271 | joint_names, 272 | actuator_names, 273 | min_effort, 274 | max_effort, 275 | min_velocity, 276 | max_velocity, 277 | kp_jnt_vel, 278 | ki_jnt_vel, 279 | kd_jnt_vel, 280 | ) 281 | self.ft_sensor = ForceTorqueSensorController( 282 | model, 283 | data, 284 | model_names, 285 | ft_sensor_site, 286 | force_sensor_name, 287 | torque_sensor_name, 288 | subtree_body_name, 289 | ft_smoothing_factor, 290 | ) 291 | self.spatial_controller = SpatialPIDController(kp, ki, kd) 292 | self.control_period = control_period 293 | 294 | def run( 295 | self, 296 | target_pose: np.ndarray, 297 | target_twist: np.ndarray, 298 | target_wrench: np.ndarray, 299 | stiffness: np.ndarray, 300 | damping: np.ndarray, 301 | error_scale: np.ndarray, 302 | ctrl: np.ndarray, 303 | ) -> None: 304 | # Get the Jacobian matrix for the end-effector. 305 | J_full = get_site_jac(self.model, self.data, self.eef_id) 306 | J = J_full[:, self.jnt_dof_ids] 307 | # Get the joint velocities 308 | dq = self.data.qvel[self.jnt_dof_ids] 309 | 310 | # Get the end-effector pose and twist (spatial velocity). 311 | ee_pos = self.data.site_xpos[self.eef_id] 312 | ee_quat = mat2quat(self.data.site_xmat[self.eef_id].reshape(3, 3)) 313 | ee_pose = np.concatenate([ee_pos, ee_quat]) 314 | ee_twist = J @ dq 315 | 316 | # ee_wrench = self.ft_sensor.get_wrench() 317 | 318 | u_task_pos = stiffness * pose_error(target_pose, ee_pose) 319 | u_task_vel = damping * (target_twist - ee_twist) 320 | # u_task_eff = target_wrench - ee_wrench 321 | 322 | # print(f"u_task_pos: {u_task_pos}, u_task_vel: {u_task_vel}, u_task_eff: {u_task_eff}") 323 | u_task = error_scale * (u_task_pos + u_task_vel) 324 | 325 | 326 | u_task = self.spatial_controller(u_task, self.control_period) 327 | 328 | # print("u_task", "{0:0.7f}".format(u_task[2])) 329 | # print("u_task_eff", "{0:0.7f}".format(u_task_eff[2])) 330 | # print("ee_wrench", "{0:0.7f}".format(ee_wrench[2])) 331 | 332 | # Compute the joint effort control signal based on the task space control signal. 333 | u = np.dot(J.T, u_task) 334 | 335 | u = 0.5 * u * self.control_period 336 | 337 | # Call the parent class's run method to apply the computed joint efforts to the robot actuators. 338 | super().run(u, ctrl) 339 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/__init__.py -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/rethink_pedestal/pedestal.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/rethink_pedestal/pedestal.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/base.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/base.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/base_mount.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/base_mount.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/coupler.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/coupler.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/driver.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/driver.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/follower.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/follower.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/pad.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/pad.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/silicone_pad.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/silicone_pad.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/spring_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/homestri_ur5e_rl/envs/assets/base_robot/assets/robotiq_2f85/spring_link.stl -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/scene.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 9 | 10 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/table.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/assets/base_robot/wire.xml: -------------------------------------------------------------------------------- 1 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 80 | 81 | 82 | 83 | 84 | 85 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/base_robot/__init__.py: -------------------------------------------------------------------------------- 1 | from homestri_ur5e_rl.envs.base_robot.base_robot_env import BaseRobot -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/base_robot/base_robot_env.py: -------------------------------------------------------------------------------- 1 | from os import path 2 | import sys 3 | 4 | import numpy as np 5 | from homestri_ur5e_rl.controllers.operational_space_controller import ImpedanceController, ComplianceController, OperationalSpaceController, TargetType 6 | from homestri_ur5e_rl.controllers.joint_effort_controller import JointEffortController 7 | from homestri_ur5e_rl.controllers.joint_velocity_controller import JointVelocityController 8 | from homestri_ur5e_rl.controllers.joint_position_controller import JointPositionController 9 | from homestri_ur5e_rl.controllers.force_torque_sensor_controller import ForceTorqueSensorController 10 | from gymnasium import spaces 11 | from homestri_ur5e_rl.envs.mujoco.mujoco_env import MujocoEnv 12 | from gymnasium_robotics.utils.mujoco_utils import MujocoModelNames, robot_get_obs 13 | 14 | np.set_printoptions(formatter={'float': lambda x: "{0:0.5f}".format(x)}) 15 | 16 | 17 | DEFAULT_CAMERA_CONFIG = { 18 | "distance": 2.2, 19 | "azimuth": 0.0, 20 | "elevation": -20.0, 21 | "lookat": np.array([0, 0, 1]), 22 | } 23 | 24 | class BaseRobot(MujocoEnv): 25 | metadata = { 26 | "render_modes": [ 27 | "human", 28 | "rgb_array", 29 | "depth_array", 30 | ], 31 | "render_fps": 12, 32 | } 33 | 34 | def __init__( 35 | self, 36 | model_path="../assets/base_robot/scene.xml", 37 | frame_skip=40, 38 | default_camera_config: dict = DEFAULT_CAMERA_CONFIG, 39 | **kwargs, 40 | ): 41 | xml_file_path = path.join( 42 | path.dirname(path.realpath(__file__)), 43 | model_path, 44 | ) 45 | 46 | observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(28,), dtype=np.float64) 47 | 48 | super().__init__( 49 | xml_file_path, 50 | frame_skip, 51 | observation_space, 52 | default_camera_config=default_camera_config, 53 | **kwargs, 54 | ) 55 | 56 | self.init_qvel = self.data.qvel.copy() 57 | self.init_ctrl = self.data.ctrl.copy() 58 | 59 | 60 | config_path = path.join( 61 | path.dirname(__file__), 62 | "../assets/base_robot/robot_config.xml", 63 | ) 64 | 65 | self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(7,), dtype=np.float64) 66 | self.model_names = MujocoModelNames(self.model) 67 | 68 | 69 | 70 | 71 | 72 | # self.controller = ImpedanceController( 73 | # model=self.model, 74 | # data=self.data, 75 | # model_names=self.model_names, 76 | # eef_name='robot0:eef_site', 77 | # joint_names=[ 78 | # 'robot0:ur5e:shoulder_pan_joint', 79 | # 'robot0:ur5e:shoulder_lift_joint', 80 | # 'robot0:ur5e:elbow_joint', 81 | # 'robot0:ur5e:wrist_1_joint', 82 | # 'robot0:ur5e:wrist_2_joint', 83 | # 'robot0:ur5e:wrist_3_joint', 84 | # ], 85 | # actuator_names=[ 86 | # 'robot0:ur5e:shoulder_pan', 87 | # 'robot0:ur5e:shoulder_lift', 88 | # 'robot0:ur5e:elbow', 89 | # 'robot0:ur5e:wrist_1', 90 | # 'robot0:ur5e:wrist_2', 91 | # 'robot0:ur5e:wrist_3', 92 | # ], 93 | # min_effort=[-150, -150, -150, -150, -150, -150], 94 | # max_effort=[150, 150, 150, 150, 150, 150], 95 | # null_damp_kv=10, 96 | # ) 97 | 98 | self.controller = OperationalSpaceController( 99 | model=self.model, 100 | data=self.data, 101 | model_names=self.model_names, 102 | eef_name='robot0:eef_site', 103 | joint_names=[ 104 | 'robot0:ur5e:shoulder_pan_joint', 105 | 'robot0:ur5e:shoulder_lift_joint', 106 | 'robot0:ur5e:elbow_joint', 107 | 'robot0:ur5e:wrist_1_joint', 108 | 'robot0:ur5e:wrist_2_joint', 109 | 'robot0:ur5e:wrist_3_joint', 110 | ], 111 | actuator_names=[ 112 | 'robot0:ur5e:shoulder_pan', 113 | 'robot0:ur5e:shoulder_lift', 114 | 'robot0:ur5e:elbow', 115 | 'robot0:ur5e:wrist_1', 116 | 'robot0:ur5e:wrist_2', 117 | 'robot0:ur5e:wrist_3', 118 | ], 119 | min_effort=[-150, -150, -150, -150, -150, -150], 120 | max_effort=[150, 150, 150, 150, 150, 150], 121 | target_type=TargetType.TWIST, 122 | kp=200.0, 123 | ko=200.0, 124 | kv=50.0, 125 | vmax_xyz=0.2, 126 | vmax_abg=1, 127 | null_damp_kv=10, 128 | ) 129 | 130 | 131 | # self.controller = JointEffortController( 132 | # model=self.model, 133 | # data=self.data, 134 | # model_names=self.model_names, 135 | # eef_name='robot0:eef_site', 136 | # joint_names=[ 137 | # 'robot0:ur5e:shoulder_pan_joint', 138 | # 'robot0:ur5e:shoulder_lift_joint', 139 | # 'robot0:ur5e:elbow_joint', 140 | # 'robot0:ur5e:wrist_1_joint', 141 | # 'robot0:ur5e:wrist_2_joint', 142 | # 'robot0:ur5e:wrist_3_joint', 143 | # ], 144 | # actuator_names=[ 145 | # 'robot0:ur5e:shoulder_pan', 146 | # 'robot0:ur5e:shoulder_lift', 147 | # 'robot0:ur5e:elbow', 148 | # 'robot0:ur5e:wrist_1', 149 | # 'robot0:ur5e:wrist_2', 150 | # 'robot0:ur5e:wrist_3', 151 | # ], 152 | # min_effort=[-150, -150, -150, -150, -150, -150], 153 | # max_effort=[150, 150, 150, 150, 150, 150], 154 | # ) 155 | 156 | # self.controller = JointVelocityController( 157 | # model=self.model, 158 | # data=self.data, 159 | # model_names=self.model_names, 160 | # eef_name='robot0:eef_site', 161 | # joint_names=[ 162 | # 'robot0:ur5e:shoulder_pan_joint', 163 | # 'robot0:ur5e:shoulder_lift_joint', 164 | # 'robot0:ur5e:elbow_joint', 165 | # 'robot0:ur5e:wrist_1_joint', 166 | # 'robot0:ur5e:wrist_2_joint', 167 | # 'robot0:ur5e:wrist_3_joint', 168 | # ], 169 | # actuator_names=[ 170 | # 'robot0:ur5e:shoulder_pan', 171 | # 'robot0:ur5e:shoulder_lift', 172 | # 'robot0:ur5e:elbow', 173 | # 'robot0:ur5e:wrist_1', 174 | # 'robot0:ur5e:wrist_2', 175 | # 'robot0:ur5e:wrist_3', 176 | # ], 177 | # min_effort=[-150, -150, -150, -150, -150, -150], 178 | # max_effort=[150, 150, 150, 150, 150, 150], 179 | # min_velocity=[-1, -1, -1, -1, -1, -1], 180 | # max_velocity=[1, 1, 1, 1, 1, 1], 181 | # kp=[100, 100, 100, 100, 100, 100], 182 | # ki=0, 183 | # kd=0, 184 | # ) 185 | 186 | # self.controller = JointPositionController( 187 | # model=self.model, 188 | # data=self.data, 189 | # model_names=self.model_names, 190 | # eef_name='robot0:eef_site', 191 | # joint_names=[ 192 | # 'robot0:ur5e:shoulder_pan_joint', 193 | # 'robot0:ur5e:shoulder_lift_joint', 194 | # 'robot0:ur5e:elbow_joint', 195 | # 'robot0:ur5e:wrist_1_joint', 196 | # 'robot0:ur5e:wrist_2_joint', 197 | # 'robot0:ur5e:wrist_3_joint', 198 | # ], 199 | # actuator_names=[ 200 | # 'robot0:ur5e:shoulder_pan', 201 | # 'robot0:ur5e:shoulder_lift', 202 | # 'robot0:ur5e:elbow', 203 | # 'robot0:ur5e:wrist_1', 204 | # 'robot0:ur5e:wrist_2', 205 | # 'robot0:ur5e:wrist_3', 206 | # ], 207 | # min_effort=[-150, -150, -150, -150, -150, -150], 208 | # max_effort=[150, 150, 150, 150, 150, 150], 209 | # min_position=[-1, -1, -1, -1, -1, -1], 210 | # max_position=[1, 1, 1, 1, 1, 1], 211 | # kp=[100, 100, 100, 100, 100, 100], 212 | # kd=[20, 20, 20, 20, 20, 20], 213 | # ) 214 | 215 | 216 | self.init_qpos_config = { 217 | "robot0:ur5e:shoulder_pan_joint": 0, 218 | "robot0:ur5e:shoulder_lift_joint": -np.pi / 2.0, 219 | "robot0:ur5e:elbow_joint": -np.pi / 2.0, 220 | "robot0:ur5e:wrist_1_joint": -np.pi / 2.0, 221 | "robot0:ur5e:wrist_2_joint": np.pi / 2.0, 222 | "robot0:ur5e:wrist_3_joint": 0, 223 | } 224 | for joint_name, joint_pos in self.init_qpos_config.items(): 225 | joint_id = self.model_names.joint_name2id[joint_name] 226 | qpos_id = self.model.jnt_qposadr[joint_id] 227 | self.init_qpos[qpos_id] = joint_pos 228 | 229 | 230 | 231 | # self.controller = ComplianceController( 232 | # model=self.model, 233 | # data=self.data, 234 | # model_names=self.model_names, 235 | # eef_name='robot0:eef_site', 236 | # joint_names=[ 237 | # 'robot0:ur5e:shoulder_pan_joint', 238 | # 'robot0:ur5e:shoulder_lift_joint', 239 | # 'robot0:ur5e:elbow_joint', 240 | # 'robot0:ur5e:wrist_1_joint', 241 | # 'robot0:ur5e:wrist_2_joint', 242 | # 'robot0:ur5e:wrist_3_joint', 243 | # ], 244 | # actuator_names=[ 245 | # 'robot0:ur5e:shoulder_pan', 246 | # 'robot0:ur5e:shoulder_lift', 247 | # 'robot0:ur5e:elbow', 248 | # 'robot0:ur5e:wrist_1', 249 | # 'robot0:ur5e:wrist_2', 250 | # 'robot0:ur5e:wrist_3', 251 | # ], 252 | # min_effort=[-150, -150, -150, -150, -150, -150], 253 | # max_effort=[150, 150, 150, 150, 150, 150], 254 | # min_velocity=[-1, -1, -1, -1, -1, -1], 255 | # max_velocity=[1, 1, 1, 1, 1, 1], 256 | # kp_jnt_vel=[100, 100, 100, 100, 100, 100], 257 | # ki_jnt_vel=0, 258 | # kd_jnt_vel=0, 259 | # kp=[10, 10, 10, 10, 10, 10], 260 | # ki=[0, 0, 0, 0, 0, 0], 261 | # kd=[3, 3, 3, 8, 8, 8], 262 | # control_period=self.model.opt.timestep, 263 | # ft_sensor_site='robot0:eef_site', 264 | # force_sensor_name='robot0:eef_force', 265 | # torque_sensor_name='robot0:eef_torque', 266 | # subtree_body_name='robot0:ur5e:wrist_3_link', 267 | # ft_smoothing_factor=0, 268 | # ) 269 | 270 | def step(self, action): 271 | # target_pose = np.array([-0.1,-0.6,1.1,0,1,0,0]) 272 | # target_twist = np.array([0,0,0,0,0,0]) 273 | # target_wrench = np.array([0,0,0,0,0,0]) 274 | 275 | for i in range(self.frame_skip): 276 | ctrl = self.data.ctrl.copy() 277 | # self.controller.run( 278 | # target_pose, 279 | # target_twist, 280 | # target_wrench, 281 | # np.array([200,200,200,200,200,200]), 282 | # np.array([50,50,50,50,50,50]), 283 | # 1, 284 | # ctrl 285 | # ) 286 | 287 | print(action) 288 | 289 | self.controller.run( 290 | action[:6], 291 | ctrl 292 | ) 293 | 294 | ctrl[6] = action[6] 295 | 296 | self.do_simulation(ctrl, n_frames=1) 297 | 298 | if self.render_mode == "human": 299 | self.render() 300 | 301 | obs = self._get_obs() 302 | 303 | return obs, 0.0, False, False, {} 304 | 305 | def _get_obs(self): 306 | # Gather simulated observation 307 | robot_qpos, robot_qvel = robot_get_obs( 308 | self.model, self.data, self.model_names.joint_names 309 | ) 310 | 311 | return np.concatenate((robot_qpos.copy(), robot_qvel.copy())) 312 | 313 | def reset_model(self): 314 | qpos = self.init_qpos 315 | qvel = self.init_qvel 316 | self.set_state(qpos, qvel) 317 | obs = self._get_obs() 318 | 319 | return obs 320 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/mujoco/mujoco_env.py: -------------------------------------------------------------------------------- 1 | from os import path 2 | from typing import Any, Dict, Optional, Tuple, Union 3 | 4 | import numpy as np 5 | from numpy.typing import NDArray 6 | 7 | import gymnasium as gym 8 | from gymnasium import error, logger, spaces 9 | from gymnasium.spaces import Space 10 | 11 | 12 | try: 13 | import mujoco_py 14 | except ImportError as e: 15 | MUJOCO_PY_IMPORT_ERROR = e 16 | else: 17 | MUJOCO_PY_IMPORT_ERROR = None 18 | 19 | try: 20 | import mujoco 21 | except ImportError as e: 22 | MUJOCO_IMPORT_ERROR = e 23 | else: 24 | MUJOCO_IMPORT_ERROR = None 25 | 26 | 27 | DEFAULT_SIZE = 480 28 | 29 | 30 | class BaseMujocoEnv(gym.Env[NDArray[np.float64], NDArray[np.float32]]): 31 | """Superclass for all MuJoCo environments.""" 32 | 33 | def __init__( 34 | self, 35 | model_path, 36 | frame_skip, 37 | observation_space: Space, 38 | render_mode: Optional[str] = None, 39 | width: int = DEFAULT_SIZE, 40 | height: int = DEFAULT_SIZE, 41 | camera_id: Optional[int] = None, 42 | camera_name: Optional[str] = None, 43 | ): 44 | """Base abstract class for mujoco based environments. 45 | 46 | Args: 47 | model_path: Path to the MuJoCo Model. 48 | frame_skip: Number of MuJoCo simulation steps per gym `step()`. 49 | observation_space: The observation space of the environment. 50 | render_mode: The `render_mode` used. 51 | width: The width of the render window. 52 | height: The height of the render window. 53 | camera_id: The camera ID used. 54 | camera_name: The name of the camera used (can not be used in conjunction with `camera_id`). 55 | 56 | Raises: 57 | OSError: when the `model_path` does not exist. 58 | error.DependencyNotInstalled: When `mujoco` is not installed. 59 | """ 60 | if model_path.startswith(".") or model_path.startswith("/"): 61 | self.fullpath = model_path 62 | elif model_path.startswith("~"): 63 | self.fullpath = path.expanduser(model_path) 64 | else: 65 | self.fullpath = path.join(path.dirname(__file__), "assets", model_path) 66 | if not path.exists(self.fullpath): 67 | raise OSError(f"File {self.fullpath} does not exist") 68 | 69 | self.width = width 70 | self.height = height 71 | # may use width and height 72 | self.model, self.data = self._initialize_simulation() 73 | 74 | self.init_qpos = self.data.qpos.ravel().copy() 75 | self.init_qvel = self.data.qvel.ravel().copy() 76 | 77 | self.frame_skip = frame_skip 78 | 79 | assert self.metadata["render_modes"] == [ 80 | "human", 81 | "rgb_array", 82 | "depth_array", 83 | ], self.metadata["render_modes"] 84 | if "render_fps" in self.metadata: 85 | assert ( 86 | int(np.round(1.0 / self.dt)) == self.metadata["render_fps"] 87 | ), f'Expected value: {int(np.round(1.0 / self.dt))}, Actual value: {self.metadata["render_fps"]}' 88 | 89 | self.observation_space = observation_space 90 | self._set_action_space() 91 | 92 | self.render_mode = render_mode 93 | self.camera_name = camera_name 94 | self.camera_id = camera_id 95 | 96 | def _set_action_space(self): 97 | bounds = self.model.actuator_ctrlrange.copy().astype(np.float32) 98 | low, high = bounds.T 99 | self.action_space = spaces.Box(low=low, high=high, dtype=np.float32) 100 | return self.action_space 101 | 102 | # methods to override: 103 | # ---------------------------- 104 | def step( 105 | self, action: NDArray[np.float32] 106 | ) -> Tuple[NDArray[np.float64], np.float64, bool, bool, Dict[str, np.float64]]: 107 | raise NotImplementedError 108 | 109 | def reset_model(self) -> NDArray[np.float64]: 110 | """ 111 | Reset the robot degrees of freedom (qpos and qvel). 112 | Implement this in each subclass. 113 | """ 114 | raise NotImplementedError 115 | 116 | def _initialize_simulation(self) -> Tuple[Any, Any]: 117 | """ 118 | Initialize MuJoCo simulation data structures mjModel and mjData. 119 | """ 120 | raise NotImplementedError 121 | 122 | def _reset_simulation(self) -> None: 123 | """ 124 | Reset MuJoCo simulation data structures, mjModel and mjData. 125 | """ 126 | raise NotImplementedError 127 | 128 | def _step_mujoco_simulation(self, ctrl, n_frames) -> None: 129 | """ 130 | Step over the MuJoCo simulation. 131 | """ 132 | raise NotImplementedError 133 | 134 | def render(self) -> Union[NDArray[np.float64], None]: 135 | """ 136 | Render a frame from the MuJoCo simulation as specified by the render_mode. 137 | """ 138 | raise NotImplementedError 139 | 140 | # ----------------------------- 141 | def _get_reset_info(self) -> Dict[str, float]: 142 | """Function that generates the `info` that is returned during a `reset()`.""" 143 | return {} 144 | 145 | def reset( 146 | self, 147 | *, 148 | seed: Optional[int] = None, 149 | options: Optional[dict] = None, 150 | ): 151 | super().reset(seed=seed) 152 | 153 | self._reset_simulation() 154 | 155 | ob = self.reset_model() 156 | info = self._get_reset_info() 157 | 158 | if self.render_mode == "human": 159 | self.render() 160 | return ob, info 161 | 162 | def set_state(self, qpos, qvel) -> None: 163 | """ 164 | Set the joints position qpos and velocity qvel of the model. Override this method depending on the MuJoCo bindings used. 165 | """ 166 | assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,) 167 | 168 | @property 169 | def dt(self) -> float: 170 | return self.model.opt.timestep * self.frame_skip 171 | 172 | def do_simulation(self, ctrl, n_frames) -> None: 173 | """ 174 | Step the simulation n number of frames and applying a control action. 175 | """ 176 | # Check control input is contained in the action space 177 | if np.array(ctrl).shape != (self.model.nu,): 178 | raise ValueError( 179 | f"Action dimension mismatch. Expected {(self.model.nu,)}, found {np.array(ctrl).shape}" 180 | ) 181 | self._step_mujoco_simulation(ctrl, n_frames) 182 | 183 | def close(self): 184 | """Close all processes like rendering contexts""" 185 | raise NotImplementedError 186 | 187 | def get_body_com(self, body_name) -> NDArray[np.float64]: 188 | """Return the cartesian position of a body frame""" 189 | raise NotImplementedError 190 | 191 | def state_vector(self) -> NDArray[np.float64]: 192 | """Return the position and velocity joint states of the model""" 193 | return np.concatenate([self.data.qpos.flat, self.data.qvel.flat]) 194 | 195 | 196 | class MuJocoPyEnv(BaseMujocoEnv): 197 | def __init__( 198 | self, 199 | model_path: str, 200 | frame_skip: int, 201 | observation_space: Space, 202 | render_mode: Optional[str] = None, 203 | width: int = DEFAULT_SIZE, 204 | height: int = DEFAULT_SIZE, 205 | camera_id: Optional[int] = None, 206 | camera_name: Optional[str] = None, 207 | ): 208 | if MUJOCO_PY_IMPORT_ERROR is not None: 209 | raise error.DependencyNotInstalled( 210 | f"{MUJOCO_PY_IMPORT_ERROR}. " 211 | "(HINT: you need to install mujoco-py, and also perform the setup instructions " 212 | "here: https://github.com/openai/mujoco-py.)" 213 | ) 214 | 215 | logger.deprecation( 216 | "This version of the mujoco environments depends " 217 | "on the mujoco-py bindings, which are no longer maintained " 218 | "and may stop working. Please upgrade to the v4 versions of " 219 | "the environments (which depend on the mujoco python bindings instead), unless " 220 | "you are trying to precisely replicate previous works)." 221 | ) 222 | 223 | self.viewer = None 224 | self._viewers = {} 225 | 226 | super().__init__( 227 | model_path, 228 | frame_skip, 229 | observation_space, 230 | render_mode, 231 | width, 232 | height, 233 | camera_id, 234 | camera_name, 235 | ) 236 | 237 | def _initialize_simulation(self): 238 | model = mujoco_py.load_model_from_path(self.fullpath) 239 | self.sim = mujoco_py.MjSim(model) 240 | data = self.sim.data 241 | return model, data 242 | 243 | def _reset_simulation(self): 244 | self.sim.reset() 245 | 246 | def set_state(self, qpos, qvel): 247 | super().set_state(qpos, qvel) 248 | state = self.sim.get_state() 249 | state = mujoco_py.MjSimState(state.time, qpos, qvel, state.act, state.udd_state) 250 | self.sim.set_state(state) 251 | self.sim.forward() 252 | 253 | def get_body_com(self, body_name): 254 | return self.data.get_body_xpos(body_name) 255 | 256 | def _step_mujoco_simulation(self, ctrl, n_frames): 257 | self.sim.data.ctrl[:] = ctrl 258 | 259 | for _ in range(n_frames): 260 | self.sim.step() 261 | 262 | def render(self): 263 | if self.render_mode is None: 264 | assert self.spec is not None 265 | gym.logger.warn( 266 | "You are calling render method without specifying any render mode. " 267 | "You can specify the render_mode at initialization, " 268 | f'e.g. gym.make("{self.spec.id}", render_mode="rgb_array")' 269 | ) 270 | return 271 | 272 | width, height = self.width, self.height 273 | camera_name, camera_id = self.camera_name, self.camera_id 274 | if self.render_mode in {"rgb_array", "depth_array"}: 275 | if camera_id is not None and camera_name is not None: 276 | raise ValueError( 277 | "Both `camera_id` and `camera_name` cannot be" 278 | " specified at the same time." 279 | ) 280 | 281 | no_camera_specified = camera_name is None and camera_id is None 282 | if no_camera_specified: 283 | camera_name = "track" 284 | 285 | if camera_id is None and camera_name in self.model._camera_name2id: 286 | if camera_name in self.model._camera_name2id: 287 | camera_id = self.model.camera_name2id(camera_name) 288 | 289 | self._get_viewer(self.render_mode).render( 290 | width, height, camera_id=camera_id 291 | ) 292 | 293 | if self.render_mode == "rgb_array": 294 | data = self._get_viewer(self.render_mode).read_pixels( 295 | width, height, depth=False 296 | ) 297 | # original image is upside-down, so flip it 298 | return data[::-1, :, :] 299 | elif self.render_mode == "depth_array": 300 | self._get_viewer(self.render_mode).render(width, height) 301 | # Extract depth part of the read_pixels() tuple 302 | data = self._get_viewer(self.render_mode).read_pixels( 303 | width, height, depth=True 304 | )[1] 305 | # original image is upside-down, so flip it 306 | return data[::-1, :] 307 | elif self.render_mode == "human": 308 | self._get_viewer(self.render_mode).render() 309 | 310 | def _get_viewer( 311 | self, mode 312 | ) -> Union["mujoco_py.MjViewer", "mujoco_py.MjRenderContextOffscreen"]: 313 | self.viewer = self._viewers.get(mode) 314 | if self.viewer is None: 315 | if mode == "human": 316 | self.viewer = mujoco_py.MjViewer(self.sim) 317 | 318 | elif mode in {"rgb_array", "depth_array"}: 319 | self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1) 320 | else: 321 | raise AttributeError( 322 | f"Unknown mode: {mode}, expected modes: {self.metadata['render_modes']}" 323 | ) 324 | 325 | self.viewer_setup() 326 | self._viewers[mode] = self.viewer 327 | 328 | return self.viewer 329 | 330 | def close(self): 331 | if self.viewer is not None: 332 | self.viewer = None 333 | self._viewers = {} 334 | 335 | def viewer_setup(self): 336 | """ 337 | This method is called when the viewer is initialized. 338 | Optionally implement this method, if you need to tinker with camera position and so forth. 339 | """ 340 | raise NotImplementedError 341 | 342 | 343 | class MujocoEnv(BaseMujocoEnv): 344 | """Superclass for MuJoCo environments.""" 345 | 346 | def __init__( 347 | self, 348 | model_path, 349 | frame_skip, 350 | observation_space: Space, 351 | render_mode: Optional[str] = None, 352 | width: int = DEFAULT_SIZE, 353 | height: int = DEFAULT_SIZE, 354 | camera_id: Optional[int] = None, 355 | camera_name: Optional[str] = None, 356 | default_camera_config: Optional[Dict[str, Union[float, int]]] = None, 357 | ): 358 | if MUJOCO_IMPORT_ERROR is not None: 359 | raise error.DependencyNotInstalled( 360 | f"{MUJOCO_IMPORT_ERROR}. " 361 | "(HINT: you need to install mujoco, run `pip install gymnasium[mujoco]`.)" 362 | ) 363 | 364 | super().__init__( 365 | model_path, 366 | frame_skip, 367 | observation_space, 368 | render_mode, 369 | width, 370 | height, 371 | camera_id, 372 | camera_name, 373 | ) 374 | 375 | from homestri_ur5e_rl.envs.mujoco.mujoco_rendering import MujocoRenderer 376 | 377 | self.mujoco_renderer = MujocoRenderer( 378 | self.model, self.data, default_camera_config 379 | ) 380 | 381 | def _initialize_simulation( 382 | self, 383 | ) -> Tuple["mujoco._structs.MjModel", "mujoco._structs.MjData"]: 384 | model = mujoco.MjModel.from_xml_path(self.fullpath) 385 | # MjrContext will copy model.vis.global_.off* to con.off* 386 | model.vis.global_.offwidth = self.width 387 | model.vis.global_.offheight = self.height 388 | data = mujoco.MjData(model) 389 | return model, data 390 | 391 | def _reset_simulation(self): 392 | mujoco.mj_resetData(self.model, self.data) 393 | 394 | def set_state(self, qpos, qvel): 395 | super().set_state(qpos, qvel) 396 | self.data.qpos[:] = np.copy(qpos) 397 | self.data.qvel[:] = np.copy(qvel) 398 | if self.model.na == 0: 399 | self.data.act[:] = None 400 | mujoco.mj_forward(self.model, self.data) 401 | 402 | def _step_mujoco_simulation(self, ctrl, n_frames): 403 | self.data.ctrl[:] = ctrl 404 | 405 | mujoco.mj_step(self.model, self.data, nstep=n_frames) 406 | 407 | # As of MuJoCo 2.0, force-related quantities like cacc are not computed 408 | # unless there's a force sensor in the model. 409 | # See https://github.com/openai/gym/issues/1541 410 | mujoco.mj_rnePostConstraint(self.model, self.data) 411 | 412 | def render(self): 413 | return self.mujoco_renderer.render( 414 | self.render_mode, self.camera_id, self.camera_name 415 | ) 416 | 417 | def close(self): 418 | if self.mujoco_renderer is not None: 419 | self.mujoco_renderer.close() 420 | 421 | def get_body_com(self, body_name): 422 | return self.data.body(body_name).xpos -------------------------------------------------------------------------------- /homestri_ur5e_rl/envs/mujoco/mujoco_rendering.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import os 3 | import time 4 | from typing import Optional 5 | 6 | import glfw 7 | import imageio 8 | import mujoco 9 | import numpy as np 10 | 11 | 12 | def _import_egl(width, height): 13 | from mujoco.egl import GLContext 14 | 15 | return GLContext(width, height) 16 | 17 | 18 | def _import_glfw(width, height): 19 | from mujoco.glfw import GLContext 20 | 21 | return GLContext(width, height) 22 | 23 | 24 | def _import_osmesa(width, height): 25 | from mujoco.osmesa import GLContext 26 | 27 | return GLContext(width, height) 28 | 29 | 30 | _ALL_RENDERERS = collections.OrderedDict( 31 | [ 32 | ("glfw", _import_glfw), 33 | ("egl", _import_egl), 34 | ("osmesa", _import_osmesa), 35 | ] 36 | ) 37 | 38 | 39 | class BaseRender: 40 | def __init__( 41 | self, model: "mujoco.MjModel", data: "mujoco.MjData", width: int, height: int 42 | ): 43 | """Render context superclass for offscreen and window rendering.""" 44 | self.model = model 45 | self.data = data 46 | 47 | self._markers = [] 48 | self._overlays = {} 49 | 50 | self.viewport = mujoco.MjrRect(0, 0, width, height) 51 | 52 | # This goes to specific visualizer 53 | self.scn = mujoco.MjvScene(self.model, 1000) 54 | self.cam = mujoco.MjvCamera() 55 | self.vopt = mujoco.MjvOption() 56 | self.pert = mujoco.MjvPerturb() 57 | 58 | self.make_context_current() 59 | 60 | # Keep in Mujoco Context 61 | self.con = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150) 62 | 63 | self._set_mujoco_buffer() 64 | 65 | def _set_mujoco_buffer(self): 66 | raise NotImplementedError 67 | 68 | def make_context_current(self): 69 | raise NotImplementedError 70 | 71 | def add_overlay(self, gridpos: int, text1: str, text2: str): 72 | """Overlays text on the scene.""" 73 | if gridpos not in self._overlays: 74 | self._overlays[gridpos] = ["", ""] 75 | self._overlays[gridpos][0] += text1 + "\n" 76 | self._overlays[gridpos][1] += text2 + "\n" 77 | 78 | def add_marker(self, **marker_params): 79 | self._markers.append(marker_params) 80 | 81 | def _add_marker_to_scene(self, marker: dict): 82 | if self.scn.ngeom >= self.scn.maxgeom: 83 | raise RuntimeError("Ran out of geoms. maxgeom: %d" % self.scn.maxgeom) 84 | 85 | g = self.scn.geoms[self.scn.ngeom] 86 | # default values. 87 | g.dataid = -1 88 | g.objtype = mujoco.mjtObj.mjOBJ_UNKNOWN 89 | g.objid = -1 90 | g.category = mujoco.mjtCatBit.mjCAT_DECOR 91 | g.texid = -1 92 | g.texuniform = 0 93 | g.texrepeat[0] = 1 94 | g.texrepeat[1] = 1 95 | g.emission = 0 96 | g.specular = 0.5 97 | g.shininess = 0.5 98 | g.reflectance = 0 99 | g.type = mujoco.mjtGeom.mjGEOM_BOX 100 | g.size[:] = np.ones(3) * 0.1 101 | g.mat[:] = np.eye(3) 102 | g.rgba[:] = np.ones(4) 103 | 104 | for key, value in marker.items(): 105 | if isinstance(value, (int, float, mujoco._enums.mjtGeom)): 106 | setattr(g, key, value) 107 | elif isinstance(value, (tuple, list, np.ndarray)): 108 | attr = getattr(g, key) 109 | attr[:] = np.asarray(value).reshape(attr.shape) 110 | elif isinstance(value, str): 111 | assert key == "label", "Only label is a string in mjtGeom." 112 | if value is None: 113 | g.label[0] = 0 114 | else: 115 | g.label = value 116 | elif hasattr(g, key): 117 | raise ValueError( 118 | "mjtGeom has attr {} but type {} is invalid".format( 119 | key, type(value) 120 | ) 121 | ) 122 | else: 123 | raise ValueError("mjtGeom doesn't have field %s" % key) 124 | 125 | self.scn.ngeom += 1 126 | 127 | def close(self): 128 | """Override close in your rendering subclass to perform any necessary cleanup 129 | after env.close() is called. 130 | """ 131 | raise NotImplementedError 132 | 133 | 134 | class OffScreenViewer(BaseRender): 135 | """Offscreen rendering class with opengl context.""" 136 | 137 | def __init__(self, model: "mujoco.MjMujoco", data: "mujoco.MjData"): 138 | width = model.vis.global_.offwidth 139 | height = model.vis.global_.offheight 140 | 141 | # We must make GLContext before MjrContext 142 | self._get_opengl_backend(width, height) 143 | 144 | super().__init__(model, data, width, height) 145 | 146 | self._init_camera() 147 | 148 | def _init_camera(self): 149 | self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE 150 | self.cam.fixedcamid = -1 151 | for i in range(3): 152 | self.cam.lookat[i] = np.median(self.data.geom_xpos[:, i]) 153 | self.cam.distance = self.model.stat.extent 154 | 155 | def _get_opengl_backend(self, width: int, height: int): 156 | self.backend = os.environ.get("MUJOCO_GL") 157 | if self.backend is not None: 158 | try: 159 | self.opengl_context = _ALL_RENDERERS[self.backend](width, height) 160 | except KeyError as e: 161 | raise RuntimeError( 162 | "Environment variable {} must be one of {!r}: got {!r}.".format( 163 | "MUJOCO_GL", _ALL_RENDERERS.keys(), self.backend 164 | ) 165 | ) from e 166 | 167 | else: 168 | for name, _ in _ALL_RENDERERS.items(): 169 | try: 170 | self.opengl_context = _ALL_RENDERERS[name](width, height) 171 | self.backend = name 172 | break 173 | except: # noqa:E722 174 | pass 175 | if self.backend is None: 176 | raise RuntimeError( 177 | "No OpenGL backend could be imported. Attempting to create a " 178 | "rendering context will result in a RuntimeError." 179 | ) 180 | 181 | def _set_mujoco_buffer(self): 182 | mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, self.con) 183 | 184 | def make_context_current(self): 185 | self.opengl_context.make_current() 186 | 187 | def free(self): 188 | self.opengl_context.free() 189 | 190 | def __del__(self): 191 | self.free() 192 | 193 | def render( 194 | self, 195 | render_mode: Optional[str], 196 | camera_id: Optional[int] = None, 197 | segmentation: bool = False, 198 | ): 199 | if camera_id is not None: 200 | if camera_id == -1: 201 | self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE 202 | else: 203 | self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED 204 | self.cam.fixedcamid = camera_id 205 | 206 | mujoco.mjv_updateScene( 207 | self.model, 208 | self.data, 209 | self.vopt, 210 | self.pert, 211 | self.cam, 212 | mujoco.mjtCatBit.mjCAT_ALL, 213 | self.scn, 214 | ) 215 | 216 | if segmentation: 217 | self.scn.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = 1 218 | self.scn.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = 1 219 | 220 | for marker_params in self._markers: 221 | self._add_marker_to_scene(marker_params) 222 | 223 | mujoco.mjr_render(self.viewport, self.scn, self.con) 224 | 225 | for gridpos, (text1, text2) in self._overlays.items(): 226 | mujoco.mjr_overlay( 227 | mujoco.mjtFontScale.mjFONTSCALE_150, 228 | gridpos, 229 | self.viewport, 230 | text1.encode(), 231 | text2.encode(), 232 | self.con, 233 | ) 234 | 235 | if segmentation: 236 | self.scn.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = 0 237 | self.scn.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = 0 238 | 239 | rgb_arr = np.zeros( 240 | 3 * self.viewport.width * self.viewport.height, dtype=np.uint8 241 | ) 242 | depth_arr = np.zeros( 243 | self.viewport.width * self.viewport.height, dtype=np.float32 244 | ) 245 | 246 | mujoco.mjr_readPixels(rgb_arr, depth_arr, self.viewport, self.con) 247 | 248 | if render_mode == "depth_array": 249 | depth_img = depth_arr.reshape(self.viewport.height, self.viewport.width) 250 | # original image is upside-down, so flip it 251 | return depth_img[::-1, :] 252 | else: 253 | rgb_img = rgb_arr.reshape(self.viewport.height, self.viewport.width, 3) 254 | 255 | if segmentation: 256 | seg_img = ( 257 | rgb_img[:, :, 0] 258 | + rgb_img[:, :, 1] * (2**8) 259 | + rgb_img[:, :, 2] * (2**16) 260 | ) 261 | seg_img[seg_img >= (self.scn.ngeom + 1)] = 0 262 | seg_ids = np.full( 263 | (self.scn.ngeom + 1, 2), fill_value=-1, dtype=np.int32 264 | ) 265 | 266 | for i in range(self.scn.ngeom): 267 | geom = self.scn.geoms[i] 268 | if geom.segid != -1: 269 | seg_ids[geom.segid + 1, 0] = geom.objtype 270 | seg_ids[geom.segid + 1, 1] = geom.objid 271 | rgb_img = seg_ids[seg_img] 272 | 273 | # original image is upside-down, so flip i 274 | return rgb_img[::-1, :, :] 275 | 276 | def close(self): 277 | self.free() 278 | glfw.terminate() 279 | 280 | 281 | class WindowViewer(BaseRender): 282 | """Class for window rendering in all MuJoCo environments.""" 283 | 284 | def __init__(self, model: "mujoco.MjModel", data: "mujoco.MjData"): 285 | glfw.init() 286 | 287 | self._button_left_pressed = False 288 | self._button_right_pressed = False 289 | self._last_mouse_x = 0 290 | self._last_mouse_y = 0 291 | self._paused = False 292 | self._transparent = False 293 | self._contacts = False 294 | self._render_every_frame = True 295 | self._image_idx = 0 296 | self._image_path = "/tmp/frame_%07d.png" 297 | self._time_per_render = 1 / 60.0 298 | self._run_speed = 1.0 299 | self._loop_count = 0 300 | self._advance_by_one_step = False 301 | self._hide_menu = False 302 | 303 | width, height = glfw.get_video_mode(glfw.get_primary_monitor()).size 304 | glfw.window_hint(glfw.VISIBLE, 1) 305 | self.window = glfw.create_window(width // 2, height // 2, "mujoco", None, None) 306 | 307 | self.width, self.height = glfw.get_framebuffer_size(self.window) 308 | window_width, _ = glfw.get_window_size(self.window) 309 | self._scale = self.width * 1.0 / window_width 310 | 311 | # set callbacks 312 | glfw.set_cursor_pos_callback(self.window, self._cursor_pos_callback) 313 | glfw.set_mouse_button_callback(self.window, self._mouse_button_callback) 314 | glfw.set_scroll_callback(self.window, self._scroll_callback) 315 | glfw.set_key_callback(self.window, self._key_callback) 316 | 317 | super().__init__(model, data, width, height) 318 | glfw.swap_interval(1) 319 | 320 | def _set_mujoco_buffer(self): 321 | mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_WINDOW, self.con) 322 | 323 | def make_context_current(self): 324 | glfw.make_context_current(self.window) 325 | 326 | def free(self): 327 | if self.window: 328 | if glfw.get_current_context() == self.window: 329 | glfw.make_context_current(None) 330 | glfw.destroy_window(self.window) 331 | self.window = None 332 | 333 | def __del__(self): 334 | """Eliminate all of the OpenGL glfw contexts and windows""" 335 | self.free() 336 | 337 | def render(self): 338 | """ 339 | Renders the environment geometries in the OpenGL glfw window: 340 | 1. Create the overlay for the left side panel menu. 341 | 2. Update the geometries used for rendering based on the current state of the model - `mujoco.mjv_updateScene()`. 342 | 3. Add markers to scene, these are additional geometries to include in the model, i.e arrows, https://mujoco.readthedocs.io/en/latest/APIreference.html?highlight=arrow#mjtgeom. 343 | These markers are added with the `add_marker()` method before rendering. 344 | 4. Render the 3D scene to the window context - `mujoco.mjr_render()`. 345 | 5. Render overlays in the window context - `mujoco.mjr_overlay()`. 346 | 6. Swap front and back buffer, https://www.glfw.org/docs/3.3/quick.html. 347 | 7. Poll events like mouse clicks or keyboard input. 348 | """ 349 | 350 | # mjv_updateScene, mjr_render, mjr_overlay 351 | def update(): 352 | # fill overlay items 353 | self._create_overlay() 354 | 355 | render_start = time.time() 356 | if self.window is None: 357 | return 358 | elif glfw.window_should_close(self.window): 359 | glfw.destroy_window(self.window) 360 | glfw.terminate() 361 | self.viewport.width, self.viewport.height = glfw.get_framebuffer_size( 362 | self.window 363 | ) 364 | # update scene 365 | mujoco.mjv_updateScene( 366 | self.model, 367 | self.data, 368 | self.vopt, 369 | mujoco.MjvPerturb(), 370 | self.cam, 371 | mujoco.mjtCatBit.mjCAT_ALL.value, 372 | self.scn, 373 | ) 374 | 375 | # marker items 376 | for marker in self._markers: 377 | self._add_marker_to_scene(marker) 378 | 379 | # render 380 | mujoco.mjr_render(self.viewport, self.scn, self.con) 381 | 382 | # overlay items 383 | if not self._hide_menu: 384 | for gridpos, [t1, t2] in self._overlays.items(): 385 | mujoco.mjr_overlay( 386 | mujoco.mjtFontScale.mjFONTSCALE_150, 387 | gridpos, 388 | self.viewport, 389 | t1, 390 | t2, 391 | self.con, 392 | ) 393 | 394 | glfw.swap_buffers(self.window) 395 | glfw.poll_events() 396 | self._time_per_render = 0.9 * self._time_per_render + 0.1 * ( 397 | time.time() - render_start 398 | ) 399 | 400 | if self._paused: 401 | while self._paused: 402 | update() 403 | if self._advance_by_one_step: 404 | self._advance_by_one_step = False 405 | break 406 | else: 407 | self._loop_count += self.model.opt.timestep / ( 408 | self._time_per_render * self._run_speed 409 | ) 410 | if self._render_every_frame: 411 | self._loop_count = 1 412 | while self._loop_count > 0: 413 | update() 414 | self._loop_count -= 1 415 | 416 | # clear overlay 417 | self._overlays.clear() 418 | # clear markers 419 | self._markers.clear() 420 | 421 | def close(self): 422 | self.free() 423 | glfw.terminate() 424 | 425 | def _key_callback(self, window, key: int, scancode, action: int, mods): 426 | if action != glfw.RELEASE: 427 | return 428 | # Switch cameras 429 | elif key == glfw.KEY_X: 430 | self.cam.fixedcamid += 1 431 | self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED 432 | if self.cam.fixedcamid >= self.model.ncam: 433 | self.cam.fixedcamid = -1 434 | self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE 435 | # Pause simulation 436 | elif key == glfw.KEY_B and self._paused is not None: 437 | self._paused = not self._paused 438 | # Advances simulation by one step. 439 | elif key == glfw.KEY_N and self._paused is not None: 440 | self._advance_by_one_step = True 441 | self._paused = True 442 | # Slows down simulation 443 | elif key == glfw.KEY_MINUS: 444 | self._run_speed /= 2.0 445 | # Speeds up simulation 446 | elif key == glfw.KEY_EQUAL: 447 | self._run_speed *= 2.0 448 | # Turn off / turn on rendering every frame. 449 | elif key == glfw.KEY_Z: 450 | self._render_every_frame = not self._render_every_frame 451 | # Capture screenshot 452 | elif key == glfw.KEY_PERIOD: 453 | img = np.zeros( 454 | ( 455 | glfw.get_framebuffer_size(self.window)[1], 456 | glfw.get_framebuffer_size(self.window)[0], 457 | 3, 458 | ), 459 | dtype=np.uint8, 460 | ) 461 | mujoco.mjr_readPixels(img, None, self.viewport, self.con) 462 | imageio.imwrite(self._image_path % self._image_idx, np.flipud(img)) 463 | self._image_idx += 1 464 | # Display contact forces 465 | elif key == glfw.KEY_C: 466 | self._contacts = not self._contacts 467 | self.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = self._contacts 468 | self.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = self._contacts 469 | # Display coordinate frames 470 | elif key == glfw.KEY_M: 471 | self.vopt.frame = 1 - self.vopt.frame 472 | # Hide overlay menu 473 | elif key == glfw.KEY_COMMA: 474 | self._hide_menu = not self._hide_menu 475 | # Make transparent 476 | elif key == glfw.KEY_V: 477 | self._transparent = not self._transparent 478 | if self._transparent: 479 | self.model.geom_rgba[:, 3] /= 5.0 480 | else: 481 | self.model.geom_rgba[:, 3] *= 5.0 482 | # Geom group visibility 483 | elif key in (glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4): 484 | self.vopt.geomgroup[key - glfw.KEY_0] ^= 1 485 | # Quit 486 | if key == glfw.KEY_ESCAPE: 487 | print("Pressed ESC") 488 | print("Quitting.") 489 | glfw.destroy_window(self.window) 490 | glfw.terminate() 491 | 492 | def _cursor_pos_callback( 493 | self, window: "glfw.LP__GLFWwindow", xpos: float, ypos: float 494 | ): 495 | if not (self._button_left_pressed or self._button_right_pressed): 496 | return 497 | 498 | mod_shift = ( 499 | glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS 500 | or glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS 501 | ) 502 | if self._button_right_pressed: 503 | action = ( 504 | mujoco.mjtMouse.mjMOUSE_MOVE_H 505 | if mod_shift 506 | else mujoco.mjtMouse.mjMOUSE_MOVE_V 507 | ) 508 | elif self._button_left_pressed: 509 | action = ( 510 | mujoco.mjtMouse.mjMOUSE_ROTATE_H 511 | if mod_shift 512 | else mujoco.mjtMouse.mjMOUSE_ROTATE_V 513 | ) 514 | else: 515 | action = mujoco.mjtMouse.mjMOUSE_ZOOM 516 | 517 | dx = int(self._scale * xpos) - self._last_mouse_x 518 | dy = int(self._scale * ypos) - self._last_mouse_y 519 | width, height = glfw.get_framebuffer_size(window) 520 | 521 | mujoco.mjv_moveCamera( 522 | self.model, action, dx / height, dy / height, self.scn, self.cam 523 | ) 524 | 525 | self._last_mouse_x = int(self._scale * xpos) 526 | self._last_mouse_y = int(self._scale * ypos) 527 | 528 | def _mouse_button_callback(self, window: "glfw.LP__GLFWwindow", button, act, mods): 529 | self._button_left_pressed = ( 530 | glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_LEFT) == glfw.PRESS 531 | ) 532 | self._button_right_pressed = ( 533 | glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_RIGHT) == glfw.PRESS 534 | ) 535 | 536 | x, y = glfw.get_cursor_pos(window) 537 | self._last_mouse_x = int(self._scale * x) 538 | self._last_mouse_y = int(self._scale * y) 539 | 540 | def _scroll_callback(self, window, x_offset, y_offset: float): 541 | mujoco.mjv_moveCamera( 542 | self.model, 543 | mujoco.mjtMouse.mjMOUSE_ZOOM, 544 | 0, 545 | -0.05 * y_offset, 546 | self.scn, 547 | self.cam, 548 | ) 549 | 550 | def _create_overlay(self): 551 | topleft = mujoco.mjtGridPos.mjGRID_TOPLEFT 552 | bottomleft = mujoco.mjtGridPos.mjGRID_BOTTOMLEFT 553 | 554 | if self._render_every_frame: 555 | self.add_overlay(topleft, "", "") 556 | else: 557 | self.add_overlay( 558 | topleft, 559 | "Run speed = %.3f x real time" % self._run_speed, 560 | "Slower [-], Faster [+]", 561 | ) 562 | self.add_overlay( 563 | topleft, "Render every frame [Z]", "On" if self._render_every_frame else "Off" 564 | ) 565 | self.add_overlay( 566 | topleft, 567 | "Switch camera (#cams = %d)" % (self.model.ncam + 1), 568 | "[X] (camera ID = %d)" % self.cam.fixedcamid, 569 | ) 570 | self.add_overlay(topleft, "Contact forces [C]", "On" if self._contacts else "Off") 571 | self.add_overlay(topleft, "Transparent [V]", "On" if self._transparent else "Off") 572 | if self._paused is not None: 573 | if not self._paused: 574 | self.add_overlay(topleft, "Stop", "[B]") 575 | else: 576 | self.add_overlay(topleft, "Start", "[B]") 577 | self.add_overlay( 578 | topleft, "Advance simulation by one step", "[N]" 579 | ) 580 | self.add_overlay( 581 | topleft, "Reference frames [M]", "On" if self.vopt.frame == 1 else "Off" 582 | ) 583 | self.add_overlay(topleft, "Hide Menu [<]", "") 584 | if self._image_idx > 0: 585 | fname = self._image_path % (self._image_idx - 1) 586 | self.add_overlay(topleft, "Capture frame [>]", "Saved as %s" % fname) 587 | else: 588 | self.add_overlay(topleft, "Capture frame [>]", "") 589 | self.add_overlay(topleft, "Toggle geomgroup visibility", "0-4") 590 | 591 | self.add_overlay(bottomleft, "FPS", "%d%s" % (1 / self._time_per_render, "")) 592 | self.add_overlay( 593 | bottomleft, "Solver iterations", str(self.data.solver_iter + 1) 594 | ) 595 | self.add_overlay( 596 | bottomleft, "Step", str(round(self.data.time / self.model.opt.timestep)) 597 | ) 598 | self.add_overlay(bottomleft, "timestep", "%.5f" % self.model.opt.timestep) 599 | 600 | 601 | class MujocoRenderer: 602 | """This is the MuJoCo renderer manager class for every MuJoCo environment. 603 | 604 | The class has two main public methods available: 605 | - :meth:`render` - Renders the environment in three possible modes: "human", "rgb_array", or "depth_array" 606 | - :meth:`close` - Closes all contexts initialized with the renderer 607 | 608 | """ 609 | 610 | def __init__( 611 | self, 612 | model: "mujoco.MjModel", 613 | data: "mujoco.MjData", 614 | default_cam_config: Optional[dict] = None, 615 | ): 616 | """A wrapper for clipping continuous actions within the valid bound. 617 | 618 | Args: 619 | model: MjModel data structure of the MuJoCo simulation 620 | data: MjData data structure of the MuJoCo simulation 621 | default_cam_config: dictionary with attribute values of the viewer's default camera, https://mujoco.readthedocs.io/en/latest/XMLreference.html?highlight=camera#visual-global 622 | """ 623 | self.model = model 624 | self.data = data 625 | self._viewers = {} 626 | self.viewer = None 627 | self.default_cam_config = default_cam_config 628 | 629 | def render( 630 | self, 631 | render_mode: Optional[str], 632 | camera_id: Optional[int] = None, 633 | camera_name: Optional[str] = None, 634 | ): 635 | """Renders a frame of the simulation in a specific format and camera view. 636 | 637 | Args: 638 | render_mode: The format to render the frame, it can be: "human", "rgb_array", or "depth_array" 639 | camera_id: The integer camera id from which to render the frame in the MuJoCo simulation 640 | camera_name: The string name of the camera from which to render the frame in the MuJoCo simulation. This argument should not be passed if using cameara_id instead and vice versa 641 | 642 | Returns: 643 | If render_mode is "rgb_array" or "depth_arra" it returns a numpy array in the specified format. "human" render mode does not return anything. 644 | """ 645 | 646 | viewer = self._get_viewer(render_mode=render_mode) 647 | 648 | if render_mode in { 649 | "rgb_array", 650 | "depth_array", 651 | }: 652 | if camera_id is not None and camera_name is not None: 653 | raise ValueError( 654 | "Both `camera_id` and `camera_name` cannot be" 655 | " specified at the same time." 656 | ) 657 | 658 | no_camera_specified = camera_name is None and camera_id is None 659 | if no_camera_specified: 660 | camera_name = "track" 661 | 662 | if camera_id is None: 663 | camera_id = mujoco.mj_name2id( 664 | self.model, 665 | mujoco.mjtObj.mjOBJ_CAMERA, 666 | camera_name, 667 | ) 668 | 669 | img = viewer.render(render_mode=render_mode, camera_id=camera_id) 670 | return img 671 | 672 | elif render_mode == "human": 673 | return viewer.render() 674 | 675 | def _get_viewer(self, render_mode: str): 676 | """Initializes and returns a viewer class depending on the render_mode 677 | - `WindowViewer` class for "human" render mode 678 | - `OffScreenViewer` class for "rgb_array" or "depth_array" render mode 679 | """ 680 | self.viewer = self._viewers.get(render_mode) 681 | if self.viewer is None: 682 | if render_mode == "human": 683 | self.viewer = WindowViewer(self.model, self.data) 684 | 685 | elif render_mode in {"rgb_array", "depth_array"}: 686 | self.viewer = OffScreenViewer(self.model, self.data) 687 | else: 688 | raise AttributeError( 689 | f"Unexpected mode: {render_mode}, expected modes: human, rgb_array, or depth_array" 690 | ) 691 | # Add default camera parameters 692 | self._set_cam_config() 693 | self._viewers[render_mode] = self.viewer 694 | 695 | if len(self._viewers.keys()) > 1: 696 | # Only one context can be current at a time 697 | self.viewer.make_context_current() 698 | 699 | return self.viewer 700 | 701 | def _set_cam_config(self): 702 | """Set the default camera parameters""" 703 | assert self.viewer is not None 704 | if self.default_cam_config is not None: 705 | for key, value in self.default_cam_config.items(): 706 | if isinstance(value, np.ndarray): 707 | getattr(self.viewer.cam, key)[:] = value 708 | else: 709 | setattr(self.viewer.cam, key, value) 710 | 711 | def close(self): 712 | """Close the OpenGL rendering contexts of all viewer modes""" 713 | for _, viewer in self._viewers.items(): 714 | viewer.close() -------------------------------------------------------------------------------- /homestri_ur5e_rl/input_devices/keyboard_input.py: -------------------------------------------------------------------------------- 1 | from pynput import keyboard 2 | import numpy as np 3 | 4 | class KeyboardInput: 5 | def __init__(self): 6 | self.linear_velocity = np.zeros(3) 7 | self.angular_velocity = np.zeros(3) 8 | self.gripper_position = 0 9 | self.listener = keyboard.Listener(on_press=self.on_press, on_release=self.on_release) 10 | self.listener.start() 11 | 12 | def on_press(self, key): 13 | try: 14 | key_char = key.char.lower() 15 | if key_char in ['w', 'a', 's', 'd', 'q', 'e']: 16 | linear_mapping = { 17 | 'w': [0.05, 0, 0], 18 | 's': [-0.05, 0, 0], 19 | 'a': [0, 0.05, 0], 20 | 'd': [0, -0.05, 0], 21 | 'q': [0, 0, 0.05], 22 | 'e': [0, 0, -0.05], 23 | } 24 | self.linear_velocity = np.array(linear_mapping[key_char]) 25 | elif key_char in ['i', 'j', 'k', 'l', 'u', 'o']: 26 | angular_mapping = { 27 | 'i': [0.25, 0, 0], 28 | 'k': [-0.25, 0, 0], 29 | 'j': [0, 0.25, 0], 30 | 'l': [0, -0.25, 0], 31 | 'u': [0, 0, 0.25], 32 | 'o': [0, 0, -0.25], 33 | } 34 | self.angular_velocity = np.array(angular_mapping[key_char]) 35 | 36 | except AttributeError: 37 | if key == keyboard.Key.space: 38 | # Toggle gripper position between 0 and 255 39 | self.gripper_position = 255 if self.gripper_position == 0 else 0 40 | 41 | def on_release(self, key): 42 | try: 43 | key_char = key.char.lower() 44 | if key_char in ['w', 'a', 's', 'd', 'q', 'e']: 45 | self.linear_velocity = np.zeros(3) 46 | elif key_char in ['i', 'j', 'k', 'l', 'u', 'o']: 47 | self.angular_velocity = np.zeros(3) 48 | except AttributeError: 49 | pass 50 | 51 | def get_action(self): 52 | action = np.concatenate((self.linear_velocity, self.angular_velocity, np.array([self.gripper_position]))) 53 | return action -------------------------------------------------------------------------------- /homestri_ur5e_rl/utils/controller_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from homestri_ur5e_rl.utils.transform_utils import ( 3 | quat_multiply, 4 | quat_conjugate, 5 | quat2mat, 6 | ) 7 | 8 | EPS = np.finfo(float).eps * 4.0 9 | 10 | def task_space_inertia_matrix(M, J, threshold=1e-3): 11 | """Generate the task-space inertia matrix 12 | 13 | Parameters 14 | ---------- 15 | M: np.array 16 | the generalized coordinates inertia matrix 17 | J: np.array 18 | the task space Jacobian 19 | threshold: scalar, optional (Default: 1e-3) 20 | singular value threshold, if the detminant of Mx_inv is less than 21 | this value then Mx is calculated using the pseudo-inverse function 22 | and all singular values < threshold * .1 are set = 0 23 | """ 24 | 25 | # calculate the inertia matrix in task space 26 | M_inv = np.linalg.inv(M) 27 | Mx_inv = np.dot(J, np.dot(M_inv, J.T)) 28 | if abs(np.linalg.det(Mx_inv)) >= threshold: 29 | # do the linalg inverse if matrix is non-singular 30 | # because it's faster and more accurate 31 | Mx = np.linalg.inv(Mx_inv) 32 | else: 33 | # using the rcond to set singular values < thresh to 0 34 | # singular values < (rcond * max(singular_values)) set to 0 35 | Mx = np.linalg.pinv(Mx_inv, rcond=threshold * 0.1) 36 | 37 | return Mx, M_inv 38 | 39 | 40 | 41 | 42 | def get_rot_angle(rotation_matrix, eps=EPS): 43 | data = rotation_matrix.reshape(9) 44 | ca = (data[0]+data[4]+data[8]-1)/2.0 45 | t = eps * eps / 2.0 46 | 47 | if ca > 1 - t: 48 | # Undefined case, choose the Z-axis and angle 0 49 | axis = np.array([0, 0, 1]) 50 | return 0, axis 51 | 52 | if ca < -1 + t: 53 | # The case of angles consisting of multiples of pi: 54 | # two solutions, choose a positive Z-component of the axis 55 | x = np.sqrt((data[0] + 1.0) / 2) 56 | y = np.sqrt((data[4] + 1.0) / 2) 57 | z = np.sqrt((data[8] + 1.0) / 2) 58 | 59 | if data[2] < 0: 60 | x = -x 61 | if data[7] < 0: 62 | y = -y 63 | if x * y * data[1] < 0: 64 | x = -x 65 | 66 | # z always >= 0 67 | # if z equals zero 68 | axis = np.array([x, y, z]) 69 | return np.pi, axis 70 | 71 | axis_x = data[7] - data[5] 72 | axis_y = data[2] - data[6] 73 | axis_z = data[3] - data[1] 74 | mod_axis = np.sqrt(axis_x * axis_x + axis_y * axis_y + axis_z * axis_z) 75 | 76 | axis = np.array([axis_x / mod_axis, axis_y / mod_axis, axis_z / mod_axis]) 77 | angle = np.arctan2(mod_axis / 2, ca) 78 | 79 | return angle, axis 80 | 81 | 82 | 83 | 84 | def pose_error(target_pose, ee_pose) -> np.ndarray: 85 | """ 86 | Calculate the rotational error (orientation difference) between the target and current orientation. 87 | 88 | Parameters: 89 | target_ori_mat (numpy.ndarray): The target orientation matrix. 90 | current_ori_mat (numpy.ndarray): The current orientation matrix. 91 | 92 | Returns: 93 | numpy.ndarray: The rotational error in axis-angle representation. 94 | """ 95 | target_pos = target_pose[:3] 96 | target_quat = target_pose[3:] 97 | ee_pos = ee_pose[:3] 98 | ee_quat = ee_pose[3:] 99 | 100 | err_pos = target_pos - ee_pos 101 | 102 | err_quat = quat_multiply(target_quat, quat_conjugate(ee_quat)) 103 | angle, axis = get_rot_angle(quat2mat(err_quat)) 104 | 105 | return np.concatenate([err_pos, angle * axis]) 106 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/utils/mujoco_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Tuple, Union 2 | 3 | import numpy as np 4 | from gymnasium import error 5 | 6 | try: 7 | import mujoco 8 | from mujoco import MjData, MjModel, mjtObj 9 | except ImportError as e: 10 | raise error.DependencyNotInstalled(f"{e}. (HINT: you need to install mujoco") 11 | 12 | MJ_OBJ_TYPES = [ 13 | "mjOBJ_BODY", 14 | "mjOBJ_JOINT", 15 | "mjOBJ_GEOM", 16 | "mjOBJ_SITE", 17 | "mjOBJ_CAMERA", 18 | "mjOBJ_ACTUATOR", 19 | "mjOBJ_SENSOR", 20 | ] 21 | 22 | def get_site_jac(model, data, site_id): 23 | """Return the Jacobian' translational component of the end-effector of 24 | the corresponding site id. 25 | """ 26 | jacp = np.zeros((3, model.nv)) 27 | jacr = np.zeros((3, model.nv)) 28 | mujoco.mj_jacSite(model, data, jacp, jacr, site_id) 29 | jac = np.vstack([jacp, jacr]) 30 | 31 | return jac 32 | 33 | def get_fullM(model, data): 34 | M = np.zeros((model.nv, model.nv)) 35 | mujoco.mj_fullM(model, M, data.qM) 36 | return M 37 | 38 | def extract_mj_names( 39 | model: MjModel, obj_type: mjtObj 40 | ) -> Tuple[Union[Tuple[str, ...], Tuple[()]], Dict[str, int], Dict[int, str]]: 41 | 42 | if obj_type == mujoco.mjtObj.mjOBJ_BODY: 43 | name_addr = model.name_bodyadr 44 | n_obj = model.nbody 45 | 46 | elif obj_type == mujoco.mjtObj.mjOBJ_JOINT: 47 | name_addr = model.name_jntadr 48 | n_obj = model.njnt 49 | 50 | elif obj_type == mujoco.mjtObj.mjOBJ_GEOM: 51 | name_addr = model.name_geomadr 52 | n_obj = model.ngeom 53 | 54 | elif obj_type == mujoco.mjtObj.mjOBJ_SITE: 55 | name_addr = model.name_siteadr 56 | n_obj = model.nsite 57 | 58 | elif obj_type == mujoco.mjtObj.mjOBJ_LIGHT: 59 | name_addr = model.name_lightadr 60 | n_obj = model.nlight 61 | 62 | elif obj_type == mujoco.mjtObj.mjOBJ_CAMERA: 63 | name_addr = model.name_camadr 64 | n_obj = model.ncam 65 | 66 | elif obj_type == mujoco.mjtObj.mjOBJ_ACTUATOR: 67 | name_addr = model.name_actuatoradr 68 | n_obj = model.nu 69 | 70 | elif obj_type == mujoco.mjtObj.mjOBJ_SENSOR: 71 | name_addr = model.name_sensoradr 72 | n_obj = model.nsensor 73 | 74 | elif obj_type == mujoco.mjtObj.mjOBJ_TENDON: 75 | name_addr = model.name_tendonadr 76 | n_obj = model.ntendon 77 | 78 | elif obj_type == mujoco.mjtObj.mjOBJ_MESH: 79 | name_addr = model.name_meshadr 80 | n_obj = model.nmesh 81 | else: 82 | raise ValueError( 83 | "`{}` was passed as the MuJoCo model object type. The MuJoCo model object type can only be of the following mjtObj enum types: {}.".format( 84 | obj_type, MJ_OBJ_TYPES 85 | ) 86 | ) 87 | 88 | id2name = {i: None for i in range(n_obj)} 89 | name2id = {} 90 | for addr in name_addr: 91 | name = model.names[addr:].split(b"\x00")[0].decode() 92 | if name: 93 | obj_id = mujoco.mj_name2id(model, obj_type, name) 94 | assert 0 <= obj_id < n_obj and id2name[obj_id] is None 95 | name2id[name] = obj_id 96 | id2name[obj_id] = name 97 | 98 | return tuple(id2name[id] for id in sorted(name2id.values())), name2id, id2name 99 | 100 | 101 | class MujocoModelNames: 102 | """Access mjtObj object names and ids of the current MuJoCo model. 103 | 104 | This class supports access to the names and ids of the following mjObj types: 105 | mjOBJ_BODY 106 | mjOBJ_JOINT 107 | mjOBJ_GEOM 108 | mjOBJ_SITE 109 | mjOBJ_CAMERA 110 | mjOBJ_ACTUATOR 111 | mjOBJ_SENSOR 112 | 113 | The properties provided for each ``mjObj`` are: 114 | ``mjObj``_names: list of the mjObj names in the model of type mjOBJ_FOO. 115 | ``mjObj``_name2id: dictionary with name of the mjObj as keys and id of the mjObj as values. 116 | ``mjObj``_id2name: dictionary with id of the mjObj as keys and name of the mjObj as values. 117 | """ 118 | 119 | def __init__(self, model: MjModel): 120 | """Access mjtObj object names and ids of the current MuJoCo model. 121 | 122 | Args: 123 | model: mjModel of the MuJoCo environment. 124 | """ 125 | ( 126 | self._body_names, 127 | self._body_name2id, 128 | self._body_id2name, 129 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_BODY) 130 | ( 131 | self._joint_names, 132 | self._joint_name2id, 133 | self._joint_id2name, 134 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_JOINT) 135 | ( 136 | self._geom_names, 137 | self._geom_name2id, 138 | self._geom_id2name, 139 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_GEOM) 140 | ( 141 | self._site_names, 142 | self._site_name2id, 143 | self._site_id2name, 144 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_SITE) 145 | ( 146 | self._camera_names, 147 | self._camera_name2id, 148 | self._camera_id2name, 149 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_CAMERA) 150 | ( 151 | self._actuator_names, 152 | self._actuator_name2id, 153 | self._actuator_id2name, 154 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_ACTUATOR) 155 | ( 156 | self._sensor_names, 157 | self._sensor_name2id, 158 | self._sensor_id2name, 159 | ) = extract_mj_names(model, mujoco.mjtObj.mjOBJ_SENSOR) 160 | 161 | @property 162 | def body_names(self): 163 | return self._body_names 164 | 165 | @property 166 | def body_name2id(self): 167 | return self._body_name2id 168 | 169 | @property 170 | def body_id2name(self): 171 | return self._body_id2name 172 | 173 | @property 174 | def joint_names(self): 175 | return self._joint_names 176 | 177 | @property 178 | def joint_name2id(self): 179 | return self._joint_name2id 180 | 181 | @property 182 | def joint_id2name(self): 183 | return self._joint_id2name 184 | 185 | @property 186 | def geom_names(self): 187 | return self._geom_names 188 | 189 | @property 190 | def geom_name2id(self): 191 | return self._geom_name2id 192 | 193 | @property 194 | def geom_id2name(self): 195 | return self._geom_id2name 196 | 197 | @property 198 | def site_names(self): 199 | return self._site_names 200 | 201 | @property 202 | def site_name2id(self): 203 | return self._site_name2id 204 | 205 | @property 206 | def site_id2name(self): 207 | return self._site_id2name 208 | 209 | @property 210 | def camera_names(self): 211 | return self._camera_names 212 | 213 | @property 214 | def camera_name2id(self): 215 | return self._camera_name2id 216 | 217 | @property 218 | def camera_id2name(self): 219 | return self._camera_id2name 220 | 221 | @property 222 | def actuator_names(self): 223 | return self._actuator_names 224 | 225 | @property 226 | def actuator_name2id(self): 227 | return self._actuator_name2id 228 | 229 | @property 230 | def actuator_id2name(self): 231 | return self._actuator_id2name 232 | 233 | @property 234 | def sensor_names(self): 235 | return self._sensor_names 236 | 237 | @property 238 | def sensor_name2id(self): 239 | return self._sensor_name2id 240 | 241 | @property 242 | def sensor_id2name(self): 243 | return self._sensor_id2name 244 | -------------------------------------------------------------------------------- /homestri_ur5e_rl/utils/pid_controller_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class PIDController: 4 | def __init__(self, kp, ki, kd): 5 | self.kp = kp # Proportional gain 6 | self.ki = ki # Integral gain 7 | self.kd = kd # Derivative gain 8 | self.prev_error = 0.0 9 | self.integral_accumulator = 0.0 10 | 11 | def __call__(self, error, dt): 12 | # Calculate the derivative of the error 13 | error_derivative = (error - self.prev_error) 14 | 15 | # Accumulate the integral component 16 | self.integral_accumulator += error 17 | 18 | # Calculate the control output 19 | control_output = self.kp * error + self.ki * self.integral_accumulator + self.kd * error_derivative 20 | 21 | # Store the current error for the next iteration 22 | self.prev_error = error 23 | 24 | return control_output 25 | 26 | 27 | class SpatialPIDController: 28 | def __init__(self, kp_values=None, ki_values=None, kd_values=None): 29 | self.m_cmd = np.zeros(6) 30 | if kp_values is None: 31 | kp_values = [0] * 6 32 | if ki_values is None: 33 | ki_values = [0] * 6 34 | if kd_values is None: 35 | kd_values = [0] * 6 36 | self.m_pid_controllers = [PIDController(kp_values[i], ki_values[i], kd_values[i]) for i in range(6)] 37 | 38 | def set_pid_gains(self, kp_values, ki_values, kd_values): 39 | # Set the P, I, and D gains for each dimension 40 | for i in range(6): 41 | self.m_pid_controllers[i].kp = kp_values[i] 42 | self.m_pid_controllers[i].ki = ki_values[i] 43 | self.m_pid_controllers[i].kd = kd_values[i] 44 | 45 | def __call__(self, error, period): 46 | # Call operator for one control cycle 47 | controlled_output = np.zeros(6) 48 | 49 | for i in range(6): 50 | # Apply each PID controller for each element in the error vector 51 | controlled_output[i] = self.m_pid_controllers[i](error[i], period) 52 | 53 | return controlled_output -------------------------------------------------------------------------------- /homestri_ur5e_rl/utils/solver_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class JacobianTransposeSolver(): 4 | def __init__(self, number_joints): 5 | self.last_positions = np.zeros(number_joints) 6 | 7 | def sync_jnt_pos(self, jnt_pos): 8 | self.last_positions = jnt_pos 9 | 10 | def compute_jnt_ctrl(self, jacobian, period, net_force): 11 | # Compute joint jacobian 12 | jacobian_transpose = np.transpose(jacobian) 13 | 14 | # Compute joint accelerations according to: \ddot{q} = H^{-1} ( J^T f) 15 | current_accelerations = np.dot(jacobian_transpose, net_force) 16 | 17 | # Integrate once, starting with zero motion 18 | current_velocities = 0.5 * current_accelerations * period 19 | 20 | # Integrate twice, starting with zero motion 21 | current_positions = self.last_positions + 0.5 * current_velocities * period 22 | 23 | # Update for the next cycle 24 | self.last_positions = current_positions 25 | 26 | return current_positions, current_velocities -------------------------------------------------------------------------------- /homestri_ur5e_rl/utils/transform_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | 4 | EPS = np.finfo(float).eps * 4.0 5 | 6 | def quat_conjugate(quaternion): 7 | """ 8 | Return conjugate of quaternion. 9 | 10 | E.g.: 11 | >>> q0 = random_quaternion() 12 | >>> q1 = quat_conjugate(q0) 13 | >>> q1[3] == q0[3] and all(q1[:3] == -q0[:3]) 14 | True 15 | 16 | Args: 17 | quaternion (np.array): (x,y,z,w) quaternion 18 | 19 | Returns: 20 | np.array: (x,y,z,w) quaternion conjugate 21 | """ 22 | return np.array( 23 | (-quaternion[0], -quaternion[1], -quaternion[2], quaternion[3]), 24 | dtype=np.float32, 25 | ) 26 | 27 | def quat_multiply(quaternion1, quaternion0): 28 | """ 29 | Return multiplication of two quaternions (q1 * q0). 30 | 31 | E.g.: 32 | >>> q = quat_multiply([1, -2, 3, 4], [-5, 6, 7, 8]) 33 | >>> np.allclose(q, [-44, -14, 48, 28]) 34 | True 35 | 36 | Args: 37 | quaternion1 (np.array): (x,y,z,w) quaternion 38 | quaternion0 (np.array): (x,y,z,w) quaternion 39 | 40 | Returns: 41 | np.array: (x,y,z,w) multiplied quaternion 42 | """ 43 | x0, y0, z0, w0 = quaternion0 44 | x1, y1, z1, w1 = quaternion1 45 | return np.array( 46 | ( 47 | x1 * w0 + y1 * z0 - z1 * y0 + w1 * x0, 48 | -x1 * z0 + y1 * w0 + z1 * x0 + w1 * y0, 49 | x1 * y0 - y1 * x0 + z1 * w0 + w1 * z0, 50 | -x1 * x0 - y1 * y0 - z1 * z0 + w1 * w0, 51 | ), 52 | dtype=np.float32, 53 | ) 54 | 55 | def mat2quat(rmat): 56 | """ 57 | Converts given rotation matrix to quaternion. 58 | 59 | Args: 60 | rmat (np.array): 3x3 rotation matrix 61 | 62 | Returns: 63 | np.array: (x,y,z,w) float quaternion angles 64 | """ 65 | M = np.asarray(rmat).astype(np.float32)[:3, :3] 66 | 67 | m00 = M[0, 0] 68 | m01 = M[0, 1] 69 | m02 = M[0, 2] 70 | m10 = M[1, 0] 71 | m11 = M[1, 1] 72 | m12 = M[1, 2] 73 | m20 = M[2, 0] 74 | m21 = M[2, 1] 75 | m22 = M[2, 2] 76 | # symmetric matrix K 77 | K = np.array( 78 | [ 79 | [m00 - m11 - m22, np.float32(0.0), np.float32(0.0), np.float32(0.0)], 80 | [m01 + m10, m11 - m00 - m22, np.float32(0.0), np.float32(0.0)], 81 | [m02 + m20, m12 + m21, m22 - m00 - m11, np.float32(0.0)], 82 | [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22], 83 | ] 84 | ) 85 | K /= 3.0 86 | # quaternion is Eigen vector of K that corresponds to largest eigenvalue 87 | w, V = np.linalg.eigh(K) 88 | inds = np.array([3, 0, 1, 2]) 89 | q1 = V[inds, np.argmax(w)] 90 | if q1[0] < 0.0: 91 | np.negative(q1, q1) 92 | inds = np.array([1, 2, 3, 0]) 93 | return q1[inds] 94 | 95 | def quat2mat(quaternion): 96 | """ 97 | Converts given quaternion to matrix. 98 | 99 | Args: 100 | quaternion (np.array): (x,y,z,w) vec4 float angles 101 | 102 | Returns: 103 | np.array: 3x3 rotation matrix 104 | """ 105 | # awkward semantics for use with numba 106 | inds = np.array([3, 0, 1, 2]) 107 | q = np.asarray(quaternion).copy().astype(np.float32)[inds] 108 | 109 | n = np.dot(q, q) 110 | if n < EPS: 111 | return np.identity(3) 112 | q *= math.sqrt(2.0 / n) 113 | q2 = np.outer(q, q) 114 | return np.array( 115 | [ 116 | [1.0 - q2[2, 2] - q2[3, 3], q2[1, 2] - q2[3, 0], q2[1, 3] + q2[2, 0]], 117 | [q2[1, 2] + q2[3, 0], 1.0 - q2[1, 1] - q2[3, 3], q2[2, 3] - q2[1, 0]], 118 | [q2[1, 3] - q2[2, 0], q2[2, 3] + q2[1, 0], 1.0 - q2[1, 1] - q2[2, 2]], 119 | ] 120 | ) 121 | 122 | def quat2axisangle(quat): 123 | """ 124 | Converts quaternion to axis-angle format. 125 | Returns a unit vector direction scaled by its angle in radians. 126 | 127 | Args: 128 | quat (np.array): (x,y,z,w) vec4 float angles 129 | 130 | Returns: 131 | np.array: (ax,ay,az) axis-angle exponential coordinates 132 | """ 133 | # clip quaternion 134 | if quat[3] > 1.0: 135 | quat[3] = 1.0 136 | elif quat[3] < -1.0: 137 | quat[3] = -1.0 138 | 139 | den = np.sqrt(1.0 - quat[3] * quat[3]) 140 | if math.isclose(den, 0.0): 141 | # This is (close to) a zero degree rotation, immediately return 142 | return np.zeros(3) 143 | 144 | return (quat[:3] * 2.0 * math.acos(quat[3])) / den 145 | -------------------------------------------------------------------------------- /media/wire_manipulation.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ian-chuang/homestri-ur5e-rl/89b7e23cc5acbd9a164222fa548cb0113684523e/media/wire_manipulation.gif -------------------------------------------------------------------------------- /scripts/circle_points.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | 4 | 5 | def half_circle_with_equal_arc_lengths(radius, num_points): 6 | if num_points < 2: 7 | raise ValueError("The number of points must be at least 2.") 8 | 9 | # Calculate the angle increment to achieve equally spaced arc lengths 10 | angle_increment = math.pi / (num_points - 1) 11 | 12 | # Generate the Cartesian coordinates for each point 13 | cartesian_coordinates = [] 14 | for i in range(num_points)[::-1]: 15 | angle = i * angle_increment 16 | x = radius * math.cos(angle) 17 | y = radius * math.sin(angle) 18 | cartesian_coordinates.append((x + radius, y)) 19 | 20 | return cartesian_coordinates 21 | 22 | # Example usage: 23 | radius = 0.25 24 | num_points = 20 25 | half_circle_points = half_circle_with_equal_arc_lengths(radius, num_points) 26 | 27 | for point in half_circle_points: 28 | print(f"{point[0]} {point[1]} 0") -------------------------------------------------------------------------------- /scripts/test_base_robot_env.py: -------------------------------------------------------------------------------- 1 | import gymnasium as gym 2 | import homestri_ur5e_rl 3 | from homestri_ur5e_rl.input_devices.keyboard_input import ( 4 | KeyboardInput, 5 | ) 6 | 7 | keyboard_input = KeyboardInput() 8 | 9 | env = gym.make("BaseRobot-v0", render_mode="human") 10 | 11 | observation, info = env.reset(seed=42) 12 | while True: 13 | action = env.action_space.sample() 14 | 15 | action = keyboard_input.get_action() 16 | 17 | observation, reward, terminated, truncated, info = env.step(action) 18 | 19 | if terminated or truncated: 20 | observation, info = env.reset() 21 | 22 | env.close() 23 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name='homestri_ur5e_rl', 5 | version='1.0.0', 6 | packages=find_packages(), 7 | install_requires=[ 8 | 'gymnasium-robotics', 9 | 'mujoco', 10 | 'gymnasium', 11 | 'pynput', 12 | ], 13 | # Other information like author, author_email, description, etc. 14 | ) 15 | --------------------------------------------------------------------------------