constrained planning, robot segmentation

2024-02-22 21:45:47 -08:00
parent 88eac64edc
commit bafdf80c05
102 changed files with 12440 additions and 8112 deletions
--- a/examples/isaac_sim/constrained_reacher.py
+++ b/examples/isaac_sim/constrained_reacher.py
@@ -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()