Python pybullet.getQuaternionFromEuler() Examples

def _reset(self):
    self.terminated = 0
    p.resetSimulation()
    p.setPhysicsEngineParameter(numSolverIterations=150)
    p.setTimeStep(self._timeStep)
    p.loadURDF(os.path.join(self._urdfRoot,"plane.urdf"),[0,0,-1])

    p.loadURDF(os.path.join(self._urdfRoot,"table/table.urdf"), 0.5000000,0.00000,-.820000,0.000000,0.000000,0.0,1.0)

    xpos = 0.5 +0.2*random.random()
    ypos = 0 +0.25*random.random()
    ang = 3.1415925438*random.random()
    orn = p.getQuaternionFromEuler([0,0,ang])
    self.blockUid =p.loadURDF(os.path.join(self._urdfRoot,"block.urdf"), xpos,ypos,-0.1,orn[0],orn[1],orn[2],orn[3])

    p.setGravity(0,0,-10)
    self._kuka = kuka.Kuka(urdfRootPath=self._urdfRoot, timeStep=self._timeStep)
    self._envStepCounter = 0
    p.stepSimulation()
    self._observation = self.getExtendedObservation()
    return np.array(self._observation) 

def _reset(self):
        # reset is called once at initialization of simulation
        self.vt = 0
        self.vd = 0
        self.maxV = 24.6 # 235RPM = 24,609142453 rad/sec
        self._envStepCounter = 0

        p.resetSimulation()
        p.setGravity(0,0,-10) # m/s^2
        p.setTimeStep(0.01) # sec
        planeId = p.loadURDF("plane.urdf")
        cubeStartPos = [0,0,0.001]
        cubeStartOrientation = p.getQuaternionFromEuler([0,0,0])

        path = os.path.abspath(os.path.dirname(__file__))
        self.botId = p.loadURDF(os.path.join(path, "balancebot_simple.xml"),
                           cubeStartPos,
                           cubeStartOrientation)

        # you *have* to compute and return the observation from reset()
        self._observation = self._compute_observation()
        return np.array(self._observation) 

def _setup(self):
        '''
        Create the mug at a random position on the ground, handle facing
        roughly towards the robot. Robot's job is to grab and lift.
        '''

        rospack = rospkg.RosPack()
        path = rospack.get_path('costar_objects')
        sdf_dir = os.path.join(path, self.sdf_dir)
        obj_to_add = os.path.join(sdf_dir, self.model, self.model_file_name)

        identity_orientation = pb.getQuaternionFromEuler([0, 0, 0])
        try:
            obj_id_list = pb.loadSDF(obj_to_add)
            for obj_id in obj_id_list:
                random_position = np.random.rand(
                    3) * self.spawn_pos_delta + self.spawn_pos_min
                pb.resetBasePositionAndOrientation(
                    obj_id, random_position, identity_orientation)
        except Exception, e:
            print e 

def ik_random_restarts(self, body, target_joint, target_pos, target_orient, world_creation, robot_arm_joint_indices, robot_lower_limits, robot_upper_limits, ik_indices=range(29, 29+7), max_iterations=1000, max_ik_random_restarts=50, random_restart_threshold=0.01, half_range=False, step_sim=False, check_env_collisions=False):
        orient_orig = target_orient
        best_ik_joints = None
        best_ik_distance = 0
        for r in range(max_ik_random_restarts):
            target_joint_positions = self.ik(body, target_joint, target_pos, target_orient, ik_indices=ik_indices, max_iterations=max_iterations, half_range=half_range)
            world_creation.setup_robot_joints(body, robot_arm_joint_indices, robot_lower_limits, robot_upper_limits, randomize_joint_positions=False, default_positions=np.array(target_joint_positions), tool=None)
            if step_sim:
                for _ in range(5):
                    p.stepSimulation(physicsClientId=self.id)
                if len(p.getContactPoints(bodyA=body, bodyB=body, physicsClientId=self.id)) > 0 and orient_orig is not None:
                    # The robot's arm is in contact with itself. Continually randomize end effector orientation until a solution is found
                    target_orient = p.getQuaternionFromEuler(p.getEulerFromQuaternion(orient_orig, physicsClientId=self.id) + np.deg2rad(self.np_random.uniform(-45, 45, size=3)), physicsClientId=self.id)
            if check_env_collisions:
                for _ in range(25):
                    p.stepSimulation(physicsClientId=self.id)
            gripper_pos, gripper_orient = p.getLinkState(body, target_joint, computeForwardKinematics=True, physicsClientId=self.id)[:2]
            if np.linalg.norm(target_pos - np.array(gripper_pos)) < random_restart_threshold and (target_orient is None or np.linalg.norm(target_orient - np.array(gripper_orient)) < random_restart_threshold or np.isclose(np.linalg.norm(target_orient - np.array(gripper_orient)), 2, atol=random_restart_threshold)):
                return True, np.array(target_joint_positions)
            if best_ik_joints is None or np.linalg.norm(target_pos - np.array(gripper_pos)) < best_ik_distance:
                best_ik_joints = target_joint_positions
                best_ik_distance = np.linalg.norm(target_pos - np.array(gripper_pos))
        world_creation.setup_robot_joints(body, robot_arm_joint_indices, robot_lower_limits, robot_upper_limits, randomize_joint_positions=False, default_positions=np.array(best_ik_joints), tool=None)
        return False, np.array(best_ik_joints) 

def reset(self, height=0.5, orientation='straight'):
        """Resets the robot for experiment (joints, robot position, ball position, simulator time)
        
        Keyword Arguments:
            height {float} -- height of the reset (m) (default: {0.55})
            orientation {str} -- orientation (straight, front or back) of the robot (default: {'straight'})
        """
        self.lines = []
        self.t = 0
        self.start = time.time()

        # Resets the robot position
        orn = [0, 0, 0]
        if orientation == 'front':
            orn = [0, math.pi/2, 0]
        elif orientation == 'back':
            orn = [0, -math.pi/2, 0]
        self.resetPose([0, 0, height], p.getQuaternionFromEuler(orn))

        # Reset the joints to 0
        for entry in self.joints.values():
            p.resetJointState(self.robot, entry, 0) 

def _reset(self):
    #print("KukaGymEnv _reset")
    self.terminated = 0
    p.resetSimulation()
    p.setPhysicsEngineParameter(numSolverIterations=150)
    p.setTimeStep(self._timeStep)
    p.loadURDF(os.path.join(self._urdfRoot,"plane.urdf"),[0,0,-1])

    p.loadURDF(os.path.join(self._urdfRoot,"table/table.urdf"), 0.5000000,0.00000,-.820000,0.000000,0.000000,0.0,1.0)

    xpos = 0.55 +0.12*random.random()
    ypos = 0 +0.2*random.random()
    ang = 3.14*0.5+3.1415925438*random.random()
    orn = p.getQuaternionFromEuler([0,0,ang])
    self.blockUid =p.loadURDF(os.path.join(self._urdfRoot,"block.urdf"), xpos,ypos,-0.15,orn[0],orn[1],orn[2],orn[3])

    p.setGravity(0,0,-10)
    self._kuka = kuka.Kuka(urdfRootPath=self._urdfRoot, timeStep=self._timeStep)
    self._envStepCounter = 0
    p.stepSimulation()
    self._observation = self.getExtendedObservation()
    return np.array(self._observation) 

def __init__(self, robot_name, print_debug, email_cred_file='', log_file='', control_rate=100, gripper_attached='default'):
        super(WidowXController, self).__init__(robot_name, print_debug, email_cred_file, log_file, control_rate, gripper_attached)
        self._redist_rate = rospy.Rate(50)

        self._arbotix = ArbotiX('/dev/ttyUSB0')
        assert self._arbotix.syncWrite(MAX_TORQUE_L, [[servo_id, 255, 0] for servo_id in SERVO_IDS]) != -1, "failure during servo configuring"
        assert self._arbotix.syncWrite(TORQUE_LIMIT, [[servo_id, 255, 0] for servo_id in SERVO_IDS]) != -1, "failure during servo configuring"

        self._joint_lock = Lock()
        self._angles, self._velocities = {}, {}
        rospy.Subscriber("/joint_states", JointState, self._joint_callback)
        time.sleep(1)        

        self._joint_pubs = [rospy.Publisher('/{}/command'.format(name), Float64, queue_size=1) for name in JOINT_NAMES[:-1]]
        self._gripper_pub = rospy.Publisher('/gripper_prismatic_joint/command', Float64, queue_size=1)

        p.connect(p.DIRECT)
        widow_x_urdf = '/'.join(__file__.split('/')[:-1]) + '/widowx/widowx.urdf'
        self._armID = p.loadURDF(widow_x_urdf, useFixedBase=True) 
        p.resetBasePositionAndOrientation(self._armID, [0, 0, 0], p.getQuaternionFromEuler([np.pi, np.pi, np.pi]))       
        self._n_errors = 0 

def _move_duck(self, pose):
    x, y, angle = pose
    orn = pybullet.getQuaternionFromEuler([0, 0, angle])
    pybullet.resetBasePositionAndOrientation(
        self._target_uid,
        [x, y, 0.],
        [orn[0], orn[1], orn[2], orn[3]],
        physicsClientId=self.cid) 

def get_extended_observation(self):
        self._observation = []
        observation_lim = []

        # ----------------------------------- #
        # --- Robot and world observation --- #
        # ----------------------------------- #
        robot_observation, robot_obs_lim = self._robot.get_observation()
        world_observation, world_obs_lim = self._world.get_observation()

        self._observation.extend(list(robot_observation))
        self._observation.extend(list(world_observation))
        observation_lim.extend(robot_obs_lim)
        observation_lim.extend(world_obs_lim)

        # ----------------------------------------- #
        # --- Object pose wrt hand c.o.m. frame --- #
        # ----------------------------------------- #
        inv_hand_pos, inv_hand_orn = p.invertTransform(robot_observation[:3],
                                                       p.getQuaternionFromEuler(robot_observation[3:6]))

        obj_pos_in_hand, obj_orn_in_hand = p.multiplyTransforms(inv_hand_pos, inv_hand_orn, world_observation[:3],
                                                                p.getQuaternionFromEuler(world_observation[3:6]))

        obj_euler_in_hand = p.getEulerFromQuaternion(obj_orn_in_hand)

        self._observation.extend(list(obj_pos_in_hand))
        self._observation.extend(list(obj_euler_in_hand))
        observation_lim.extend([[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]])
        observation_lim.extend([[0, 2 * m.pi], [0, 2 * m.pi], [0, 2 * m.pi]])

        return np.array(self._observation), observation_lim 

def get_extended_observation(self):
        self._observation = []
        observation_lim = []

        # ----------------------------------- #
        # --- Robot and world observation --- #
        # ----------------------------------- #
        robot_observation, robot_obs_lim = self._robot.get_observation()
        world_observation, world_obs_lim = self._world.get_observation()

        self._observation.extend(list(robot_observation))
        self._observation.extend(list(world_observation))
        observation_lim.extend(robot_obs_lim)
        observation_lim.extend(world_obs_lim)

        # ----------------------------------------- #
        # --- Object pose wrt hand c.o.m. frame --- #
        # ----------------------------------------- #
        inv_hand_pos, inv_hand_orn = p.invertTransform(robot_observation[:3],
                                                       p.getQuaternionFromEuler(robot_observation[3:6]))

        obj_pos_in_hand, obj_orn_in_hand = p.multiplyTransforms(inv_hand_pos, inv_hand_orn, world_observation[:3],
                                                                p.getQuaternionFromEuler(world_observation[3:6]))

        obj_euler_in_hand = p.getEulerFromQuaternion(obj_orn_in_hand)

        self._observation.extend(list(obj_pos_in_hand))
        self._observation.extend(list(obj_euler_in_hand))
        observation_lim.extend([[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]])
        observation_lim.extend([[0, 2*m.pi], [0, 2*m.pi], [0, 2*m.pi]])

        # ------------------- #
        # --- Target pose --- #
        # ------------------- #
        self._observation.extend(self._target_pose)
        observation_lim.extend(world_obs_lim[:3])

        return np.array(self._observation), observation_lim 

def _sample_pose(self):
        # ws_lim = self._ws_lim
        x_min = self._ws_lim[0][0] + 0.05
        x_max = self._ws_lim[0][1] - 0.1
        y_min = self._ws_lim[1][0] + 0.05
        y_max = self._ws_lim[1][1] - 0.05

        px = x_min + 0.5 * (x_max - x_min)
        py = y_min + 0.5 * (y_max - y_min)
        pz = self._h_table + 0.07

        quat = p.getQuaternionFromEuler([0.0, 0.0, 1/4*m.pi])

        if self._obj_pose_rnd_std > 0:
            # Add a Gaussian noise to position
            mu, sigma = 0, self._obj_pose_rnd_std
            # noise = self.np_random.normal(mu, sigma, 2)
            noise = [self.np_random.uniform(low=-self._obj_pose_rnd_std, high=self._obj_pose_rnd_std),
                     self.np_random.uniform(low=-self._obj_pose_rnd_std, high=self._obj_pose_rnd_std)]

            px = px + noise[0]
            py = py + noise[1]

            # Add uniform noise to yaw orientation
            quat = p.getQuaternionFromEuler([0, 0, self.np_random.uniform(low=-m.pi/4, high=m.pi/4)])

        px = np.clip(px, x_min, x_max)
        py = np.clip(py, y_min, y_max)

        obj_pose = (px, py, pz) + quat

        return obj_pose 

def get_extended_observation(self):
        self._observation = []
        observation_lim = []

        # ----------------------------------- #
        # --- Robot and world observation --- #
        # ----------------------------------- #
        robot_observation, robot_obs_lim = self._robot.get_observation()
        world_observation, world_obs_lim = self._world.get_observation()

        self._observation.extend(list(robot_observation))
        self._observation.extend(list(world_observation))
        observation_lim.extend(robot_obs_lim)
        observation_lim.extend(world_obs_lim)

        # ----------------------------------------- #
        # --- Object pose wrt hand c.o.m. frame --- #
        # ----------------------------------------- #
        inv_hand_pos, inv_hand_orn = p.invertTransform(robot_observation[:3],
                                                       p.getQuaternionFromEuler(robot_observation[3:6]))

        obj_pos_in_hand, obj_orn_in_hand = p.multiplyTransforms(inv_hand_pos, inv_hand_orn, world_observation[:3],
                                                                p.getQuaternionFromEuler(world_observation[3:6]))

        obj_euler_in_hand = p.getEulerFromQuaternion(obj_orn_in_hand)

        self._observation.extend(list(obj_pos_in_hand))
        self._observation.extend(list(obj_euler_in_hand))
        observation_lim.extend([[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]])
        observation_lim.extend([[0, 2*m.pi], [0, 2*m.pi], [0, 2*m.pi]])

        # ------------------- #
        # --- Target pose --- #
        # ------------------- #
        self._observation.extend(self._tg_pose)
        observation_lim.extend(world_obs_lim[:3])

        return np.array(self._observation), observation_lim 

def get_extended_observation(self):
        self._observation = []
        observation_lim = []

        # ----------------------------------- #
        # --- Robot and world observation --- #
        # ----------------------------------- #
        robot_observation, robot_obs_lim = self._robot.get_observation()
        world_observation, world_obs_lim = self._world.get_observation()

        self._observation.extend(list(robot_observation))
        self._observation.extend(list(world_observation))
        observation_lim.extend(robot_obs_lim)
        observation_lim.extend(world_obs_lim)

        # ----------------------------------------- #
        # --- Object pose wrt hand c.o.m. frame --- #
        # ----------------------------------------- #
        inv_hand_pos, inv_hand_orn = p.invertTransform(robot_observation[:3],
                                                       p.getQuaternionFromEuler(robot_observation[3:6]))
        obj_pos_in_hand, obj_orn_in_hand = p.multiplyTransforms(inv_hand_pos, inv_hand_orn, world_observation[:3],
                                                                p.getQuaternionFromEuler(world_observation[3:6]))
        obj_euler_in_hand = p.getEulerFromQuaternion(obj_orn_in_hand)

        self._observation.extend(list(obj_pos_in_hand))
        self._observation.extend(list(obj_euler_in_hand))
        observation_lim.extend([[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]])
        observation_lim.extend([[0, 2 * m.pi], [0, 2 * m.pi], [0, 2 * m.pi]])

        return np.array(self._observation), observation_lim 

def _updateConstraint(self):
        """
        INTERNAL METHOD, update the robot's constraint.
        """
        # Change the constraint to the requested position and orientation
        pybullet.changeConstraint(
            self.motion_constraint,
            self.pose_goal["position"],
            jointChildFrameOrientation=pybullet.getQuaternionFromEuler(
                self.pose_goal["orientation"]),
            maxForce=self.force,
            physicsClientId=self.physics_client) 

def _broadcastOdometry(self, odometry_publisher):
        """
        INTERNAL METHOD, computes an odometry message based on the robot's
        position, and broadcast it

        Parameters:
            odometry_publisher - The ROS publisher for the odometry message
        """
        # Send Transform odom
        x, y, theta = self.robot.getPosition()
        odom_trans = TransformStamped()
        odom_trans.header.frame_id = "odom"
        odom_trans.child_frame_id = "base_link"
        odom_trans.header.stamp = rospy.get_rostime()
        odom_trans.transform.translation.x = x
        odom_trans.transform.translation.y = y
        odom_trans.transform.translation.z = 0
        quaternion = pybullet.getQuaternionFromEuler([0, 0, theta])
        odom_trans.transform.rotation.x = quaternion[0]
        odom_trans.transform.rotation.y = quaternion[1]
        odom_trans.transform.rotation.z = quaternion[2]
        odom_trans.transform.rotation.w = quaternion[3]
        self.transform_broadcaster.sendTransform(odom_trans)
        # Set up the odometry
        odom = Odometry()
        odom.header.stamp = rospy.get_rostime()
        odom.header.frame_id = "odom"
        odom.pose.pose.position.x = x
        odom.pose.pose.position.y = y
        odom.pose.pose.position.z = 0.0
        odom.pose.pose.orientation = odom_trans.transform.rotation
        odom.child_frame_id = "base_link"
        [vx, vy, vz], [wx, wy, wz] = pybullet.getBaseVelocity(
            self.robot.getRobotModel(),
            self.robot.getPhysicsClientId())
        odom.twist.twist.linear.x = vx
        odom.twist.twist.linear.y = vy
        odom.twist.twist.angular.z = wz
        odometry_publisher.publish(odom) 

def apply_action(self, a):
        yaw= a
        pos, ornQuaternion = p.getBasePositionAndOrientation(self.robot_ids[0])
        ornEuler = p.getEulerFromQuaternion(ornQuaternion)
        currAngle = ornEuler[2]-self.config['yaw_constant']*yaw
        print(pos)
        print(ornEuler)
        pos = [pos[0]+ self.config['vx_constant']*self.vx*math.cos(currAngle)+random.uniform(self.config['wind_limits'][0], self.config['wind_limits'][1]), 
        pos[1]+self.config['vx_constant']*self.vx*math.sin(currAngle)+random.uniform(self.config['wind_limits'][0], self.config['wind_limits'][1]), self.height]
        pitch = min(max( ornEuler[0] + random.uniform(self.config['d_roll_per_step'][0], self.config['d_roll_per_step'][1]), self.config['roll_limits'][0]), self.config['roll_limits'][1])
        ornEuler = [pitch, ornEuler[1], currAngle]
        ornQuaternion = p.getQuaternionFromEuler(ornEuler)
        self.robot_body.set_pose(pos, ornQuaternion)
        return pos[0], pos[1], ornEuler[2], self.height, self.vx 

def _setup(self):
        '''
        Create random objects at random positions. Load random objects from the
        scene and create them in different places. In the future we may want to
        switch from using the list of "all" objects to a subset that we can
        actually pick up and manipulate.
        '''

        rospack = rospkg.RosPack()
        path = rospack.get_path('costar_objects')
        sdf_dir = os.path.join(path, self.sdf_dir)
        objs = [obj for obj in os.listdir(
            sdf_dir) if os.path.isdir(os.path.join(sdf_dir, obj))]

        randn = np.random.randint(1, len(objs))

        objs_name_to_add = np.random.choice(objs, randn)
        objs_to_add = [os.path.join(sdf_dir, obj, self.model_file_name)
                       for obj in objs_name_to_add]

        identity_orientation = pb.getQuaternionFromEuler([0, 0, 0])
        # load sdfs for all objects and initialize positions
        for obj_index, obj in enumerate(objs_to_add):
            if objs_name_to_add[obj_index] in self.models:
                try:
                    print 'Loading object: ', obj
                    obj_id_list = pb.loadSDF(obj)
                    for obj_id in obj_id_list:
                        self.objs.append(obj_id)
                        random_position = np.random.rand(
                            3) * self.spawn_pos_delta + self.spawn_pos_min
                        pb.resetBasePositionAndOrientation(
                            obj_id, random_position, identity_orientation)
                except Exception, e:
                    print e 

def _randomly_place_objects(self, urdfList):
    """Randomly places the objects in the bin.

    Args:
      urdfList: The list of urdf files to place in the bin.

    Returns:
      The list of object unique ID's.
    """


    # Randomize positions of each object urdf.
    objectUids = []
    for urdf_name in urdfList:
      xpos = 0.4 +self._blockRandom*random.random()
      ypos = self._blockRandom*(random.random()-.5)
      angle = np.pi/2 + self._blockRandom * np.pi * random.random()
      orn = p.getQuaternionFromEuler([0, 0, angle])
      urdf_path = os.path.join(self._urdfRoot, urdf_name)
      uid = p.loadURDF(urdf_path, [xpos, ypos, .15],
        [orn[0], orn[1], orn[2], orn[3]])
      objectUids.append(uid)
      # Let each object fall to the tray individual, to prevent object
      # intersection.
      for _ in range(500):
        p.stepSimulation()
    return objectUids 

def _setup(self):
        '''
        Create task by adding objects to the scene
        '''

        rospack = rospkg.RosPack()
        path = rospack.get_path('costar_objects')
        urdf_dir = os.path.join(path, self.urdf_dir)
        sdf_dir = os.path.join(path, self.sdf_dir)
        objs = [obj for obj in os.listdir(
            sdf_dir) if os.path.isdir(os.path.join(sdf_dir, obj))]

        randn = np.random.randint(1, len(objs))

        objs_name_to_add = np.random.choice(objs, randn)
        objs_to_add = [os.path.join(sdf_dir, obj, self.model_file_name)
                       for obj in objs_name_to_add]

        identity_orientation = pb.getQuaternionFromEuler([0, 0, 0])
        # load sdfs for all objects and initialize positions
        for obj_index, obj in enumerate(objs_to_add):
            if objs_name_to_add[obj_index] in self.models:
                try:
                    print 'Loading object: ', obj
                    obj_id_list = pb.loadSDF(obj)
                    for obj_id in obj_id_list:
                        self.objs.append(obj_id)
                        random_position = np.random.rand(
                            3) * self.spawn_pos_delta + self.spawn_pos_min
                        pb.resetBasePositionAndOrientation(
                            obj_id, random_position, identity_orientation)
                except Exception, e:
                    print e 

def display_cup_points(self):
        sphere_collision = -1
        sphere_visual = p.createVisualShape(shapeType=p.GEOM_SPHERE, radius=0.01, rgbaColor=[0, 1, 1, 1], physicsClientId=self.id)
        cup_pos, cup_orient = p.getBasePositionAndOrientation(self.cup, physicsClientId=self.id)
        cup_pos, cup_orient = p.multiplyTransforms(cup_pos, cup_orient, [0, 0.06, 0], p.getQuaternionFromEuler([np.pi/2.0, 0, 0], physicsClientId=self.id), physicsClientId=self.id)
        top_center_pos, _ = p.multiplyTransforms(cup_pos, cup_orient, self.cup_top_center_offset, [0, 0, 0, 1], physicsClientId=self.id)
        bottom_center_pos, _ = p.multiplyTransforms(cup_pos, cup_orient, self.cup_bottom_center_offset, [0, 0, 0, 1], physicsClientId=self.id)
        self.cup_top_center = p.createMultiBody(baseMass=0.0, baseCollisionShapeIndex=sphere_collision, baseVisualShapeIndex=sphere_visual, basePosition=top_center_pos, useMaximalCoordinates=False, physicsClientId=self.id)
        self.cup_bottom_center = p.createMultiBody(baseMass=0.0, baseCollisionShapeIndex=sphere_collision, baseVisualShapeIndex=sphere_visual, basePosition=bottom_center_pos, useMaximalCoordinates=False, physicsClientId=self.id)

        cylinder_collision = -1
        cylinder_visual = p.createVisualShape(shapeType=p.GEOM_CYLINDER, radius=0.05, length=0.14, rgbaColor=[0, 1, 1, 0.25], physicsClientId=self.id)
        self.cup_cylinder = p.createMultiBody(baseMass=0.0, baseCollisionShapeIndex=cylinder_collision, baseVisualShapeIndex=cylinder_visual, basePosition=cup_pos, baseOrientation=cup_orient, useMaximalCoordinates=False, physicsClientId=self.id) 

def step(self, action):
        self.take_step(action, robot_arm='right', gains=self.config('robot_gains'), forces=self.config('robot_forces'), human_gains=0.0005)

        robot_force_on_human, cup_force_on_human = self.get_total_force()
        total_force_on_human = robot_force_on_human + cup_force_on_human
        reward_water, water_mouth_velocities, water_hit_human_reward = self.get_water_rewards()
        end_effector_velocity = np.linalg.norm(p.getBaseVelocity(self.cup, physicsClientId=self.id)[0])
        obs = self._get_obs([cup_force_on_human], [robot_force_on_human, cup_force_on_human])

        # Get human preferences
        preferences_score = self.human_preferences(end_effector_velocity=end_effector_velocity, total_force_on_human=robot_force_on_human, tool_force_at_target=cup_force_on_human, food_hit_human_reward=water_hit_human_reward, food_mouth_velocities=water_mouth_velocities)

        cup_pos, cup_orient = p.getBasePositionAndOrientation(self.cup, physicsClientId=self.id)
        cup_pos, cup_orient = p.multiplyTransforms(cup_pos, cup_orient, [0, 0.06, 0], p.getQuaternionFromEuler([np.pi/2.0, 0, 0], physicsClientId=self.id), physicsClientId=self.id)
        cup_top_center_pos, _ = p.multiplyTransforms(cup_pos, cup_orient, self.cup_top_center_offset, [0, 0, 0, 1], physicsClientId=self.id)
        reward_distance = -np.linalg.norm(self.target_pos - np.array(cup_top_center_pos)) # Penalize distances between top of cup and mouth
        reward_action = -np.sum(np.square(action)) # Penalize actions
        # Encourage robot to have a tilted end effector / cup
        cup_euler = p.getEulerFromQuaternion(cup_orient, physicsClientId=self.id)
        reward_tilt = -abs(cup_euler[0] + np.pi/2) if self.robot_type == 'jaco' else -abs(cup_euler[0] - np.pi/2)

        reward = self.config('distance_weight')*reward_distance + self.config('action_weight')*reward_action + self.config('cup_tilt_weight')*reward_tilt + self.config('drinking_reward_weight')*reward_water + preferences_score

        if self.gui and reward_water != 0:
            print('Task success:', self.task_success, 'Water reward:', reward_water)

        info = {'total_force_on_human': total_force_on_human, 'task_success': int(self.task_success >= self.total_water_count*self.config('task_success_threshold')), 'action_robot_len': self.action_robot_len, 'action_human_len': self.action_human_len, 'obs_robot_len': self.obs_robot_len, 'obs_human_len': self.obs_human_len}
        done = False

        return obs, reward, done, info 

def set_cstr_dof(cstr, pos, euler, max_force):
        pos = list(pos)
        euler = list(euler)
        quat = p.getQuaternionFromEuler(euler)
        p.changeConstraint(
                userConstraintUniqueId=cstr,
                jointChildPivot=pos,
                jointChildFrameOrientation=quat,
                maxForce=max_force) 

def pos_in_frame(pos, frame):
        frame_xyz = frame[0]
        frame_rpy = frame[1]
        quat = p.getQuaternionFromEuler(frame_rpy)
        mat33 = p.getMatrixFromQuaternion(quat)
        mat33 = np.reshape(mat33, [3, 3])
        pos_in_frame = frame_xyz + np.dot(pos, mat33.T)
        return pos_in_frame 

def mat33_from_euler(euler):
        euler = list(euler)
        quat = p.getQuaternionFromEuler(euler)
        mat33 = p.getMatrixFromQuaternion(quat)
        return np.reshape(mat33, [3, 3]) 

def quat_from_euler(euler):
        euler = list(euler)
        quat = p.getQuaternionFromEuler(euler)
        return np.array(quat, dtype=np.float32) 

def __init__(self, path, position, orientation):
        self.orientation = orientation
        self.position = position
        self.path = path
        self.plane = p.loadURDF(self.path, basePosition=self.position,
                                baseOrientation=p.getQuaternionFromEuler(self.orientation)) 

def setBallPos(self, x, y):
        """Sets the ball position on the field"""
        if self.ball is not None:
            # Putting the ball on the ground at given position
            p.resetBasePositionAndOrientation(
                self.ball, [x, y, 0.06], p.getQuaternionFromEuler([0, 0, 0]))

            # Stopping the ball speed
            p.changeDynamics(self.ball, 0,
                             linearDamping=0, angularDamping=0.1) 

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 joinUrdf(self, childEditor, parentLinkIndex=0, jointPivotXYZInParent=[0,0,0], jointPivotRPYInParent=[0,0,0], jointPivotXYZInChild=[0,0,0], jointPivotRPYInChild=[0,0,0], parentPhysicsClientId=0, childPhysicsClientId=0):

		childLinkIndex = len(self.urdfLinks)
		insertJointIndex = len(self.urdfJoints)
		
		#combine all links, and add a joint

		for link in childEditor.urdfLinks:
			self.linkNameToIndex[link.link_name]=len(self.urdfLinks)
			self.urdfLinks.append(link)
		for joint in childEditor.urdfJoints:
			self.urdfJoints.append(joint)
		#add a new joint between a particular


		
		
		jointPivotQuatInChild = p.getQuaternionFromEuler(jointPivotRPYInChild)
		invJointPivotXYZInChild, invJointPivotQuatInChild = p.invertTransform(jointPivotXYZInChild,jointPivotQuatInChild)

		#apply this invJointPivot***InChild to all inertial, visual and collision element in the child link
		#inertial
		pos, orn = p.multiplyTransforms(self.urdfLinks[childLinkIndex].urdf_inertial.origin_xyz,p.getQuaternionFromEuler(self.urdfLinks[childLinkIndex].urdf_inertial.origin_rpy),invJointPivotXYZInChild,invJointPivotQuatInChild, physicsClientId=parentPhysicsClientId)
		self.urdfLinks[childLinkIndex].urdf_inertial.origin_xyz = pos
		self.urdfLinks[childLinkIndex].urdf_inertial.origin_rpy = p.getEulerFromQuaternion(orn)
		#all visual
		for v in self.urdfLinks[childLinkIndex].urdf_visual_shapes:
			pos, orn = p.multiplyTransforms(v.origin_xyz,p.getQuaternionFromEuler(v.origin_rpy),invJointPivotXYZInChild,invJointPivotQuatInChild)
			v.origin_xyz = pos
			v.origin_rpy = p.getEulerFromQuaternion(orn)
		#all collision
		for c in self.urdfLinks[childLinkIndex].urdf_collision_shapes:
			pos, orn = p.multiplyTransforms(c.origin_xyz,p.getQuaternionFromEuler(c.origin_rpy),invJointPivotXYZInChild,invJointPivotQuatInChild)
			c.origin_xyz = pos
			c.origin_rpy = p.getEulerFromQuaternion(orn)


		childLink = self.urdfLinks[childLinkIndex]
		parentLink = self.urdfLinks[parentLinkIndex]

		joint = UrdfJoint()
		joint.link = childLink
		joint.joint_name = "joint_dummy1"
		joint.joint_type = p.JOINT_REVOLUTE
		joint.joint_lower_limit = 0
		joint.joint_upper_limit = -1
		joint.parent_name = parentLink.link_name
		joint.child_name = childLink.link_name
		joint.joint_origin_xyz = jointPivotXYZInParent
		joint.joint_origin_rpy = jointPivotRPYInParent
		joint.joint_axis_xyz = [0,0,1]

		#the following commented line would crash PyBullet, it messes up the joint indexing/ordering
		#self.urdfJoints.append(joint)

		#so make sure to insert the joint in the right place, to links/joints match
		self.urdfJoints.insert(insertJointIndex,joint)
		return joint 

def get_water_rewards(self):
        # Check all water particles to see if they have entered the person's mouth or have left the scene
        # Delete such particles and give the robot a reward or penalty depending on particle status
        cup_pos, cup_orient = p.getBasePositionAndOrientation(self.cup, physicsClientId=self.id)
        cup_pos, cup_orient = p.multiplyTransforms(cup_pos, cup_orient, [0, 0.06, 0], p.getQuaternionFromEuler([np.pi/2.0, 0, 0], physicsClientId=self.id), physicsClientId=self.id)
        top_center_pos, _ = p.multiplyTransforms(cup_pos, cup_orient, self.cup_top_center_offset, [0, 0, 0, 1], physicsClientId=self.id)
        bottom_center_pos, _ = p.multiplyTransforms(cup_pos, cup_orient, self.cup_bottom_center_offset, [0, 0, 0, 1], physicsClientId=self.id)
        top_center_pos = np.array(top_center_pos)
        bottom_center_pos = np.array(bottom_center_pos)
        if self.cup_top_center is not None:
            p.resetBasePositionAndOrientation(self.cup_top_center, top_center_pos, [0, 0, 0, 1], physicsClientId=self.id)
            p.resetBasePositionAndOrientation(self.cup_bottom_center, bottom_center_pos, [0, 0, 0, 1], physicsClientId=self.id)
            p.resetBasePositionAndOrientation(self.cup_cylinder, cup_pos, cup_orient, physicsClientId=self.id)
        water_reward = 0
        water_hit_human_reward = 0
        water_mouth_velocities = []
        waters_to_remove = []
        for w in self.waters:
            water_pos, water_orient = p.getBasePositionAndOrientation(w, physicsClientId=self.id)
            if not self.util.points_in_cylinder(top_center_pos, bottom_center_pos, 0.05, np.array(water_pos)):
                distance_to_mouth = np.linalg.norm(self.target_pos - water_pos)
                if distance_to_mouth < 0.03: # hard
                # if distance_to_mouth < 0.05: # easy
                    # Delete particle and give robot a reward
                    water_reward += 10
                    self.task_success += 1
                    p.resetBasePositionAndOrientation(w, self.np_random.uniform(1000, 2000, size=3), [0, 0, 0, 1], physicsClientId=self.id)
                    water_velocity = np.linalg.norm(p.getBaseVelocity(w, physicsClientId=self.id)[0])
                    water_mouth_velocities.append(water_velocity)
                    waters_to_remove.append(w)
                    continue
                elif water_pos[-1] < 0.5:
                    # Delete particle and give robot a penalty for spilling water
                    water_reward -= 1
                    waters_to_remove.append(w)
                    continue
                if len(p.getContactPoints(bodyA=w, bodyB=self.human, physicsClientId=self.id)) > 0:
                    # Record that this water particle just hit the person, so that we can penalize the robot
                    waters_to_remove.append(w)
                    water_hit_human_reward -= 1
        self.waters = [w for w in self.waters if w not in waters_to_remove]
        return water_reward, water_mouth_velocities, water_hit_human_reward