Skip to content

Commit

Permalink
Run isort & black for all python files
Browse files Browse the repository at this point in the history
  • Loading branch information
@KolinGuo
KolinGuo committed Jan 12, 2024
1 parent a15c156 commit 3333cf2
Show file tree
Hide file tree
Showing 10 changed files with 1,037 additions and 678 deletions.
535 changes: 353 additions & 182 deletions dev/stubgen.py
View file

Large diffs are not rendered by default.

43 changes: 25 additions & 18 deletions mplib/examples/collision_avoidance.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
import sapien.core as sapien

from .demo_setup import DemoSetup


class PlanningDemo(DemoSetup):
""" The shows the planner's ability to generate a collision free path with the straight path causes collisions """
"""The shows the planner's ability to generate a collision free path with the straight path causes collisions"""

def __init__(self):
"""
Same setup as demo.py except the boxes are of difference sizes and different use
Expand All @@ -23,75 +26,79 @@ def __init__(self):
builder = self.scene.create_actor_builder()
builder.add_box_collision(half_size=[0.4, 0.4, 0.025])
builder.add_box_visual(half_size=[0.4, 0.4, 0.025])
self.table = builder.build_kinematic(name='table')
self.table.set_pose(sapien.Pose([0.56, 0, - 0.025]))
self.table = builder.build_kinematic(name="table")
self.table.set_pose(sapien.Pose([0.56, 0, -0.025]))

# red box is the target we want to grab
builder = self.scene.create_actor_builder()
builder.add_box_collision(half_size=[0.02, 0.02, 0.06])
builder.add_box_visual(half_size=[0.02, 0.02, 0.06], color=[1, 0, 0])
self.red_cube = builder.build(name='red_cube')
self.red_cube = builder.build(name="red_cube")
self.red_cube.set_pose(sapien.Pose([0.7, 0, 0.06]))

# green box is the landing pad on which we want to place the red box
builder = self.scene.create_actor_builder()
builder.add_box_collision(half_size=[0.04, 0.04, 0.005])
builder.add_box_visual(half_size=[0.04, 0.04, 0.005], color=[0, 1, 0])
self.green_cube = builder.build(name='green_cube')
self.green_cube = builder.build(name="green_cube")
self.green_cube.set_pose(sapien.Pose([0.4, 0.3, 0.005]))

# blue box is the obstacle we want to avoid
builder = self.scene.create_actor_builder()
builder.add_box_collision(half_size=[0.05, 0.2, 0.1])
builder.add_box_visual(half_size=[0.05, 0.2, 0.1], color=[0, 0, 1])
self.blue_cube = builder.build(name='blue_cube')
self.blue_cube = builder.build(name="blue_cube")
self.blue_cube.set_pose(sapien.Pose([0.55, 0, 0.1]))

def add_point_cloud(self):
""" we tell the planner about the obstacle through a point cloud """
"""we tell the planner about the obstacle through a point cloud"""
import trimesh

box = trimesh.creation.box([0.1, 0.4, 0.2])
points, _ = trimesh.sample.sample_surface(box, 1000)
points += [0.55, 0, 0.1]
self.planner.update_point_cloud(points, radius=0.02)
return
return

def demo(self, with_screw=True, use_point_cloud=True, use_attach=True):
"""
We pick up the red box while avoiding the blue box and place it back down on top of the green box
"""
pickup_pose = [0.7, 0, 0.12, 0, 1, 0, 0]
delivery_pose = [0.4, 0.3, 0.13, 0, 1, 0, 0]

# tell the planner where the obstacle is
if use_point_cloud:
self.add_point_cloud()

# move to the pickup pose
pickup_pose[2] += 0.2
# no need to check collision against attached object since nothing picked up yet
self.move_to_pose(pickup_pose, with_screw, use_point_cloud, use_attach=False)
self.move_to_pose(pickup_pose, with_screw, use_point_cloud, use_attach=False)
self.open_gripper()
pickup_pose[2] -= 0.12
# no attach since nothing picked up yet
self.move_to_pose(pickup_pose, with_screw, use_point_cloud, use_attach=False)
self.move_to_pose(pickup_pose, with_screw, use_point_cloud, use_attach=False)
self.close_gripper()

if use_attach:
self.planner.update_attached_box([0.04, 0.04, 0.12], [0, 0, 0.14, 1, 0, 0, 0])
self.planner.update_attached_box(
[0.04, 0.04, 0.12], [0, 0, 0.14, 1, 0, 0, 0]
)

# move to the delivery pose
pickup_pose[2] += 0.12
self.move_to_pose(pickup_pose, with_screw, use_point_cloud, use_attach)
self.move_to_pose(pickup_pose, with_screw, use_point_cloud, use_attach)
delivery_pose[2] += 0.2
self.move_to_pose(delivery_pose, with_screw, use_point_cloud, use_attach)
self.move_to_pose(delivery_pose, with_screw, use_point_cloud, use_attach)
delivery_pose[2] -= 0.12
self.move_to_pose(delivery_pose, with_screw, use_point_cloud, use_attach)
self.move_to_pose(delivery_pose, with_screw, use_point_cloud, use_attach)
self.open_gripper()
delivery_pose[2] += 0.12
self.move_to_pose(delivery_pose, with_screw, use_point_cloud, use_attach=False)
self.move_to_pose(delivery_pose, with_screw, use_point_cloud, use_attach=False)


if __name__ == '__main__':
if __name__ == "__main__":
"""
As you change some of the parameters, you will see different behaviors
In particular, when point cloud is not used, the robot will attemt to go through the blue box
Expand Down
233 changes: 128 additions & 105 deletions mplib/examples/constrained_planning.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,114 +2,137 @@

import numpy as np
import transforms3d

from .demo_setup import DemoSetup


class ConstrainedPlanningDemo(DemoSetup):
"""
This demo shows the planner's ability to plan with constraints.
For this particular demo, we move to several poses while pointing the end effector roughly 15 degrees w.r.t. -z axis
"""
def __init__(self):
""" set up the scene and load the robot """
super().__init__()
self.setup_scene()
self.load_robot()
self.setup_planner()

def add_point_cloud(self):
""" add some random obstacles to make the planning more challenging """
import trimesh
box = trimesh.creation.box([0.1, 0.4, 0.2])
points, _ = trimesh.sample.sample_surface(box, 1000)
all_pts = np.concatenate([points+[-0.65, -0.1, 0.4], points+[0.55, 0, 0.1]], axis=0)
self.planner.update_point_cloud(all_pts, radius=0.02)
return

def get_eef_z(self):
""" helper function for constraint """
ee_idx = self.planner.link_name_2_idx[self.planner.move_group]
ee_pose = self.planner.robot.get_pinocchio_model().get_link_pose(ee_idx)
mat = transforms3d.quaternions.quat2mat(ee_pose[3:])
return mat[:,2]

def make_f(self):
"""
create a constraint function that takes in a qpos and outputs a scalar
A valid constraint function should evaluates to 0 when the constraint is satisfied
See [ompl constrained planning](https://ompl.kavrakilab.org/constrainedPlanning.html) for more details
This demo shows the planner's ability to plan with constraints.
For this particular demo, we move to several poses while pointing the end effector roughly 15 degrees w.r.t. -z axis
"""
def f(x, out):
self.planner.robot.set_qpos(x)
out[0] = self.get_eef_z().dot(np.array([0,0,-1]))-0.966 # maintain 15 degrees w.r.t. -z axis
return f

def make_j(self):
"""
create the jacobian of the constraint function w.r.t. qpos
This is needed because the planner uses the jacobian to project a random sample to the constraint manifold
"""
def j(x, out):
full_qpos = self.planner.pad_qpos(x)
jac = self.planner.robot.get_pinocchio_model().compute_single_link_jacobian(full_qpos, len(self.planner.move_group_joint_indices)-1)
rot_jac = jac[3:, self.planner.move_group_joint_indices]
for i in range(len(self.planner.move_group_joint_indices)):
out[i] = np.cross(rot_jac[:,i], self.get_eef_z()).dot(np.array([0,0,-1]))
return j

def demo(self):
"""
We first plan with constraints to three poses, then plan without constraints to the same poses
While not always the case, sometimes without constraints, the end effector will tilt almost upside down
"""
# this starting pose has the end effector tilted roughly 15 degrees
starting_qpos = [0, 0.19, 0.0, -2.61, 0.0, 2.88, 0.78, 0, 0]
self.robot.set_qpos(starting_qpos)
self.planner.robot.set_qpos(starting_qpos[:7])
# all these poses are constrain compatible (roughly 15 degrees w.r.t. -z axis)
poses = [[-0.4, -0.3, 0.28, 0.0704682, -0.5356872, 0.8342834, 0.1097478],
[0.6, 0.1, 0.44, 0.0704682, -0.5356872, -0.8342834, -0.1097478],
[0, -0.3, 0.5, 0.1304237, -0.9914583, 0, 0]]

# add some point cloud to make the planning more challenging so we can see the effect of no constraint
self.add_point_cloud()

# with constraint
print("with constraint. all movements roughly maintain 15 degrees w.r.t. -z axis")
for pose in poses:
result = self.planner.plan_qpos_to_pose(
pose,
self.robot.get_qpos(),
time_step=1/250,
use_point_cloud=True,
use_attach=False,
planner_name="RRTConnect",
constraint_function=self.make_f(),
constraint_jacobian=self.make_j(),
constraint_tolerance=0.05,
)
if result['status'] != "Success":
print(result['status'])
return -1
self.follow_path(result)

# without constraint
print("without constraint. certain movements can sometimes tilt the end effector almost upside down")
for pose in poses:
result = self.planner.plan_qpos_to_pose(
pose,
self.robot.get_qpos(),
time_step=1/250,
use_point_cloud=True,
use_attach=False,
planner_name="RRTConnect",
no_simplification=True
)
if result['status'] != "Success":
print(result['status'])
return -1
self.follow_path(result)


if __name__ == '__main__':
demo = ConstrainedPlanningDemo()
demo.demo()
def __init__(self):
"""set up the scene and load the robot"""
super().__init__()
self.setup_scene()
self.load_robot()
self.setup_planner()

def add_point_cloud(self):
"""add some random obstacles to make the planning more challenging"""
import trimesh

box = trimesh.creation.box([0.1, 0.4, 0.2])
points, _ = trimesh.sample.sample_surface(box, 1000)
all_pts = np.concatenate(
[points + [-0.65, -0.1, 0.4], points + [0.55, 0, 0.1]], axis=0
)
self.planner.update_point_cloud(all_pts, radius=0.02)
return

def get_eef_z(self):
"""helper function for constraint"""
ee_idx = self.planner.link_name_2_idx[self.planner.move_group]
ee_pose = self.planner.robot.get_pinocchio_model().get_link_pose(ee_idx)
mat = transforms3d.quaternions.quat2mat(ee_pose[3:])
return mat[:, 2]

def make_f(self):
"""
create a constraint function that takes in a qpos and outputs a scalar
A valid constraint function should evaluates to 0 when the constraint is satisfied
See [ompl constrained planning](https://ompl.kavrakilab.org/constrainedPlanning.html) for more details
"""

def f(x, out):
self.planner.robot.set_qpos(x)
out[0] = (
self.get_eef_z().dot(np.array([0, 0, -1])) - 0.966
) # maintain 15 degrees w.r.t. -z axis

return f

def make_j(self):
"""
create the jacobian of the constraint function w.r.t. qpos
This is needed because the planner uses the jacobian to project a random sample to the constraint manifold
"""

def j(x, out):
full_qpos = self.planner.pad_qpos(x)
jac = self.planner.robot.get_pinocchio_model().compute_single_link_jacobian(
full_qpos, len(self.planner.move_group_joint_indices) - 1
)
rot_jac = jac[3:, self.planner.move_group_joint_indices]
for i in range(len(self.planner.move_group_joint_indices)):
out[i] = np.cross(rot_jac[:, i], self.get_eef_z()).dot(
np.array([0, 0, -1])
)

return j

def demo(self):
"""
We first plan with constraints to three poses, then plan without constraints to the same poses
While not always the case, sometimes without constraints, the end effector will tilt almost upside down
"""
# this starting pose has the end effector tilted roughly 15 degrees
starting_qpos = [0, 0.19, 0.0, -2.61, 0.0, 2.88, 0.78, 0, 0]
self.robot.set_qpos(starting_qpos)
self.planner.robot.set_qpos(starting_qpos[:7])
# all these poses are constrain compatible (roughly 15 degrees w.r.t. -z axis)
poses = [
[-0.4, -0.3, 0.28, 0.0704682, -0.5356872, 0.8342834, 0.1097478],
[0.6, 0.1, 0.44, 0.0704682, -0.5356872, -0.8342834, -0.1097478],
[0, -0.3, 0.5, 0.1304237, -0.9914583, 0, 0],
]

# add some point cloud to make the planning more challenging so we can see the effect of no constraint
self.add_point_cloud()

# with constraint
print(
"with constraint. all movements roughly maintain 15 degrees w.r.t. -z axis"
)
for pose in poses:
result = self.planner.plan_qpos_to_pose(
pose,
self.robot.get_qpos(),
time_step=1 / 250,
use_point_cloud=True,
use_attach=False,
planner_name="RRTConnect",
constraint_function=self.make_f(),
constraint_jacobian=self.make_j(),
constraint_tolerance=0.05,
)
if result["status"] != "Success":
print(result["status"])
return -1
self.follow_path(result)

# without constraint
print(
"without constraint. certain movements can sometimes tilt the end effector"
" almost upside down"
)
for pose in poses:
result = self.planner.plan_qpos_to_pose(
pose,
self.robot.get_qpos(),
time_step=1 / 250,
use_point_cloud=True,
use_attach=False,
planner_name="RRTConnect",
no_simplification=True,
)
if result["status"] != "Success":
print(result["status"])
return -1
self.follow_path(result)


if __name__ == "__main__":
demo = ConstrainedPlanningDemo()
demo.demo()
Loading

0 comments on commit 3333cf2

Please sign in to comment.