constrained planning, robot segmentation

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()

examples/isaac_sim/motion_gen_reacher.py

@@ -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()

examples/isaac_sim/simple_stacking.py

@@ -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())

examples/kinematics_example.py

@@ -15,7 +15,7 @@
 import torch
 
 # CuRobo
-from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
 from curobo.types.base import TensorDeviceType
 from curobo.types.robot import RobotConfig
 from curobo.util.logger import setup_curobo_logger

examples/motion_gen_example.py

@@ -42,8 +42,9 @@ def plot_traj(trajectory, dt):
         axs[3].plot(timesteps, qddd[:, i], label=str(i))
 
     plt.legend()
-    # plt.savefig("test.png")
-    plt.show()
+    plt.savefig("test.png")
+    plt.close()
+    # plt.show()
 
 
 def plot_iters_traj(trajectory, d_id=1, dof=7, seed=0):
@@ -140,37 +141,54 @@ def demo_motion_gen(js=False):
         world_file,
         tensor_args,
         interpolation_dt=0.01,
+        # trajopt_dt=0.15,
+        # velocity_scale=0.1,
+        use_cuda_graph=True,
+        # finetune_dt_scale=2.5,
+        interpolation_steps=10000,
     )
 
     motion_gen = MotionGen(motion_gen_config)
+    motion_gen.warmup(parallel_finetune=True)
 
-    motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js)
-    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
-    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    # motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js, parallel_finetune=True)
+    # robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    # robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
     retract_cfg = motion_gen.get_retract_config()
     state = motion_gen.rollout_fn.compute_kinematics(
         JointState.from_position(retract_cfg.view(1, -1))
     )
 
     retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
-    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.1)
+    start_state = JointState.from_position(retract_cfg.view(1, -1))
     goal_state = start_state.clone()
-    goal_state.position[..., 3] -= 0.1
+
+    start_state.position[0, 0] += 0.25
+    # goal_state.position[0,0] += 0.5
     if js:
         result = motion_gen.plan_single_js(
             start_state,
             goal_state,
-            MotionGenPlanConfig(
-                max_attempts=1, enable_graph=False, enable_opt=True, enable_finetune_trajopt=True
-            ),
+            MotionGenPlanConfig(max_attempts=1, parallel_finetune=True),
         )
     else:
         result = motion_gen.plan_single(
-            start_state, retract_pose, MotionGenPlanConfig(max_attempts=1)
+            start_state,
+            retract_pose,
+            MotionGenPlanConfig(
+                max_attempts=1, parallel_finetune=True, enable_finetune_trajopt=True
+            ),
         )
-    traj = result.get_interpolated_plan()
-    print("Trajectory Generated: ", result.success, result.solve_time, result.status)
+    print(
+        "Trajectory Generated: ",
+        result.success,
+        result.solve_time,
+        result.status,
+        result.optimized_dt,
+    )
     if PLOT and result.success.item():
+        traj = result.get_interpolated_plan()
+
         plot_traj(traj, result.interpolation_dt)
         # plt.save("test.png")
         # plt.close()
@@ -261,6 +279,45 @@ def demo_motion_gen_batch():
         plot_traj(traj[1, : result.path_buffer_last_tstep[1], :].cpu().numpy())
 
 
+def demo_motion_gen_goalset():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_cubby.yml"
+    robot_file = "franka.yml"
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_file,
+        world_file,
+        tensor_args,
+        collision_checker_type=CollisionCheckerType.PRIMITIVE,
+        use_cuda_graph=True,
+        num_trajopt_seeds=12,
+        num_graph_seeds=1,
+        num_ik_seeds=30,
+    )
+    motion_gen = MotionGen(motion_gen_config)
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    retract_cfg = motion_gen.get_retract_config()
+    state = motion_gen.rollout_fn.compute_kinematics(
+        JointState.from_position(retract_cfg.view(1, -1))
+    )
+
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.6)
+
+    state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+    goal_pose = Pose(
+        state.ee_pos_seq.repeat(2, 1).view(1, -1, 3),
+        quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
+    )
+    goal_pose.position[0, 0, 0] -= 0.1
+
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+    m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10)
+
+    result = motion_gen.plan_goalset(start_state, goal_pose, m_config)
+
+
 def demo_motion_gen_api():
     tensor_args = TensorDeviceType(device=torch.device("cuda:0"))
     interpolation_dt = 0.02
@@ -373,8 +430,9 @@ def demo_motion_gen_batch_env(n_envs: int = 10):
 
 if __name__ == "__main__":
     setup_curobo_logger("error")
-    # demo_motion_gen(js=True)
-    demo_motion_gen_batch()
+    demo_motion_gen(js=False)
+    # demo_motion_gen_batch()
+    # demo_motion_gen_goalset()
     # n = [2, 10]
     # for n_envs in n:
     #    demo_motion_gen_batch_env(n_envs=n_envs)

examples/robot_image_segmentation_example.py

@@ -0,0 +1,238 @@
+#
+# 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.
+#
+""" This example shows how to use cuRobo's kinematics to generate a mask. """
+
+
+# Standard Library
+import time
+
+# Third Party
+import imageio
+import numpy as np
+import torch
+import torch.autograd.profiler as profiler
+from nvblox_torch.datasets.mesh_dataset import MeshDataset
+from torch.profiler import ProfilerActivity, profile, record_function
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
+from curobo.geom.types import PointCloud, WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.camera import CameraObservation
+from curobo.types.math import Pose
+from curobo.types.robot import 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_segmenter import RobotSegmenter
+from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
+
+torch.manual_seed(30)
+
+torch.backends.cudnn.benchmark = True
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+
+def create_render_dataset(robot_file, save_debug_data: bool = False):
+    # load robot:
+    robot_dict = load_yaml(join_path(get_robot_configs_path(), robot_file))
+    robot_dict["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
+    robot_dict["robot_cfg"]["kinematics"]["load_meshes"] = True
+
+    robot_cfg = RobotConfig.from_dict(robot_dict["robot_cfg"])
+
+    kin_model = CudaRobotModel(robot_cfg.kinematics)
+
+    q = kin_model.retract_config
+
+    meshes = kin_model.get_robot_as_mesh(q)
+
+    world = WorldConfig(mesh=meshes[:])
+    world_table = WorldConfig.from_dict(
+        load_yaml(join_path(get_world_configs_path(), "collision_test.yml"))
+    )
+    world_table.cuboid[0].dims = [0.5, 0.5, 0.1]
+    world.add_obstacle(world_table.objects[0])
+    world.add_obstacle(world_table.objects[1])
+    if save_debug_data:
+        world.save_world_as_mesh("scene.stl", process_color=False)
+    robot_mesh = (
+        WorldConfig.create_merged_mesh_world(world, process_color=False).mesh[0].get_trimesh_mesh()
+    )
+
+    mesh_dataset = MeshDataset(
+        None,
+        n_frames=20,
+        image_size=480,
+        save_data_dir=None,
+        trimesh_mesh=robot_mesh,
+    )
+    q_js = JointState(position=q, joint_names=kin_model.joint_names)
+
+    return mesh_dataset, q_js
+
+
+def mask_image(robot_file="ur5e.yml"):
+    save_debug_data = False
+    # create robot segmenter:
+    tensor_args = TensorDeviceType()
+
+    curobo_segmenter = RobotSegmenter.from_robot_file(
+        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
+    )
+
+    mesh_dataset, q_js = create_render_dataset(robot_file, save_debug_data)
+
+    if save_debug_data:
+        visualize_scale = 10.0
+        data = mesh_dataset[0]
+        cam_obs = CameraObservation(
+            depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
+            intrinsics=data["intrinsics"],
+            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+        )
+        # save depth image
+        imageio.imwrite(
+            "camera_depth.png",
+            (cam_obs.depth_image * visualize_scale)
+            .squeeze()
+            .detach()
+            .cpu()
+            .numpy()
+            .astype(np.uint16),
+        )
+
+        # save robot spheres in current joint configuration
+        robot_kinematics = curobo_segmenter._robot_world.kinematics
+        sph = robot_kinematics.get_robot_as_spheres(q_js.position)
+        WorldConfig(sphere=sph[0]).save_world_as_mesh("robot_spheres.stl")
+
+        # save world pointcloud in robot origin
+
+        pc = cam_obs.get_pointcloud()
+        pc_obs = PointCloud("world", pose=cam_obs.pose.to_list(), points=pc)
+        pc_obs.save_as_mesh("camera_pointcloud.stl", transform_with_pose=True)
+
+        # run segmentation:
+        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
+            cam_obs,
+            q_js,
+        )
+        # save robot points as mesh
+
+        robot_mask = cam_obs.clone()
+        robot_mask.depth_image[~depth_mask] = 0.0
+        robot_mesh = PointCloud(
+            "world", pose=robot_mask.pose.to_list(), points=robot_mask.get_pointcloud()
+        )
+        robot_mesh.save_as_mesh("robot_segmented.stl", transform_with_pose=True)
+        # save depth image
+        imageio.imwrite(
+            "robot_depth.png",
+            (robot_mask.depth_image * visualize_scale)
+            .detach()
+            .squeeze()
+            .cpu()
+            .numpy()
+            .astype(np.uint16),
+        )
+
+        # save world points as mesh
+
+        world_mask = cam_obs.clone()
+        world_mask.depth_image[depth_mask] = 0.0
+        world_mesh = PointCloud(
+            "world", pose=world_mask.pose.to_list(), points=world_mask.get_pointcloud()
+        )
+        world_mesh.save_as_mesh("world_segmented.stl", transform_with_pose=True)
+
+        imageio.imwrite(
+            "world_depth.png",
+            (world_mask.depth_image * visualize_scale)
+            .detach()
+            .squeeze()
+            .cpu()
+            .numpy()
+            .astype(np.uint16),
+        )
+
+    dt_list = []
+
+    for i in range(len(mesh_dataset)):
+
+        data = mesh_dataset[i]
+        cam_obs = CameraObservation(
+            depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
+            intrinsics=data["intrinsics"],
+            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+        )
+        if not curobo_segmenter.ready:
+            curobo_segmenter.update_camera_projection(cam_obs)
+        st_time = time.time()
+
+        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
+            cam_obs,
+            q_js,
+        )
+
+        torch.cuda.synchronize()
+        dt_list.append(time.time() - st_time)
+
+    print("Segmentation Time (ms), (hz)", np.mean(dt_list[5:]) * 1000.0, 1.0 / np.mean(dt_list[5:]))
+
+
+def profile_mask_image(robot_file="ur5e.yml"):
+    # create robot segmenter:
+    tensor_args = TensorDeviceType()
+
+    curobo_segmenter = RobotSegmenter.from_robot_file(
+        robot_file, collision_sphere_buffer=0.0, distance_threshold=0.05, use_cuda_graph=False
+    )
+
+    mesh_dataset, q_js = create_render_dataset(robot_file)
+
+    dt_list = []
+    data = mesh_dataset[0]
+    cam_obs = CameraObservation(
+        depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
+        intrinsics=data["intrinsics"],
+        pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+    )
+    if not curobo_segmenter.ready:
+        curobo_segmenter.update_camera_projection(cam_obs)
+    depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(cam_obs, q_js)
+
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+
+        for i in range(len(mesh_dataset)):
+            with profiler.record_function("get_data"):
+
+                data = mesh_dataset[i]
+                cam_obs = CameraObservation(
+                    depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
+                    intrinsics=data["intrinsics"],
+                    pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+                )
+            st_time = time.time()
+            with profiler.record_function("segmentation"):
+
+                depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
+                    cam_obs, q_js
+                )
+
+    print("Exporting the trace..")
+    prof.export_chrome_trace("segmentation.json")
+
+
+if __name__ == "__main__":
+    mask_image()
+    # profile_mask_image()