367 lines
12 KiB
Python
367 lines
12 KiB
Python
#
|
|
# 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.
|
|
#
|
|
|
|
# Standard Library
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# 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.util.logger import setup_curobo_logger
|
|
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
|
|
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
|
|
|
|
|
|
def plot_traj(trajectory, dt):
|
|
# Third Party
|
|
import matplotlib.pyplot as plt
|
|
|
|
_, axs = plt.subplots(4, 1)
|
|
q = trajectory.position.cpu().numpy()
|
|
qd = trajectory.velocity.cpu().numpy()
|
|
qdd = trajectory.acceleration.cpu().numpy()
|
|
qddd = trajectory.jerk.cpu().numpy()
|
|
timesteps = [i * dt for i in range(q.shape[0])]
|
|
for i in range(q.shape[-1]):
|
|
axs[0].plot(timesteps, q[:, i], label=str(i))
|
|
axs[1].plot(timesteps, qd[:, i], label=str(i))
|
|
axs[2].plot(timesteps, qdd[:, i], label=str(i))
|
|
axs[3].plot(timesteps, qddd[:, i], label=str(i))
|
|
|
|
plt.legend()
|
|
# plt.savefig("test.png")
|
|
plt.show()
|
|
|
|
|
|
def plot_iters_traj(trajectory, d_id=1, dof=7, seed=0):
|
|
# Third Party
|
|
import matplotlib.pyplot as plt
|
|
|
|
_, axs = plt.subplots(len(trajectory), 1)
|
|
if len(trajectory) == 1:
|
|
axs = [axs]
|
|
for k in range(len(trajectory)):
|
|
q = trajectory[k]
|
|
|
|
for i in range(len(q)):
|
|
axs[k].plot(
|
|
q[i][seed, :-1, d_id].cpu(),
|
|
"r+-",
|
|
label=str(i),
|
|
alpha=0.1 + min(0.9, float(i) / (len(q))),
|
|
)
|
|
plt.legend()
|
|
plt.show()
|
|
|
|
|
|
def plot_iters_traj_3d(trajectory, d_id=1, dof=7, seed=0):
|
|
# Third Party
|
|
import matplotlib.pyplot as plt
|
|
|
|
ax = plt.axes(projection="3d")
|
|
c = 0
|
|
h = trajectory[0][0].shape[1] - 1
|
|
x = [x for x in range(h)]
|
|
|
|
for k in range(len(trajectory)):
|
|
q = trajectory[k]
|
|
|
|
for i in range(len(q)):
|
|
# ax.plot3D(x,[c for _ in range(h)], q[i][seed, :, d_id].cpu())#, 'r')
|
|
ax.scatter3D(
|
|
x, [c for _ in range(h)], q[i][seed, :h, d_id].cpu(), c=q[i][seed, :, d_id].cpu()
|
|
)
|
|
# @plt.show()
|
|
c += 1
|
|
# plt.legend()
|
|
plt.show()
|
|
|
|
|
|
def demo_motion_gen_mesh():
|
|
PLOT = True
|
|
tensor_args = TensorDeviceType()
|
|
world_file = "collision_mesh_scene.yml"
|
|
robot_file = "franka.yml"
|
|
motion_gen_config = MotionGenConfig.load_from_robot_config(
|
|
robot_file,
|
|
world_file,
|
|
tensor_args,
|
|
trajopt_tsteps=40,
|
|
collision_checker_type=CollisionCheckerType.MESH,
|
|
use_cuda_graph=False,
|
|
)
|
|
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 = robot_cfg.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.5)
|
|
result = motion_gen.plan(
|
|
start_state,
|
|
retract_pose,
|
|
enable_graph=False,
|
|
enable_opt=True,
|
|
max_attempts=1,
|
|
num_trajopt_seeds=10,
|
|
num_graph_seeds=10,
|
|
)
|
|
print(result.status, result.attempts, result.trajopt_time)
|
|
traj = result.raw_plan # optimized plan
|
|
print("Trajectory Generated: ", result.success)
|
|
if PLOT:
|
|
plot_traj(traj.cpu().numpy())
|
|
|
|
|
|
def demo_motion_gen():
|
|
# Standard Library
|
|
PLOT = False
|
|
tensor_args = TensorDeviceType()
|
|
world_file = "collision_table.yml"
|
|
robot_file = "franka.yml"
|
|
motion_gen_config = MotionGenConfig.load_from_robot_config(
|
|
robot_file,
|
|
world_file,
|
|
tensor_args,
|
|
interpolation_dt=0.01,
|
|
)
|
|
|
|
motion_gen = MotionGen(motion_gen_config)
|
|
|
|
motion_gen.warmup(enable_graph=False)
|
|
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.3)
|
|
result = motion_gen.plan_single(start_state, retract_pose, MotionGenPlanConfig(max_attempts=1))
|
|
traj = result.get_interpolated_plan()
|
|
print("Trajectory Generated: ", result.success, result.optimized_dt.item())
|
|
if PLOT:
|
|
plot_traj(traj, result.interpolation_dt)
|
|
# plt.save("test.png")
|
|
# plt.close()
|
|
# traj = result.debug_info["opt_solution"].position.view(-1,7)
|
|
# plot_traj(traj.cpu().numpy())
|
|
|
|
|
|
def demo_motion_gen_debug():
|
|
PLOT = True
|
|
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,
|
|
trajopt_tsteps=24,
|
|
collision_checker_type=CollisionCheckerType.PRIMITIVE,
|
|
use_cuda_graph=True,
|
|
num_trajopt_seeds=1,
|
|
num_graph_seeds=1,
|
|
store_ik_debug=True,
|
|
store_trajopt_debug=True,
|
|
trajopt_particle_opt=False,
|
|
grad_trajopt_iters=100,
|
|
)
|
|
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 = robot_cfg.cspace.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.4)
|
|
result = motion_gen.plan(start_state, retract_pose, enable_graph=True, enable_opt=True)
|
|
if result.status not in [None, "Opt Fail"]:
|
|
return
|
|
traj = result.plan.view(-1, 7) # optimized plan
|
|
print("Trajectory Generated: ", result.success)
|
|
trajectory_iter_steps = result.debug_info["trajopt_result"].debug_info["solver"]["steps"]
|
|
|
|
if PLOT:
|
|
plot_iters_traj_3d(trajectory_iter_steps, d_id=6)
|
|
# plot_traj(traj.cpu().numpy())
|
|
|
|
|
|
def demo_motion_gen_batch():
|
|
PLOT = False
|
|
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,
|
|
trajopt_tsteps=30,
|
|
collision_checker_type=CollisionCheckerType.PRIMITIVE,
|
|
use_cuda_graph=True,
|
|
num_trajopt_seeds=4,
|
|
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))
|
|
)
|
|
|
|
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.6)
|
|
|
|
retract_pose = retract_pose.repeat_seeds(2)
|
|
retract_pose.position[0, 0] = -0.3
|
|
for _ in range(2):
|
|
result = motion_gen.plan_batch(
|
|
start_state.repeat_seeds(2),
|
|
retract_pose,
|
|
MotionGenPlanConfig(enable_graph=False, enable_opt=True),
|
|
)
|
|
traj = result.optimized_plan.position.view(2, -1, 7) # optimized plan
|
|
print("Trajectory Generated: ", result.success)
|
|
if PLOT:
|
|
plot_traj(traj[0, : result.path_buffer_last_tstep[0], :].cpu().numpy())
|
|
plot_traj(traj[1, : result.path_buffer_last_tstep[1], :].cpu().numpy())
|
|
|
|
|
|
def demo_motion_gen_api():
|
|
tensor_args = TensorDeviceType(device=torch.device("cuda:0"))
|
|
interpolation_dt = 0.02
|
|
# create motion gen with a cuboid cache to be able to load obstacles later:
|
|
motion_gen_cfg = MotionGenConfig.load_from_robot_config(
|
|
"franka.yml",
|
|
"collision_table.yml",
|
|
tensor_args,
|
|
trajopt_tsteps=34,
|
|
interpolation_steps=5000,
|
|
num_ik_seeds=50,
|
|
num_trajopt_seeds=6,
|
|
collision_checker_type=CollisionCheckerType.PRIMITIVE,
|
|
grad_trajopt_iters=500,
|
|
trajopt_dt=0.5,
|
|
interpolation_dt=interpolation_dt,
|
|
evaluate_interpolated_trajectory=True,
|
|
js_trajopt_dt=0.5,
|
|
js_trajopt_tsteps=34,
|
|
)
|
|
motion_gen = MotionGen(motion_gen_cfg)
|
|
|
|
motion_gen.warmup(warmup_js_trajopt=False)
|
|
|
|
# create world representation:
|
|
|
|
cuboids = [Cuboid(name="obs_1", pose=[0, 0, 0, 1, 0, 0, 0], dims=[0.1, 0.1, 0.1])]
|
|
world = WorldConfig(cuboid=cuboids)
|
|
|
|
motion_gen.update_world(world)
|
|
|
|
q_start = JointState.from_position(
|
|
tensor_args.to_device([[0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0.0]]),
|
|
joint_names=[
|
|
"panda_joint1",
|
|
"panda_joint2",
|
|
"panda_joint3",
|
|
"panda_joint4",
|
|
"panda_joint5",
|
|
"panda_joint6",
|
|
"panda_joint7",
|
|
],
|
|
)
|
|
|
|
goal_pose = Pose(
|
|
position=tensor_args.to_device([[0.5, 0.0, 0.3]]),
|
|
quaternion=tensor_args.to_device([[1, 0, 0, 0]]),
|
|
)
|
|
print(goal_pose)
|
|
|
|
result = motion_gen.plan_single(q_start, goal_pose)
|
|
|
|
# get result:
|
|
|
|
interpolated_solution = result.get_interpolated_plan()
|
|
|
|
print(result.get_interpolated_plan())
|
|
|
|
|
|
def demo_motion_gen_batch_env(n_envs: int = 10):
|
|
tensor_args = TensorDeviceType()
|
|
world_files = ["collision_table.yml" for _ in range(int(n_envs / 2))] + [
|
|
"collision_test.yml" for _ in range(int(n_envs / 2))
|
|
]
|
|
|
|
world_cfg = [
|
|
WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
|
|
for world_file in world_files
|
|
]
|
|
robot_file = "franka.yml"
|
|
motion_gen_config = MotionGenConfig.load_from_robot_config(
|
|
robot_file,
|
|
world_cfg,
|
|
tensor_args,
|
|
trajopt_tsteps=30,
|
|
use_cuda_graph=False,
|
|
num_trajopt_seeds=4,
|
|
num_ik_seeds=30,
|
|
num_batch_ik_seeds=30,
|
|
evaluate_interpolated_trajectory=True,
|
|
interpolation_dt=0.05,
|
|
interpolation_steps=500,
|
|
grad_trajopt_iters=30,
|
|
)
|
|
motion_gen_batch_env = MotionGen(motion_gen_config)
|
|
motion_gen_batch_env.reset()
|
|
motion_gen_batch_env.warmup(
|
|
enable_graph=False, batch=n_envs, warmup_js_trajopt=False, batch_env_mode=True
|
|
)
|
|
retract_cfg = motion_gen_batch_env.get_retract_config().clone()
|
|
state = motion_gen_batch_env.compute_kinematics(
|
|
JointState.from_position(retract_cfg.view(1, -1))
|
|
)
|
|
|
|
goal_pose = Pose(
|
|
state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
|
|
).repeat_seeds(n_envs)
|
|
|
|
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(n_envs)
|
|
|
|
goal_pose.position[1, 0] -= 0.2
|
|
|
|
m_config = MotionGenPlanConfig(
|
|
False, True, max_attempts=1, enable_graph_attempt=None, enable_finetune_trajopt=False
|
|
)
|
|
result = motion_gen_batch_env.plan_batch_env(start_state, goal_pose, m_config)
|
|
|
|
print(n_envs, result.total_time, result.total_time / n_envs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
setup_curobo_logger("error")
|
|
demo_motion_gen()
|
|
n = [2, 10]
|
|
for n_envs in n:
|
|
demo_motion_gen_batch_env(n_envs=n_envs)
|