hofee/gen_data_curobo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

constrained planning, robot segmentation

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

88
examples/motion_gen_example.py
View File

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