release repository
This commit is contained in:
Balakumar Sundaralingam
366
examples/motion_gen_example.py
Normal file
366
examples/motion_gen_example.py
Normal file
@@ -0,0 +1,366 @@
|
||||
#
|
||||
# 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)
|
||||
Reference in New Issue
Block a user