Python pybullet.getJointStates() Examples

def getAnglesVelocity(self, joint_names):
        """
        Gets the velocity of the robot's joints in rad/s. If one of the joint
        doesn't exist, the method will raise a KeyError.

        Parameters:
            joint_names - List of string containing the names of the joints

        Returns:
            joint_velocities - List of floats containing the joint's velocities
        """
        indexes = [self.joint_dict[name].getIndex() for name in joint_names]

        return [state[1] for state in pybullet.getJointStates(
            self.robot_model,
            indexes,
            physicsClientId=self.physics_client)] 

def _get_obs(self, forces, forces_human):
        torso_pos = np.array(p.getLinkState(self.robot, 15 if self.robot_type == 'pr2' else 0, computeForwardKinematics=True, physicsClientId=self.id)[0])
        tool_pos, tool_orient = p.getBasePositionAndOrientation(self.cup, physicsClientId=self.id)
        robot_joint_states = p.getJointStates(self.robot, jointIndices=self.robot_right_arm_joint_indices, physicsClientId=self.id)
        robot_joint_positions = np.array([x[0] for x in robot_joint_states])
        robot_pos, robot_orient = p.getBasePositionAndOrientation(self.robot, physicsClientId=self.id)
        if self.human_control:
            human_pos = np.array(p.getBasePositionAndOrientation(self.human, physicsClientId=self.id)[0])
            human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
            human_joint_positions = np.array([x[0] for x in human_joint_states])

        head_pos, head_orient = p.getLinkState(self.human, 23, computeForwardKinematics=True, physicsClientId=self.id)[:2]

        robot_obs = np.concatenate([tool_pos-torso_pos, tool_orient, tool_pos - self.target_pos, robot_joint_positions, head_pos-torso_pos, head_orient, forces]).ravel()
        if self.human_control:
            human_obs = np.concatenate([tool_pos-human_pos, tool_orient, tool_pos - self.target_pos, human_joint_positions, head_pos-human_pos, head_orient, forces_human]).ravel()
        else:
            human_obs = []

        return np.concatenate([robot_obs, human_obs]).ravel() 

def _get_obs(self, forces, forces_human):
        torso_pos = np.array(p.getLinkState(self.robot, 15 if self.robot_type == 'pr2' else 0, computeForwardKinematics=True, physicsClientId=self.id)[0])
        spoon_pos, spoon_orient = p.getBasePositionAndOrientation(self.spoon, physicsClientId=self.id)
        robot_right_joint_states = p.getJointStates(self.robot, jointIndices=self.robot_right_arm_joint_indices, physicsClientId=self.id)
        robot_right_joint_positions = np.array([x[0] for x in robot_right_joint_states])
        robot_pos, robot_orient = p.getBasePositionAndOrientation(self.robot, physicsClientId=self.id)
        if self.human_control:
            human_pos = np.array(p.getBasePositionAndOrientation(self.human, physicsClientId=self.id)[0])
            human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
            human_joint_positions = np.array([x[0] for x in human_joint_states])

        head_pos, head_orient = p.getLinkState(self.human, 23, computeForwardKinematics=True, physicsClientId=self.id)[:2]

        robot_obs = np.concatenate([spoon_pos-torso_pos, spoon_orient, spoon_pos-self.target_pos, robot_right_joint_positions, head_pos-torso_pos, head_orient, forces]).ravel()
        if self.human_control:
            human_obs = np.concatenate([spoon_pos-human_pos, spoon_orient, spoon_pos-self.target_pos, human_joint_positions, head_pos-human_pos, head_orient, forces_human]).ravel()
        else:
            human_obs = []

        return np.concatenate([robot_obs, human_obs]).ravel() 

def enforce_joint_limits(self, body):
        # Enforce joint limits
        joint_states = p.getJointStates(body, jointIndices=list(range(p.getNumJoints(body, physicsClientId=self.id))), physicsClientId=self.id)
        joint_positions = np.array([x[0] for x in joint_states])
        lower_limits = []
        upper_limits = []
        for j in range(p.getNumJoints(body, physicsClientId=self.id)):
            joint_info = p.getJointInfo(body, j, physicsClientId=self.id)
            joint_name = joint_info[1]
            joint_pos = joint_positions[j]
            lower_limit = joint_info[8]
            upper_limit = joint_info[9]
            if lower_limit == 0 and upper_limit == -1:
                lower_limit = -1e10
                upper_limit = 1e10
            lower_limits.append(lower_limit)
            upper_limits.append(upper_limit)
            # print(joint_name, joint_pos, lower_limit, upper_limit)
            if joint_pos < lower_limit:
                p.resetJointState(body, jointIndex=j, targetValue=lower_limit, targetVelocity=0, physicsClientId=self.id)
            elif joint_pos > upper_limit:
                p.resetJointState(body, jointIndex=j, targetValue=upper_limit, targetVelocity=0, physicsClientId=self.id)
        lower_limits = np.array(lower_limits)
        upper_limits = np.array(upper_limits)
        return lower_limits, upper_limits 

def getAnglesPosition(self, joint_names):
        """
        Gets the position of the robot's joints in radians. If one of the joint
        doesn't exist, the method will raise a KeyError.

        Parameters:
            joint_names - List of string containing the names of the joints

        Returns:
            joint_positions - List of floats containing the joint's positions
        """
        indexes = [self.joint_dict[name].getIndex() for name in joint_names]

        return [state[0] for state in pybullet.getJointStates(
            self.robot_model,
            indexes,
            physicsClientId=self.physics_client)] 

def enforce_hard_human_joint_limits(self):
        if not self.human_controllable_joint_indices:
            return
        # Enforce joint limits. Sometimes, external forces and break the person's hard joint limits.
        joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
        joint_positions = np.array([x[0] for x in joint_states])
        if self.human_joint_lower_limits is None:
            self.human_joint_lower_limits = []
            self.human_joint_upper_limits = []
            for i, j in enumerate(self.human_controllable_joint_indices):
                joint_info = p.getJointInfo(self.human, j, physicsClientId=self.id)
                joint_name = joint_info[1]
                joint_pos = joint_positions[i]
                lower_limit = joint_info[8]
                upper_limit = joint_info[9]
                self.human_joint_lower_limits.append(lower_limit)
                self.human_joint_upper_limits.append(upper_limit)
                # print(joint_name, joint_pos, lower_limit, upper_limit)
        for i, j in enumerate(self.human_controllable_joint_indices):
            if joint_positions[i] < self.human_joint_lower_limits[i]:
                p.resetJointState(self.human, jointIndex=j, targetValue=self.human_joint_lower_limits[i], targetVelocity=0, physicsClientId=self.id)
            elif joint_positions[i] > self.human_joint_upper_limits[i]:
                p.resetJointState(self.human, jointIndex=j, targetValue=self.human_joint_upper_limits[i], targetVelocity=0, physicsClientId=self.id) 

def _get_obs(self, forces, forces_human):
        torso_pos = np.array(p.getLinkState(self.robot, 15 if self.robot_type == 'pr2' else 0, computeForwardKinematics=True, physicsClientId=self.id)[0])
        state = p.getLinkState(self.tool, 1, computeForwardKinematics=True, physicsClientId=self.id)
        tool_pos = np.array(state[0])
        tool_orient = np.array(state[1]) # Quaternions
        robot_joint_states = p.getJointStates(self.robot, jointIndices=self.robot_left_arm_joint_indices, physicsClientId=self.id)
        robot_joint_positions = np.array([x[0] for x in robot_joint_states])
        robot_pos, robot_orient = p.getBasePositionAndOrientation(self.robot, physicsClientId=self.id)
        if self.human_control:
            human_pos = np.array(p.getBasePositionAndOrientation(self.human, physicsClientId=self.id)[0])
            human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
            human_joint_positions = np.array([x[0] for x in human_joint_states])

        # Human shoulder, elbow, and wrist joint locations
        shoulder_pos, shoulder_orient = p.getLinkState(self.human, 5, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        elbow_pos, elbow_orient = p.getLinkState(self.human, 7, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        wrist_pos, wrist_orient = p.getLinkState(self.human, 9, computeForwardKinematics=True, physicsClientId=self.id)[:2]

        robot_obs = np.concatenate([tool_pos-torso_pos, tool_orient, robot_joint_positions, shoulder_pos-torso_pos, elbow_pos-torso_pos, wrist_pos-torso_pos, forces]).ravel()
        if self.human_control:
            human_obs = np.concatenate([tool_pos-human_pos, tool_orient, human_joint_positions, shoulder_pos-human_pos, elbow_pos-human_pos, wrist_pos-human_pos, forces_human]).ravel()
        else:
            human_obs = []

        return np.concatenate([robot_obs, human_obs]).ravel() 

def _get_obs(self, forces, forces_human):
        torso_pos = np.array(p.getLinkState(self.robot, 15 if self.robot_type == 'pr2' else 0, computeForwardKinematics=True, physicsClientId=self.id)[0])
        state = p.getLinkState(self.tool, 1, computeForwardKinematics=True, physicsClientId=self.id)
        tool_pos = np.array(state[0])
        tool_orient = np.array(state[1]) # Quaternions
        robot_joint_states = p.getJointStates(self.robot, jointIndices=self.robot_left_arm_joint_indices, physicsClientId=self.id)
        robot_joint_positions = np.array([x[0] for x in robot_joint_states])
        robot_pos, robot_orient = p.getBasePositionAndOrientation(self.robot, physicsClientId=self.id)
        if self.human_control:
            human_pos = np.array(p.getBasePositionAndOrientation(self.human, physicsClientId=self.id)[0])
            human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
            human_joint_positions = np.array([x[0] for x in human_joint_states])

        # Human shoulder, elbow, and wrist joint locations
        shoulder_pos, shoulder_orient = p.getLinkState(self.human, 5, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        elbow_pos, elbow_orient = p.getLinkState(self.human, 7, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        wrist_pos, wrist_orient = p.getLinkState(self.human, 9, computeForwardKinematics=True, physicsClientId=self.id)[:2]

        robot_obs = np.concatenate([tool_pos-torso_pos, tool_orient, tool_pos - self.target_pos, self.target_pos-torso_pos, robot_joint_positions, shoulder_pos-torso_pos, elbow_pos-torso_pos, wrist_pos-torso_pos, forces]).ravel()
        if self.human_control:
            human_obs = np.concatenate([tool_pos-human_pos, tool_orient, tool_pos - self.target_pos, self.target_pos-human_pos, human_joint_positions, shoulder_pos-human_pos, elbow_pos-human_pos, wrist_pos-human_pos, forces_human]).ravel()
        else:
            human_obs = []

        return np.concatenate([robot_obs, human_obs]).ravel() 

def getMotorJointStates(robot):
	joint_states = p.getJointStates(robot, range(p.getNumJoints(robot)))
	joint_infos = [p.getJointInfo(robot, i) for i in range(p.getNumJoints(robot))]
	joint_states = [j for j, i in zip(joint_states, joint_infos) if i[3] > -1]
	joint_positions = [state[0] for state in joint_states]
	joint_velocities = [state[1] for state in joint_states]
	joint_torques = [state[3] for state in joint_states]
	return joint_positions, joint_velocities, joint_torques 

def grasp(self, pos=None):
        # close fingers

        if self._control_arm is 'l':
            idx_thumb = self._joint_name_to_ids['l_hand::l_tj2']
            idx_fingers = [self._joint_name_to_ids[jn] for jn in self.joint_groups['l_hand']]
        else:
            idx_thumb = self._joint_name_to_ids['r_hand::r_tj2']
            idx_fingers = [self._joint_name_to_ids[jn] for jn in self.joint_groups['r_hand']]

        # # set also position to other joints to avoid weird movements
        # not_idx_fingers = [idx for idx in self._joints_to_control if idx not in idx_fingers]
        #
        # joint_states = p.getJointStates(self.robot_id, not_idx_fingers)
        # joint_poses = [x[0] for x in joint_states]
        # p.setJointMotorControlArray(self.robot_id, not_idx_fingers, p.POSITION_CONTROL,
        #                             targetPositions=joint_poses,
        #                             positionGains=[0.1] * len(not_idx_fingers),
        #                             velocityGains=[1.0] * len(not_idx_fingers),
        #                             physicsClientId=self._physics_client_id)

        position_control = True
        if position_control:
            if pos is None:
                pos = self._grasp_pos
            p.setJointMotorControlArray(self.robot_id, idx_fingers, p.POSITION_CONTROL,
                                        targetPositions=pos,
                                        positionGains=[0.1] * len(idx_fingers),
                                        velocityGains=[1.0] * len(idx_fingers),
                                        forces=[10] * len(idx_fingers),
                                        physicsClientId=self._physics_client_id)

        else:
            # vel = [0, 1, 1, 1.2, 0, 1, 1, 1.2, 0, 1, 1, 1.2, 0, 1, 1, 1.2, 1.57, 1.5, 1.1, 1.1]
            vel = [0, 0.5, 0.6, 0.6, 0, 0.5, 0.6, 0.6, 0, 0.5, 0.6, 0.6, 0, 0.5, 0.6, 0.6, 0, 0.5, 0.6, 0.6]

            p.setJointMotorControlArray(self.robot_id, idx_fingers, p.VELOCITY_CONTROL,
                                        targetVelocities=vel,
                                        positionGains=[0.1] * len(idx_fingers),
                                        velocityGains=[1.0] * len(idx_fingers),
                                        physicsClientId=self._physics_client_id) 

def _getGripper(self):
        vs = [v[0] for v in pb.getJointStates(self.handle,
            self.stable_gripper_indices)]
        return np.array([np.round(-np.mean(vs),1)]) 

def _get_obs(self, forces, forces_human):
        torso_pos = np.array(p.getLinkState(self.robot, 15 if self.robot_type == 'pr2' else 0, computeForwardKinematics=True, physicsClientId=self.id)[0])
        tool_left_pos, tool_left_orient = p.getLinkState(self.robot, 78 if self.robot_type=='pr2' else 24 if self.robot_type=='sawyer' else 54 if self.robot_type=='baxter' else 9 if self.robot_type=='jaco' else 7, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        tool_right_pos, tool_right_orient = p.getLinkState(self.robot, 55 if self.robot_type=='pr2' else 24 if self.robot_type=='sawyer' else 31 if self.robot_type=='baxter' else 9 if self.robot_type=='jaco' else 7, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        robot_joint_states = p.getJointStates(self.robot, jointIndices=self.robot_both_arm_joint_indices, physicsClientId=self.id)
        robot_joint_positions = np.array([x[0] for x in robot_joint_states])
        robot_pos, robot_orient = p.getBasePositionAndOrientation(self.robot, physicsClientId=self.id)
        if self.human_control:
            human_pos = np.array(p.getBasePositionAndOrientation(self.human, physicsClientId=self.id)[0])
            human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
            human_joint_positions = np.array([x[0] for x in human_joint_states])

        # Human shoulder, elbow, and wrist joint locations
        shoulder_pos, shoulder_orient = p.getLinkState(self.human, 5, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        elbow_pos, elbow_orient = p.getLinkState(self.human, 7, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        wrist_pos, wrist_orient = p.getLinkState(self.human, 9, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        waist_pos = np.array(p.getLinkState(self.human, 24, computeForwardKinematics=True, physicsClientId=self.id)[0])
        hips_pos = np.array(p.getLinkState(self.human, 27, computeForwardKinematics=True, physicsClientId=self.id)[0])
        shoulder_pos, elbow_pos, wrist_pos, waist_pos, hips_pos = np.array(shoulder_pos), np.array(elbow_pos), np.array(wrist_pos), np.array(waist_pos), np.array(hips_pos)

        robot_obs = np.concatenate([tool_left_pos-torso_pos, tool_left_orient, tool_right_pos-torso_pos, tool_right_orient, robot_joint_positions, shoulder_pos-torso_pos, elbow_pos-torso_pos, wrist_pos-torso_pos, waist_pos-torso_pos, hips_pos-torso_pos, forces]).ravel()
        if self.human_control:
            human_obs = np.concatenate([tool_left_pos-human_pos, tool_left_orient, tool_right_pos-human_pos, tool_right_orient, human_joint_positions, shoulder_pos-human_pos, elbow_pos-human_pos, wrist_pos-human_pos, waist_pos-human_pos, hips_pos-human_pos, forces_human]).ravel()
        else:
            human_obs = []

        return np.concatenate([robot_obs, human_obs]).ravel() 

def get_motor_joint_states(self, robot):
        num_joints = p.getNumJoints(robot, physicsClientId=self.id)
        joint_states = p.getJointStates(robot, range(num_joints), physicsClientId=self.id)
        joint_infos = [p.getJointInfo(robot, i, physicsClientId=self.id) for i in range(num_joints)]
        joint_states = [j for j, i in zip(joint_states, joint_infos) if i[2] != p.JOINT_FIXED]
        joint_positions = [state[0] for state in joint_states]
        joint_velocities = [state[1] for state in joint_states]
        joint_torques = [state[3] for state in joint_states]
        return joint_positions, joint_velocities, joint_torques 

def setup_human_joints(self, human, joints_positions, controllable_joints, use_static_joints=True, human_reactive_force=None, human_reactive_gain=0.05):
        if self.human_impairment != 'tremor':
            self.human_tremors = np.zeros(len(controllable_joints))
        elif len(controllable_joints) == 4:
            self.human_tremors = self.np_random.uniform(np.deg2rad(-20), np.deg2rad(20), size=len(controllable_joints))
        else:
            self.human_tremors = self.np_random.uniform(np.deg2rad(-10), np.deg2rad(10), size=len(controllable_joints))
        # Set starting joint positions
        human_joint_states = p.getJointStates(human, jointIndices=list(range(p.getNumJoints(human, physicsClientId=self.id))), physicsClientId=self.id)
        human_joint_positions = np.array([x[0] for x in human_joint_states])
        for j in range(p.getNumJoints(human, physicsClientId=self.id)):
            set_position = None
            for j_index, j_angle in joints_positions:
                if j == j_index:
                    p.resetJointState(human, jointIndex=j, targetValue=j_angle, targetVelocity=0, physicsClientId=self.id)
                    set_position = j_angle
                    break
            if use_static_joints and j not in controllable_joints:
                # Make all non controllable joints on the person static by setting mass of each link (joint) to 0
                p.changeDynamics(human, j, mass=0, physicsClientId=self.id)
                # Set velocities to 0
                p.resetJointState(human, jointIndex=j, targetValue=human_joint_positions[j] if set_position is None else set_position, targetVelocity=0, physicsClientId=self.id)

        # By default, all joints have motors enabled by default that prevent free motion. Disable these motors in human
        for j in range(p.getNumJoints(human, physicsClientId=self.id)):
            p.setJointMotorControl2(human, jointIndex=j, controlMode=p.VELOCITY_CONTROL, force=0, physicsClientId=self.id)

        self.enforce_joint_limits(human)

        if human_reactive_force is not None:
            # NOTE: This runs a Position / Velocity PD controller for each joint motor on the human
            human_joint_states = p.getJointStates(human, jointIndices=controllable_joints, physicsClientId=self.id)
            target_human_joint_positions = np.array([x[0] for x in human_joint_states])
            forces = [human_reactive_force * self.human_strength] * len(target_human_joint_positions)
            p.setJointMotorControlArray(human, jointIndices=controllable_joints, controlMode=p.POSITION_CONTROL, targetPositions=target_human_joint_positions, positionGains=np.array([human_reactive_gain]*len(forces)), forces=forces, physicsClientId=self.id) 

def getMotorJointStates(self, robot):
        import pybullet as p
        joint_states = p.getJointStates(robot, range(p.getNumJoints(robot)))
        joint_infos = [p.getJointInfo(robot, i) for i in range(p.getNumJoints(robot))]
        joint_states = [j for j, i in zip(joint_states, joint_infos) if i[3] > -1]
        joint_positions = [state[0] for state in joint_states]
        joint_velocities = [state[1] for state in joint_states]
        joint_torques = [state[3] for state in joint_states]
        return joint_positions, joint_velocities, joint_torques 

def getJointStates(robot):
	joint_states = p.getJointStates(robot, range(p.getNumJoints(robot)))
	joint_positions = [state[0] for state in joint_states]
	joint_velocities = [state[1] for state in joint_states]
	joint_torques = [state[3] for state in joint_states]
	return joint_positions, joint_velocities, joint_torques 

def reset(self):
        self.setup_timing()
        self.task_success = 0
        self.prev_target_contact_pos = np.zeros(3)
        self.human, self.wheelchair, self.robot, self.robot_lower_limits, self.robot_upper_limits, self.human_lower_limits, self.human_upper_limits, self.robot_right_arm_joint_indices, self.robot_left_arm_joint_indices, self.gender = self.world_creation.create_new_world(furniture_type='wheelchair', static_human_base=True, human_impairment='random', print_joints=False, gender='random')
        self.robot_lower_limits = self.robot_lower_limits[self.robot_left_arm_joint_indices]
        self.robot_upper_limits = self.robot_upper_limits[self.robot_left_arm_joint_indices]
        self.reset_robot_joints()
        if self.robot_type == 'jaco':
            wheelchair_pos, wheelchair_orient = p.getBasePositionAndOrientation(self.wheelchair, physicsClientId=self.id)
            p.resetBasePositionAndOrientation(self.robot, np.array(wheelchair_pos) + np.array([-0.35, -0.3, 0.3]), p.getQuaternionFromEuler([0, 0, -np.pi/2.0], physicsClientId=self.id), physicsClientId=self.id)

        joints_positions = [(3, np.deg2rad(30)), (6, np.deg2rad(-90)), (16, np.deg2rad(-90)), (28, np.deg2rad(-90)), (31, np.deg2rad(80)), (35, np.deg2rad(-90)), (38, np.deg2rad(80))]
        self.human_controllable_joint_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        self.world_creation.setup_human_joints(self.human, joints_positions, self.human_controllable_joint_indices, use_static_joints=True, human_reactive_force=None if self.human_control else 1, human_reactive_gain=0.01)
        p.resetBasePositionAndOrientation(self.human, [0, 0.03, 0.89 if self.gender == 'male' else 0.86], [0, 0, 0, 1], physicsClientId=self.id)
        human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
        self.target_human_joint_positions = np.array([x[0] for x in human_joint_states])
        self.human_lower_limits = self.human_lower_limits[self.human_controllable_joint_indices]
        self.human_upper_limits = self.human_upper_limits[self.human_controllable_joint_indices]

        shoulder_pos, shoulder_orient = p.getLinkState(self.human, 5, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        elbow_pos, elbow_orient = p.getLinkState(self.human, 7, computeForwardKinematics=True, physicsClientId=self.id)[:2]
        wrist_pos, wrist_orient = p.getLinkState(self.human, 9, computeForwardKinematics=True, physicsClientId=self.id)[:2]

        if self.robot_type == 'pr2':
            target_pos = np.array([-0.55, 0, 0.8]) + self.np_random.uniform(-0.05, 0.05, size=3)
            target_orient = np.array(p.getQuaternionFromEuler(np.array([0, 0, 0]), physicsClientId=self.id))
            self.position_robot_toc(self.robot, 76, [(target_pos, target_orient)], [(shoulder_pos, None), (elbow_pos, None), (wrist_pos, None)], self.robot_left_arm_joint_indices, self.robot_lower_limits, self.robot_upper_limits, ik_indices=range(29, 29+7), pos_offset=np.array([0.1, 0, 0]), max_ik_iterations=200, step_sim=True, check_env_collisions=False, human_joint_indices=self.human_controllable_joint_indices, human_joint_positions=self.target_human_joint_positions)
            self.world_creation.set_gripper_open_position(self.robot, position=0.25, left=True, set_instantly=True)
            self.tool = self.world_creation.init_tool(self.robot, mesh_scale=[0.001]*3, pos_offset=[0, 0, 0], orient_offset=p.getQuaternionFromEuler([0, 0, 0], physicsClientId=self.id), maximal=False)
        elif self.robot_type == 'jaco':
            target_pos = np.array([-0.5, 0, 0.8]) + self.np_random.uniform(-0.05, 0.05, size=3)
            target_orient = p.getQuaternionFromEuler(np.array([0, np.pi/2.0, 0]), physicsClientId=self.id)
            self.util.ik_random_restarts(self.robot, 8, target_pos, target_orient, self.world_creation, self.robot_left_arm_joint_indices, self.robot_lower_limits, self.robot_upper_limits, ik_indices=[0, 1, 2, 3, 4, 5, 6], max_iterations=1000, max_ik_random_restarts=40, random_restart_threshold=0.03, step_sim=True)
            self.world_creation.set_gripper_open_position(self.robot, position=1, left=True, set_instantly=True)
            self.tool = self.world_creation.init_tool(self.robot, mesh_scale=[0.001]*3, pos_offset=[0, 0, 0.02], orient_offset=p.getQuaternionFromEuler([0, -np.pi/2.0, 0], physicsClientId=self.id), maximal=False)
        else:
            target_pos = np.array([-0.55, 0, 0.8]) + self.np_random.uniform(-0.05, 0.05, size=3)
            target_orient = p.getQuaternionFromEuler(np.array([0, np.pi/2.0, 0]), physicsClientId=self.id)
            if self.robot_type == 'baxter':
                self.position_robot_toc(self.robot, 48, [(target_pos, target_orient)], [(shoulder_pos, None), (elbow_pos, None), (wrist_pos, None)], self.robot_left_arm_joint_indices, self.robot_lower_limits, self.robot_upper_limits, ik_indices=range(10, 17), pos_offset=np.array([0, 0, 0.975]), max_ik_iterations=200, step_sim=True, check_env_collisions=False, human_joint_indices=self.human_controllable_joint_indices, human_joint_positions=self.target_human_joint_positions)
            else:
                self.position_robot_toc(self.robot, 19, [(target_pos, target_orient)], [(shoulder_pos, None), (elbow_pos, None), (wrist_pos, None)], self.robot_left_arm_joint_indices, self.robot_lower_limits, self.robot_upper_limits, ik_indices=[0, 2, 3, 4, 5, 6, 7], pos_offset=np.array([-0.1, 0, 0.975]), max_ik_iterations=200, step_sim=True, check_env_collisions=False, human_joint_indices=self.human_controllable_joint_indices, human_joint_positions=self.target_human_joint_positions)
            self.world_creation.set_gripper_open_position(self.robot, position=0.015, left=True, set_instantly=True)
            self.tool = self.world_creation.init_tool(self.robot, mesh_scale=[0.001]*3, pos_offset=[0, 0.125, 0], orient_offset=p.getQuaternionFromEuler([0, 0, np.pi/2.0], physicsClientId=self.id), maximal=False)

        self.generate_target()

        p.setGravity(0, 0, 0, physicsClientId=self.id)
        p.setGravity(0, 0, -1, body=self.human, physicsClientId=self.id)

        # Enable rendering
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1, physicsClientId=self.id)

        return self._get_obs([0], [0, 0]) 

def take_step(self, action, robot_arm='left', gains=0.05, forces=1, human_gains=0.1, human_forces=1, step_sim=True):
        action = np.clip(action, a_min=self.action_space.low, a_max=self.action_space.high)
        # print('cameraYaw=%.2f, cameraPitch=%.2f, distance=%.2f' % p.getDebugVisualizerCamera(physicsClientId=self.id)[-4:-1])

        # print('Total time:', self.total_time)
        # self.total_time += 0.1
        self.iteration += 1
        if self.last_sim_time is None:
            self.last_sim_time = time.time()

        action *= 0.05
        action_robot = action
        indices = self.robot_left_arm_joint_indices if robot_arm == 'left' else self.robot_right_arm_joint_indices if robot_arm == 'right' else self.robot_both_arm_joint_indices

        if self.human_control or (self.world_creation.human_impairment == 'tremor' and self.human_controllable_joint_indices):
            human_len = len(self.human_controllable_joint_indices)
            if self.human_control:
                action_robot = action[:self.action_robot_len]
                action_human = action[self.action_robot_len:]
            else:
                action_human = np.zeros(human_len)
            if len(action_human) != human_len:
                print('Received human actions of length %d does not match expected action length of %d' % (len(action_human), human_len))
                exit()
            human_joint_states = p.getJointStates(self.human, jointIndices=self.human_controllable_joint_indices, physicsClientId=self.id)
            human_joint_positions = np.array([x[0] for x in human_joint_states])

        robot_joint_states = p.getJointStates(self.robot, jointIndices=indices, physicsClientId=self.id)
        robot_joint_positions = np.array([x[0] for x in robot_joint_states])

        for _ in range(self.frame_skip):
            action_robot[robot_joint_positions + action_robot < self.robot_lower_limits] = 0
            action_robot[robot_joint_positions + action_robot > self.robot_upper_limits] = 0
            robot_joint_positions += action_robot
            if self.human_control or (self.world_creation.human_impairment == 'tremor' and self.human_controllable_joint_indices):
                action_human[human_joint_positions + action_human < self.human_lower_limits] = 0
                action_human[human_joint_positions + action_human > self.human_upper_limits] = 0
                if self.world_creation.human_impairment == 'tremor':
                    human_joint_positions = self.target_human_joint_positions + self.world_creation.human_tremors * (1 if self.iteration % 2 == 0 else -1)
                    self.target_human_joint_positions += action_human
                human_joint_positions += action_human

        p.setJointMotorControlArray(self.robot, jointIndices=indices, controlMode=p.POSITION_CONTROL, targetPositions=robot_joint_positions, positionGains=np.array([gains]*self.action_robot_len), forces=[forces]*self.action_robot_len, physicsClientId=self.id)
        if self.human_control or (self.world_creation.human_impairment == 'tremor' and self.human_controllable_joint_indices):
            p.setJointMotorControlArray(self.human, jointIndices=self.human_controllable_joint_indices, controlMode=p.POSITION_CONTROL, targetPositions=human_joint_positions, positionGains=np.array([human_gains]*human_len), forces=[human_forces*self.world_creation.human_strength]*human_len, physicsClientId=self.id)

        if step_sim:
            # Update robot position
            for _ in range(self.frame_skip):
                p.stepSimulation(physicsClientId=self.id)
                if self.human_control:
                    self.enforce_realistic_human_joint_limits()
                self.enforce_hard_human_joint_limits()
                self.update_targets()
                if self.gui:
                    # Slow down time so that the simulation matches real time
                    self.slow_time()
            self.record_video_frame() 

def get_observation(self):
        # Create observation state
        observation = []
        observation_lim = []

        # Get state of the end-effector link
        state = p.getLinkState(self.robot_id, self.end_eff_idx, computeLinkVelocity=1,
                                computeForwardKinematics=1, physicsClientId=self._physics_client_id)

        # ------------------------- #
        # --- Cartesian 6D pose --- #
        # ------------------------- #
        pos = state[0]
        quat = state[1]

        observation.extend(list(pos))
        observation_lim.extend(list(self._workspace_lim))

        # cartesian orientation
        if self._control_eu_or_quat is 0:
            euler = p.getEulerFromQuaternion(quat)

            observation.extend(list(euler))
            observation_lim.extend(self._eu_lim)

        else:
            observation.extend(list(quat))
            observation_lim.extend([[-1, 1], [-1, 1], [-1, 1], [-1, 1]])

        # --------------------------------- #
        # --- Cartesian linear velocity --- #
        # --------------------------------- #
        vel_l = state[6]

        observation.extend(list(vel_l))
        observation_lim.extend([[-1, 1], [-1, 1], [-1, 1]])

        # ------------------- #
        # --- Joint poses --- #
        # ------------------- #
        joint_states = p.getJointStates(self.robot_id, self._joints_to_control, physicsClientId=self._physics_client_id)
        joint_poses = [x[0] for x in joint_states]

        observation.extend(list(joint_poses))
        observation_lim.extend([[self.ll[i], self.ul[i]] for i, idx in enumerate(self._joint_name_to_ids.values())
                                if idx in self._joints_to_control])

        return observation, observation_lim 

def get_observation(self):
        # Create observation state
        observation = []
        observation_lim = []

        # Get state of the end-effector link
        state = p.getLinkState(self.robot_id, self.end_eff_idx, computeLinkVelocity=1,
                               computeForwardKinematics=1, physicsClientId=self._physics_client_id)

        # ------------------------- #
        # --- Cartesian 6D pose --- #
        # ------------------------- #
        pos = state[0]
        orn = state[1]

        observation.extend(list(pos))
        observation_lim.extend(list(self._workspace_lim))

        # cartesian orientation
        if self._control_eu_or_quat is 0:
            euler = p.getEulerFromQuaternion(orn)
            observation.extend(list(euler))  # roll, pitch, yaw
            observation_lim.extend(self._eu_lim)
        else:
            observation.extend(list(orn))
            observation_lim.extend([[-1, 1], [-1, 1], [-1, 1], [-1, 1]])

        # --------------------------------- #
        # --- Cartesian linear velocity --- #
        # --------------------------------- #
        if self._include_vel_obs:
            # standardize by subtracting the mean and dividing by the std

            vel_std = [0.04, 0.07, 0.03]
            vel_mean = [0.0, 0.01, 0.0]

            vel_l = np.subtract(state[6], vel_mean)
            vel_l = np.divide(vel_l, vel_std)

            observation.extend(list(vel_l))
            observation_lim.extend([[-1, 1], [-1, 1], [-1, 1]])

        # ------------------- #
        # --- Joint poses --- #
        # ------------------- #

        jointStates = p.getJointStates(self.robot_id, self._joint_name_to_ids.values(), physicsClientId=self._physics_client_id)
        jointPoses = [x[0] for x in jointStates]

        observation.extend(list(jointPoses))
        observation_lim.extend([[self.ll[i], self.ul[i]] for i in range(0, len(self._joint_name_to_ids.values()))])

        return observation, observation_lim