constrained planning, robot segmentation

This commit is contained in:
Balakumar Sundaralingam
2024-02-22 21:45:47 -08:00
parent 88eac64edc
commit bafdf80c05
102 changed files with 12440 additions and 8112 deletions

View File

@@ -0,0 +1,355 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import torch
a = torch.zeros(4, device="cuda:0")
# Third Party
from omni.isaac.kit import SimulationApp
simulation_app = SimulationApp(
{
"headless": False,
"width": "1920",
"height": "1080",
}
)
# Third Party
import numpy as np
import torch
from helper import add_extensions, add_robot_to_scene
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType
from curobo.geom.types import Cuboid, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig
from curobo.types.state import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
from curobo.wrap.reacher.motion_gen import (
MotionGen,
MotionGenConfig,
MotionGenPlanConfig,
PoseCostMetric,
)
simulation_app.update()
# Standard Library
import argparse
# Third Party
import carb
from omni.isaac.core import World
from omni.isaac.core.materials import OmniGlass, OmniPBR
from omni.isaac.core.objects import cuboid, sphere
from omni.isaac.core.utils.types import ArticulationAction
# CuRobo
from curobo.util.usd_helper import UsdHelper
parser = argparse.ArgumentParser()
parser.add_argument("--robot", type=str, default="franka.yml", help="robot configuration to load")
args = parser.parse_args()
if __name__ == "__main__":
my_world = World(stage_units_in_meters=1.0)
stage = my_world.stage
n_obstacle_cuboids = 10
n_obstacle_mesh = 10
stage = my_world.stage
my_world.scene.add_default_ground_plane()
xform = stage.DefinePrim("/World", "Xform")
stage.SetDefaultPrim(xform)
target_material = OmniPBR("/World/looks/t", color=np.array([0, 1, 0]))
target_material_2 = OmniPBR("/World/looks/t2", color=np.array([0, 1, 0]))
target_material_plane = OmniGlass(
"/World/looks/t3", color=np.array([0, 1, 0]), ior=1.25, depth=0.001, thin_walled=False
)
target_material_line = OmniGlass(
"/World/looks/t4", color=np.array([0, 1, 0]), ior=1.25, depth=0.001, thin_walled=True
)
# target_orient = [0,0,0.707,0.707]
target_orient = [0.5, -0.5, 0.5, 0.5]
target = cuboid.VisualCuboid(
"/World/target_1",
position=np.array([0.55, -0.3, 0.5]),
orientation=np.array(target_orient),
size=0.04,
visual_material=target_material,
)
# Make a target to follow
target_2 = cuboid.VisualCuboid(
"/World/target_2",
position=np.array([0.55, 0.4, 0.5]),
orientation=np.array(target_orient),
size=0.04,
visual_material=target_material_2,
)
x_plane = cuboid.VisualCuboid(
"/World/constraint_plane",
position=np.array([0.55, 0.05, 0.5]),
orientation=np.array(target_orient),
scale=[1.1, 0.001, 1.0],
visual_material=target_material_plane,
)
xz_line = cuboid.VisualCuboid(
"/World/constraint_line",
position=np.array([0.55, 0.05, 0.5]),
orientation=np.array(target_orient),
scale=[0.04, 0.04, 0.65],
visual_material=target_material_line,
)
collision_checker_type = CollisionCheckerType.BLOX
tensor_args = TensorDeviceType()
robot_cfg = load_yaml(join_path(get_robot_configs_path(), args.robot))["robot_cfg"]
j_names = robot_cfg["kinematics"]["cspace"]["joint_names"]
default_config = robot_cfg["kinematics"]["cspace"]["retract_config"]
robot, robot_prim_path = add_robot_to_scene(robot_cfg, my_world, "/World/world_robot/")
world_cfg_table = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
)
world_cfg_table.cuboid[0].pose[2] -= 0.01
world_cfg1 = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_mesh_world()
world_cfg1.mesh[0].name += "_mesh"
world_cfg1.mesh[0].pose[2] = -10.5
world_cfg = WorldConfig(cuboid=world_cfg_table.cuboid, mesh=world_cfg1.mesh)
usd_help = UsdHelper()
usd_help.load_stage(my_world.stage)
usd_help.add_world_to_stage(world_cfg_table.get_mesh_world(), base_frame="/World")
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg,
world_cfg,
tensor_args,
collision_checker_type=CollisionCheckerType.MESH,
collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
velocity_scale=0.75,
interpolation_dt=0.02,
ee_link_name="right_gripper",
)
motion_gen = MotionGen(motion_gen_config)
print("warming up..")
motion_gen.warmup(warmup_js_trajopt=False)
world_model = motion_gen.world_collision
i = 0
tensor_args = TensorDeviceType()
target_list = [target, target_2]
target_material_list = [target_material, target_material_2]
for material in target_material_list:
material.set_color(np.array([0.1, 0.1, 0.1]))
target_material_plane.set_color(np.array([1, 1, 1]))
target_material_line.set_color(np.array([1, 1, 1]))
target_idx = 0
cmd_idx = 0
cmd_plan = None
articulation_controller = robot.get_articulation_controller()
plan_config = MotionGenPlanConfig(
enable_graph=False, enable_graph_attempt=4, max_attempts=2, enable_finetune_trajopt=True
)
plan_idx = 0
cmd_step_idx = 0
pose_cost_metric = None
while simulation_app.is_running():
my_world.step(render=True)
if not my_world.is_playing():
if i % 100 == 0:
print("**** Click Play to start simulation *****")
i += 1
# if step_index == 0:
# my_world.play()
continue
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
robot._articulation_view.set_max_efforts(
values=np.array([5000 for i in range(len(idx_list))]), joint_indices=idx_list
)
if False and step_index % 50 == 0.0: # No obstacle update
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
reference_prim_path=robot_prim_path,
ignore_substring=[
robot_prim_path,
"/World/target",
"/World/defaultGroundPlane",
"/curobo",
],
).get_collision_check_world()
# print(len(obstacles.objects))
motion_gen.update_world(obstacles)
# print("Updated World")
carb.log_info("Synced CuRobo world from stage.")
linear_color = np.ravel([249, 87, 56]) / 255.0
orient_color = np.ravel([103, 148, 54]) / 255.0
disable_color = np.ravel([255, 255, 255]) / 255.0
if cmd_plan is None and step_index % 10 == 0:
if plan_idx == 4:
print("Constrained: Holding tool linear-y")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([0, 0, 0, 0, 1, 0]),
)
target_material_plane.set_color(linear_color)
if plan_idx == 8:
print("Constrained: Holding tool Orientation and linear-y")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 0, 1, 0]),
)
target_material_plane.set_color(orient_color)
if plan_idx == 12:
print("Constrained: Holding tool linear-y, linear-x")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([0, 0, 0, 1, 1, 0]),
)
target_material_line.set_color(linear_color)
target_material_plane.set_color(disable_color)
if plan_idx == 16:
print("Constrained: Holding tool Orientation and linear-y, linear-x")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 1, 1, 0]),
)
target_material_line.set_color(orient_color)
target_material_plane.set_color(disable_color)
if plan_idx > 20:
plan_idx = 0
if plan_idx == 0:
print("Constrained: Reset")
target_material_line.set_color(disable_color)
target_material_plane.set_color(disable_color)
pose_cost_metric = None
plan_config.pose_cost_metric = pose_cost_metric
# motion generation:
for ks in range(len(target_material_list)):
if ks == target_idx:
target_material_list[ks].set_color(np.ravel([0, 1.0, 0]))
else:
target_material_list[ks].set_color(np.ravel([0.1, 0.1, 0.1]))
sim_js = robot.get_joints_state()
sim_js_names = robot.dof_names
cu_js = JointState(
position=tensor_args.to_device(sim_js.positions),
velocity=tensor_args.to_device(sim_js.velocities) * 0.0,
acceleration=tensor_args.to_device(sim_js.velocities) * 0.0,
jerk=tensor_args.to_device(sim_js.velocities) * 0.0,
joint_names=sim_js_names,
)
cu_js = cu_js.get_ordered_joint_state(motion_gen.kinematics.joint_names)
cube_position, cube_orientation = target_list[target_idx].get_world_pose()
# Set EE teleop goals, use cube for simple non-vr init:
ee_translation_goal = cube_position
ee_orientation_teleop_goal = cube_orientation
# compute curobo solution:
ik_goal = Pose(
position=tensor_args.to_device(ee_translation_goal),
quaternion=tensor_args.to_device(ee_orientation_teleop_goal),
)
result = motion_gen.plan_single(cu_js.unsqueeze(0), ik_goal, plan_config)
# ik_result = ik_solver.solve_single(ik_goal, cu_js.position.view(1,-1), cu_js.position.view(1,1,-1))
succ = result.success.item() # ik_result.success.item()
plan_idx += 1
if succ:
cmd_plan = result.get_interpolated_plan()
cmd_plan = motion_gen.get_full_js(cmd_plan)
# get only joint names that are in both:
idx_list = []
common_js_names = []
for x in sim_js_names:
if x in cmd_plan.joint_names:
idx_list.append(robot.get_dof_index(x))
common_js_names.append(x)
# idx_list = [robot.get_dof_index(x) for x in sim_js_names]
cmd_plan = cmd_plan.get_ordered_joint_state(common_js_names)
cmd_idx = 0
target_idx += 1
if target_idx >= len(target_list):
target_idx = 0
else:
carb.log_warn("Plan did not converge to a solution. No action is being taken.")
if cmd_plan is not None:
cmd_state = cmd_plan[cmd_idx]
# get full dof state
art_action = ArticulationAction(
cmd_state.position.cpu().numpy(),
cmd_state.velocity.cpu().numpy(),
joint_indices=idx_list,
)
# set desired joint angles obtained from IK:
articulation_controller.apply_action(art_action)
cmd_step_idx += 1
if cmd_step_idx == 2:
cmd_idx += 1
cmd_step_idx = 0
# for _ in range(2):
# my_world.step(render=False)
if cmd_idx >= len(cmd_plan.position):
cmd_idx = 0
cmd_plan = None
print("finished program")
simulation_app.close()

View File

@@ -51,6 +51,32 @@ parser.add_argument(
help="When True, runs in reactive mode",
default=False,
)
parser.add_argument(
"--constrain_grasp_approach",
action="store_true",
help="When True, approaches grasp with fixed orientation and motion only along z axis.",
default=False,
)
parser.add_argument(
"--reach_partial_pose",
nargs=6,
metavar=("qx", "qy", "qz", "x", "y", "z"),
help="Reach partial pose",
type=float,
default=None,
)
parser.add_argument(
"--hold_partial_pose",
nargs=6,
metavar=("qx", "qy", "qz", "x", "y", "z"),
help="Hold partial pose while moving to goal",
type=float,
default=None,
)
args = parser.parse_args()
############################################################
@@ -97,7 +123,12 @@ from curobo.util_file import (
join_path,
load_yaml,
)
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
from curobo.wrap.reacher.motion_gen import (
MotionGen,
MotionGenConfig,
MotionGenPlanConfig,
PoseCostMetric,
)
############################################################
@@ -107,7 +138,7 @@ from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGen
def main():
# create a curobo motion gen instance:
num_targets = 0
# assuming obstacles are in objects_path:
my_world = World(stage_units_in_meters=1.0)
stage = my_world.stage
@@ -152,12 +183,12 @@ def main():
robot, robot_prim_path = add_robot_to_scene(robot_cfg, my_world)
articulation_controller = robot.get_articulation_controller()
articulation_controller = None
world_cfg_table = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
)
world_cfg_table.cuboid[0].pose[2] -= 0.04
world_cfg_table.cuboid[0].pose[2] -= 0.02
world_cfg1 = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_mesh_world()
@@ -171,12 +202,14 @@ def main():
trajopt_tsteps = 32
trim_steps = None
max_attempts = 4
interpolation_dt = 0.05
if args.reactive:
trajopt_tsteps = 36
trajopt_tsteps = 40
trajopt_dt = 0.05
optimize_dt = False
max_attemtps = 1
max_attempts = 1
trim_steps = [1, None]
interpolation_dt = trajopt_dt
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg,
world_cfg,
@@ -184,12 +217,15 @@ def main():
collision_checker_type=CollisionCheckerType.MESH,
num_trajopt_seeds=12,
num_graph_seeds=12,
interpolation_dt=0.05,
interpolation_dt=interpolation_dt,
collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
optimize_dt=optimize_dt,
trajopt_dt=trajopt_dt,
trajopt_tsteps=trajopt_tsteps,
trim_steps=trim_steps,
# velocity_scale=0.1,
# self_collision_check=False,
# self_collision_opt=False,
)
motion_gen = MotionGen(motion_gen_config)
print("warming up...")
@@ -216,6 +252,9 @@ def main():
i = 0
spheres = None
past_cmd = None
target_orientation = None
past_orientation = None
pose_metric = None
while simulation_app.is_running():
my_world.step(render=True)
if not my_world.is_playing():
@@ -228,6 +267,9 @@ def main():
step_index = my_world.current_time_step_index
# print(step_index)
if articulation_controller is None:
# robot.initialize()
articulation_controller = robot.get_articulation_controller()
if step_index < 2:
my_world.reset()
robot._articulation_view.initialize()
@@ -265,6 +307,10 @@ def main():
past_pose = cube_position
if target_pose is None:
target_pose = cube_position
if target_orientation is None:
target_orientation = cube_orientation
if past_orientation is None:
past_orientation = cube_orientation
sim_js = robot.get_joints_state()
sim_js_names = robot.dof_names
@@ -313,8 +359,12 @@ def main():
if (np.max(np.abs(sim_js.velocities)) < 0.2) or args.reactive:
robot_static = True
if (
np.linalg.norm(cube_position - target_pose) > 1e-3
(
np.linalg.norm(cube_position - target_pose) > 1e-3
or np.linalg.norm(cube_orientation - target_orientation) > 1e-3
)
and np.linalg.norm(past_pose - cube_position) == 0.0
and np.linalg.norm(past_orientation - cube_orientation) == 0.0
and robot_static
):
# Set EE teleop goals, use cube for simple non-vr init:
@@ -326,12 +376,24 @@ def main():
position=tensor_args.to_device(ee_translation_goal),
quaternion=tensor_args.to_device(ee_orientation_teleop_goal),
)
plan_config.pose_cost_metric = pose_metric
result = motion_gen.plan_single(cu_js.unsqueeze(0), ik_goal, plan_config)
# ik_result = ik_solver.solve_single(ik_goal, cu_js.position.view(1,-1), cu_js.position.view(1,1,-1))
succ = result.success.item() # ik_result.success.item()
if num_targets == 1:
if args.constrain_grasp_approach:
pose_metric = PoseCostMetric.create_grasp_approach_metric()
if args.reach_partial_pose is not None:
reach_vec = motion_gen.tensor_args.to_device(args.reach_partial_pose)
pose_metric = PoseCostMetric(
reach_partial_pose=True, reach_vec_weight=reach_vec
)
if args.hold_partial_pose is not None:
hold_vec = motion_gen.tensor_args.to_device(args.hold_partial_pose)
pose_metric = PoseCostMetric(hold_partial_pose=True, hold_vec_weight=hold_vec)
if succ:
num_targets += 1
cmd_plan = result.get_interpolated_plan()
cmd_plan = motion_gen.get_full_js(cmd_plan)
# get only joint names that are in both:
@@ -350,7 +412,9 @@ def main():
else:
carb.log_warn("Plan did not converge to a solution. No action is being taken.")
target_pose = cube_position
target_orientation = cube_orientation
past_pose = cube_position
past_orientation = cube_orientation
if cmd_plan is not None:
cmd_state = cmd_plan[cmd_idx]
past_cmd = cmd_state.clone()

View File

@@ -34,6 +34,13 @@ parser.add_argument(
default=None,
help="To run headless, use one of [native, websocket], webrtc might not work.",
)
parser.add_argument(
"--constrain_grasp_approach",
action="store_true",
help="When True, approaches grasp with fixed orientation and motion only along z axis.",
default=False,
)
args = parser.parse_args()
# Third Party
@@ -78,6 +85,7 @@ from curobo.wrap.reacher.motion_gen import (
MotionGenConfig,
MotionGenPlanConfig,
MotionGenResult,
PoseCostMetric,
)
@@ -87,6 +95,7 @@ class CuroboController(BaseController):
my_world: World,
my_task: BaseStacking,
name: str = "curobo_controller",
constrain_grasp_approach: bool = False,
) -> None:
BaseController.__init__(self, name=name)
self._save_log = False
@@ -145,13 +154,16 @@ class CuroboController(BaseController):
collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
store_ik_debug=self._save_log,
store_trajopt_debug=self._save_log,
state_finite_difference_mode="CENTRAL",
acceleration_scale=0.5,
fixed_iters_trajopt=True,
velocity_scale=0.75,
)
self.motion_gen = MotionGen(motion_gen_config)
print("warming up...")
self.motion_gen.warmup()
self.motion_gen.warmup(parallel_finetune=True)
pose_metric = None
if constrain_grasp_approach:
pose_metric = PoseCostMetric.create_grasp_approach_metric(
offset_position=0.1, tstep_fraction=0.6
)
self.plan_config = MotionGenPlanConfig(
enable_graph=False,
@@ -159,6 +171,8 @@ class CuroboController(BaseController):
enable_graph_attempt=None,
enable_finetune_trajopt=True,
partial_ik_opt=False,
parallel_finetune=True,
pose_cost_metric=pose_metric,
)
self.usd_help.load_stage(self.my_world.stage)
self.cmd_plan = None
@@ -185,7 +199,7 @@ class CuroboController(BaseController):
cu_js,
[cube_name],
sphere_fit_type=SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
world_objects_pose_offset=Pose.from_list([0, 0, 0.005, 1, 0, 0, 0], self.tensor_args),
world_objects_pose_offset=Pose.from_list([0, 0, 0.01, 1, 0, 0, 0], self.tensor_args),
)
def detach_obj(self) -> None:
@@ -259,7 +273,7 @@ class CuroboController(BaseController):
# get full dof state
art_action = ArticulationAction(
cmd_state.position.cpu().numpy(),
# cmd_state.velocity.cpu().numpy(),
cmd_state.velocity.cpu().numpy() * 0.0,
joint_indices=self.idx_list,
)
if self.cmd_idx >= len(self.cmd_plan.position):
@@ -417,7 +431,9 @@ my_world.add_task(my_task)
my_world.reset()
robot_name = my_task.get_params()["robot_name"]["value"]
my_franka = my_world.scene.get_object(robot_name)
my_controller = CuroboController(my_world=my_world, my_task=my_task)
my_controller = CuroboController(
my_world=my_world, my_task=my_task, constrain_grasp_approach=args.constrain_grasp_approach
)
articulation_controller = my_franka.get_articulation_controller()
set_camera_view(eye=[2, 0, 1], target=[0.00, 0.00, 0.00], camera_prim_path="/OmniverseKit_Persp")
wait_steps = 8
@@ -439,17 +455,17 @@ print(articulation_controller.get_gains())
print(articulation_controller.get_max_efforts())
robot = my_franka
print("**********************")
robot.enable_gravity()
articulation_controller.set_gains(
kps=np.array(
[100000000, 6000000.0, 10000000, 600000.0, 25000.0, 15000.0, 50000.0, 6000.0, 6000.0]
if True:
robot.enable_gravity()
articulation_controller.set_gains(
kps=np.array(
[100000000, 6000000.0, 10000000, 600000.0, 25000.0, 15000.0, 50000.0, 6000.0, 6000.0]
)
)
)
articulation_controller.set_max_efforts(
values=np.array([100000, 52.199997, 100000, 52.199997, 7.2, 7.2, 7.2, 50.0, 50])
)
articulation_controller.set_max_efforts(
values=np.array([100000, 52.199997, 100000, 52.199997, 7.2, 7.2, 7.2, 50.0, 50])
)
print("Updated gains:")
print(articulation_controller.get_gains())