440 lines
17 KiB
Python
440 lines
17 KiB
Python
|
from torch_robotics.isaac_gym_envs.motion_planning_envs import PandaMotionPlanningIsaacGymEnv, MotionPlanningController
|
||
|
|
|
||
|
|
import os
|
||
|
|
import pickle
|
||
|
|
from math import ceil
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import einops
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import torch
|
||
|
|
from einops._torch_specific import allow_ops_in_compiled_graph # requires einops>=0.6.1
|
||
|
|
|
||
|
|
from experiment_launcher import single_experiment_yaml, run_experiment
|
||
|
|
from mp_baselines.planners.costs.cost_functions import CostCollision, CostComposite, CostGPTrajectory
|
||
|
|
from mpd.models import TemporalUnet, UNET_DIM_MULTS
|
||
|
|
from mpd.models.diffusion_models.guides import GuideManagerTrajectoriesWithVelocity
|
||
|
|
from mpd.models.diffusion_models.sample_functions import guide_gradient_steps, ddpm_sample_fn
|
||
|
|
from mpd.trainer import get_dataset, get_model
|
||
|
|
from mpd.utils.loading import load_params_from_yaml
|
||
|
|
from torch_robotics.robots import RobotPanda
|
||
|
|
from torch_robotics.torch_utils.seed import fix_random_seed
|
||
|
|
from torch_robotics.torch_utils.torch_timer import TimerCUDA
|
||
|
|
from torch_robotics.torch_utils.torch_utils import get_torch_device, freeze_torch_model_params
|
||
|
|
from torch_robotics.trajectory.metrics import compute_smoothness, compute_path_length, compute_variance_waypoints
|
||
|
|
from torch_robotics.trajectory.utils import interpolate_traj_via_points
|
||
|
|
from torch_robotics.visualizers.planning_visualizer import PlanningVisualizer
|
||
|
|
|
||
|
|
allow_ops_in_compiled_graph()
|
||
|
|
|
||
|
|
|
||
|
|
TRAINED_MODELS_DIR = '../../data_trained_models/'
|
||
|
|
|
||
|
|
|
||
|
|
@single_experiment_yaml
|
||
|
|
def experiment(
|
||
|
|
########################################################################################################################
|
||
|
|
# Experiment configuration
|
||
|
|
# model_id: str = 'EnvDense2D-RobotPointMass',
|
||
|
|
# model_id: str = 'EnvNarrowPassageDense2D-RobotPointMass',
|
||
|
|
# model_id: str = 'EnvSimple2D-RobotPointMass',
|
||
|
|
model_id: str = 'EnvSpheres3D-RobotPanda',
|
||
|
|
|
||
|
|
# planner_alg: str = 'diffusion_prior',
|
||
|
|
# planner_alg: str = 'diffusion_prior_then_guide',
|
||
|
|
planner_alg: str = 'mpd',
|
||
|
|
|
||
|
|
use_guide_on_extra_objects_only: bool = False,
|
||
|
|
|
||
|
|
n_samples: int = 50,
|
||
|
|
|
||
|
|
start_guide_steps_fraction: float = 0.25,
|
||
|
|
n_guide_steps: int = 5,
|
||
|
|
n_diffusion_steps_without_noise: int = 5,
|
||
|
|
|
||
|
|
weight_grad_cost_collision: float = 1e-2,
|
||
|
|
weight_grad_cost_smoothness: float = 1e-7,
|
||
|
|
|
||
|
|
factor_num_interpolated_points_for_collision: float = 1.5,
|
||
|
|
|
||
|
|
trajectory_duration: float = 5.0, # currently fixed
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
device: str = 'cuda',
|
||
|
|
|
||
|
|
debug: bool = True,
|
||
|
|
|
||
|
|
render: bool = True,
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
# MANDATORY
|
||
|
|
seed: int = 30,
|
||
|
|
results_dir: str = 'logs',
|
||
|
|
|
||
|
|
########################################################################
|
||
|
|
**kwargs
|
||
|
|
):
|
||
|
|
########################################################################################################################
|
||
|
|
fix_random_seed(seed)
|
||
|
|
|
||
|
|
device = get_torch_device(device)
|
||
|
|
tensor_args = {'device': device, 'dtype': torch.float32}
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
print(f'##########################################################################################################')
|
||
|
|
print(f'Model -- {model_id}')
|
||
|
|
print(f'Algorithm -- {planner_alg}')
|
||
|
|
run_prior_only = False
|
||
|
|
run_prior_then_guidance = False
|
||
|
|
if planner_alg == 'mpd':
|
||
|
|
pass
|
||
|
|
elif planner_alg == 'diffusion_prior_then_guide':
|
||
|
|
run_prior_then_guidance = True
|
||
|
|
elif planner_alg == 'diffusion_prior':
|
||
|
|
run_prior_only = True
|
||
|
|
else:
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
model_dir = os.path.join(TRAINED_MODELS_DIR, model_id)
|
||
|
|
results_dir = os.path.join(model_dir, 'results_inference', str(seed))
|
||
|
|
os.makedirs(results_dir, exist_ok=True)
|
||
|
|
|
||
|
|
args = load_params_from_yaml(os.path.join(model_dir, "args.yaml"))
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Load dataset with env, robot, task
|
||
|
|
train_subset, train_dataloader, val_subset, val_dataloader = get_dataset(
|
||
|
|
dataset_class='TrajectoryDataset',
|
||
|
|
use_extra_objects=True,
|
||
|
|
obstacle_cutoff_margin=0.05,
|
||
|
|
**args,
|
||
|
|
tensor_args=tensor_args
|
||
|
|
)
|
||
|
|
dataset = train_subset.dataset
|
||
|
|
n_support_points = dataset.n_support_points
|
||
|
|
env = dataset.env
|
||
|
|
robot = dataset.robot
|
||
|
|
task = dataset.task
|
||
|
|
|
||
|
|
dt = trajectory_duration / n_support_points # time interval for finite differences
|
||
|
|
|
||
|
|
# set robot's dt
|
||
|
|
robot.dt = dt
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Load prior model
|
||
|
|
diffusion_configs = dict(
|
||
|
|
variance_schedule=args['variance_schedule'],
|
||
|
|
n_diffusion_steps=args['n_diffusion_steps'],
|
||
|
|
predict_epsilon=args['predict_epsilon'],
|
||
|
|
)
|
||
|
|
unet_configs = dict(
|
||
|
|
state_dim=dataset.state_dim,
|
||
|
|
n_support_points=dataset.n_support_points,
|
||
|
|
unet_input_dim=args['unet_input_dim'],
|
||
|
|
dim_mults=UNET_DIM_MULTS[args['unet_dim_mults_option']],
|
||
|
|
)
|
||
|
|
diffusion_model = get_model(
|
||
|
|
model_class=args['diffusion_model_class'],
|
||
|
|
model=TemporalUnet(**unet_configs),
|
||
|
|
tensor_args=tensor_args,
|
||
|
|
**diffusion_configs,
|
||
|
|
**unet_configs
|
||
|
|
)
|
||
|
|
diffusion_model.load_state_dict(
|
||
|
|
torch.load(os.path.join(model_dir, 'checkpoints', 'ema_model_current_state_dict.pth' if args['use_ema'] else 'model_current_state_dict.pth'),
|
||
|
|
map_location=tensor_args['device'])
|
||
|
|
)
|
||
|
|
diffusion_model.eval()
|
||
|
|
model = diffusion_model
|
||
|
|
|
||
|
|
freeze_torch_model_params(model)
|
||
|
|
model = torch.compile(model)
|
||
|
|
model.warmup(horizon=n_support_points, device=device)
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Random initial and final positions
|
||
|
|
n_tries = 100
|
||
|
|
start_state_pos, goal_state_pos = None, None
|
||
|
|
for _ in range(n_tries):
|
||
|
|
q_free = task.random_coll_free_q(n_samples=2)
|
||
|
|
start_state_pos = q_free[0]
|
||
|
|
goal_state_pos = q_free[1]
|
||
|
|
|
||
|
|
if torch.linalg.norm(start_state_pos - goal_state_pos) > dataset.threshold_start_goal_pos:
|
||
|
|
break
|
||
|
|
|
||
|
|
if start_state_pos is None or goal_state_pos is None:
|
||
|
|
raise ValueError(f"No collision free configuration was found\n"
|
||
|
|
f"start_state_pos: {start_state_pos}\n"
|
||
|
|
f"goal_state_pos: {goal_state_pos}\n")
|
||
|
|
|
||
|
|
print(f'start_state_pos: {start_state_pos}')
|
||
|
|
print(f'goal_state_pos: {goal_state_pos}')
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Run motion planning inference
|
||
|
|
|
||
|
|
########
|
||
|
|
# normalize start and goal positions
|
||
|
|
hard_conds = dataset.get_hard_conditions(torch.vstack((start_state_pos, goal_state_pos)), normalize=True)
|
||
|
|
context = None
|
||
|
|
|
||
|
|
########
|
||
|
|
# Set up the planning costs
|
||
|
|
|
||
|
|
# Cost collisions
|
||
|
|
cost_collision_l = []
|
||
|
|
weights_grad_cost_l = [] # for guidance, the weights_cost_l are the gradient multipliers (after gradient clipping)
|
||
|
|
if use_guide_on_extra_objects_only:
|
||
|
|
collision_fields = task.get_collision_fields_extra_objects()
|
||
|
|
else:
|
||
|
|
collision_fields = task.get_collision_fields()
|
||
|
|
|
||
|
|
for collision_field in collision_fields:
|
||
|
|
cost_collision_l.append(
|
||
|
|
CostCollision(
|
||
|
|
robot, n_support_points,
|
||
|
|
field=collision_field,
|
||
|
|
sigma_coll=1.0,
|
||
|
|
tensor_args=tensor_args
|
||
|
|
)
|
||
|
|
)
|
||
|
|
weights_grad_cost_l.append(weight_grad_cost_collision)
|
||
|
|
|
||
|
|
# Cost smoothness
|
||
|
|
cost_smoothness_l = [
|
||
|
|
CostGPTrajectory(
|
||
|
|
robot, n_support_points, dt, sigma_gp=1.0,
|
||
|
|
tensor_args=tensor_args
|
||
|
|
)
|
||
|
|
]
|
||
|
|
weights_grad_cost_l.append(weight_grad_cost_smoothness)
|
||
|
|
|
||
|
|
####### Cost composition
|
||
|
|
cost_func_list = [
|
||
|
|
*cost_collision_l,
|
||
|
|
*cost_smoothness_l
|
||
|
|
]
|
||
|
|
|
||
|
|
cost_composite = CostComposite(
|
||
|
|
robot, n_support_points, cost_func_list,
|
||
|
|
weights_cost_l=weights_grad_cost_l,
|
||
|
|
tensor_args=tensor_args
|
||
|
|
)
|
||
|
|
|
||
|
|
########
|
||
|
|
# Guiding manager
|
||
|
|
guide = GuideManagerTrajectoriesWithVelocity(
|
||
|
|
dataset,
|
||
|
|
cost_composite,
|
||
|
|
clip_grad=True,
|
||
|
|
interpolate_trajectories_for_collision=True,
|
||
|
|
num_interpolated_points=ceil(n_support_points * factor_num_interpolated_points_for_collision),
|
||
|
|
tensor_args=tensor_args,
|
||
|
|
)
|
||
|
|
|
||
|
|
t_start_guide = ceil(start_guide_steps_fraction * model.n_diffusion_steps)
|
||
|
|
sample_fn_kwargs = dict(
|
||
|
|
guide=None if run_prior_then_guidance or run_prior_only else guide,
|
||
|
|
n_guide_steps=n_guide_steps,
|
||
|
|
t_start_guide=t_start_guide,
|
||
|
|
noise_std_extra_schedule_fn=lambda x: 0.5,
|
||
|
|
)
|
||
|
|
|
||
|
|
########
|
||
|
|
# Sample trajectories with the diffusion/cvae model
|
||
|
|
with TimerCUDA() as timer_model_sampling:
|
||
|
|
trajs_normalized_iters = model.run_inference(
|
||
|
|
context, hard_conds,
|
||
|
|
n_samples=n_samples, horizon=n_support_points,
|
||
|
|
return_chain=True,
|
||
|
|
sample_fn=ddpm_sample_fn,
|
||
|
|
**sample_fn_kwargs,
|
||
|
|
n_diffusion_steps_without_noise=n_diffusion_steps_without_noise,
|
||
|
|
# ddim=True
|
||
|
|
)
|
||
|
|
print(f't_model_sampling: {timer_model_sampling.elapsed:.3f} sec')
|
||
|
|
t_total = timer_model_sampling.elapsed
|
||
|
|
|
||
|
|
########
|
||
|
|
# run extra guiding steps without diffusion
|
||
|
|
if run_prior_then_guidance:
|
||
|
|
n_post_diffusion_guide_steps = (t_start_guide + n_diffusion_steps_without_noise) * n_guide_steps
|
||
|
|
with TimerCUDA() as timer_post_model_sample_guide:
|
||
|
|
trajs = trajs_normalized_iters[-1]
|
||
|
|
trajs_post_diff_l = []
|
||
|
|
for i in range(n_post_diffusion_guide_steps):
|
||
|
|
trajs = guide_gradient_steps(
|
||
|
|
trajs,
|
||
|
|
hard_conds=hard_conds,
|
||
|
|
guide=guide,
|
||
|
|
n_guide_steps=1,
|
||
|
|
unnormalize_data=False
|
||
|
|
)
|
||
|
|
trajs_post_diff_l.append(trajs)
|
||
|
|
|
||
|
|
chain = torch.stack(trajs_post_diff_l, dim=1)
|
||
|
|
chain = einops.rearrange(chain, 'b post_diff_guide_steps h d -> post_diff_guide_steps b h d')
|
||
|
|
trajs_normalized_iters = torch.cat((trajs_normalized_iters, chain))
|
||
|
|
print(f't_post_diffusion_guide: {timer_post_model_sample_guide.elapsed:.3f} sec')
|
||
|
|
t_total = timer_model_sampling.elapsed + timer_post_model_sample_guide.elapsed
|
||
|
|
|
||
|
|
# unnormalize trajectory samples from the models
|
||
|
|
trajs_iters = dataset.unnormalize_trajectories(trajs_normalized_iters)
|
||
|
|
|
||
|
|
trajs_final = trajs_iters[-1]
|
||
|
|
trajs_final_coll, trajs_final_coll_idxs, trajs_final_free, trajs_final_free_idxs, _ = task.get_trajs_collision_and_free(trajs_final, return_indices=True)
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Compute motion planning metrics
|
||
|
|
print(f'\n----------------METRICS----------------')
|
||
|
|
print(f't_total: {t_total:.3f} sec')
|
||
|
|
|
||
|
|
success_free_trajs = task.compute_success_free_trajs(trajs_final)
|
||
|
|
fraction_free_trajs = task.compute_fraction_free_trajs(trajs_final)
|
||
|
|
collision_intensity_trajs = task.compute_collision_intensity_trajs(trajs_final)
|
||
|
|
|
||
|
|
print(f'success: {success_free_trajs}')
|
||
|
|
print(f'percentage free trajs: {fraction_free_trajs*100:.2f}')
|
||
|
|
print(f'percentage collision intensity: {collision_intensity_trajs*100:.2f}')
|
||
|
|
|
||
|
|
# compute costs only on collision-free trajectories
|
||
|
|
traj_final_free_best = None
|
||
|
|
idx_best_traj = None
|
||
|
|
cost_best_free_traj = None
|
||
|
|
cost_smoothness = None
|
||
|
|
cost_path_length = None
|
||
|
|
cost_all = None
|
||
|
|
variance_waypoint_trajs_final_free = None
|
||
|
|
if trajs_final_free is not None:
|
||
|
|
cost_smoothness = compute_smoothness(trajs_final_free, robot)
|
||
|
|
print(f'cost smoothness: {cost_smoothness.mean():.4f}, {cost_smoothness.std():.4f}')
|
||
|
|
|
||
|
|
cost_path_length = compute_path_length(trajs_final_free, robot)
|
||
|
|
print(f'cost path length: {cost_path_length.mean():.4f}, {cost_path_length.std():.4f}')
|
||
|
|
|
||
|
|
# compute best trajectory
|
||
|
|
cost_all = cost_path_length + cost_smoothness
|
||
|
|
idx_best_traj = torch.argmin(cost_all).item()
|
||
|
|
traj_final_free_best = trajs_final_free[idx_best_traj]
|
||
|
|
cost_best_free_traj = torch.min(cost_all).item()
|
||
|
|
print(f'cost best: {cost_best_free_traj:.3f}')
|
||
|
|
|
||
|
|
# variance of waypoints
|
||
|
|
variance_waypoint_trajs_final_free = compute_variance_waypoints(trajs_final_free, robot)
|
||
|
|
print(f'variance waypoint: {variance_waypoint_trajs_final_free:.4f}')
|
||
|
|
|
||
|
|
print(f'\n--------------------------------------\n')
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Save data
|
||
|
|
results_data_dict = {
|
||
|
|
'trajs_iters': trajs_iters,
|
||
|
|
'trajs_final_coll': trajs_final_coll,
|
||
|
|
'trajs_final_coll_idxs': trajs_final_coll_idxs,
|
||
|
|
'trajs_final_free': trajs_final_free,
|
||
|
|
'trajs_final_free_idxs': trajs_final_free_idxs,
|
||
|
|
'success_free_trajs': success_free_trajs,
|
||
|
|
'fraction_free_trajs': fraction_free_trajs,
|
||
|
|
'collision_intensity_trajs': collision_intensity_trajs,
|
||
|
|
'idx_best_traj': idx_best_traj,
|
||
|
|
'traj_final_free_best': traj_final_free_best,
|
||
|
|
'cost_best_free_traj': cost_best_free_traj,
|
||
|
|
'cost_path_length_trajs_final_free': cost_smoothness,
|
||
|
|
'cost_smoothness_trajs_final_free': cost_path_length,
|
||
|
|
'cost_all_trajs_final_free': cost_all,
|
||
|
|
'variance_waypoint_trajs_final_free': variance_waypoint_trajs_final_free,
|
||
|
|
't_total': t_total
|
||
|
|
}
|
||
|
|
with open(os.path.join(results_dir, 'results_data_dict.pickle'), 'wb') as handle:
|
||
|
|
pickle.dump(results_data_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
|
||
|
|
|
||
|
|
########################################################################################################################
|
||
|
|
# Render results
|
||
|
|
if render:
|
||
|
|
# Render
|
||
|
|
planner_visualizer = PlanningVisualizer(
|
||
|
|
task=task,
|
||
|
|
)
|
||
|
|
|
||
|
|
base_file_name = Path(os.path.basename(__file__)).stem
|
||
|
|
|
||
|
|
pos_trajs_iters = robot.get_position(trajs_iters)
|
||
|
|
|
||
|
|
planner_visualizer.animate_opt_iters_joint_space_state(
|
||
|
|
trajs=trajs_iters,
|
||
|
|
pos_start_state=start_state_pos, pos_goal_state=goal_state_pos,
|
||
|
|
vel_start_state=torch.zeros_like(start_state_pos), vel_goal_state=torch.zeros_like(goal_state_pos),
|
||
|
|
traj_best=traj_final_free_best,
|
||
|
|
video_filepath=os.path.join(results_dir, f'{base_file_name}-joint-space-opt-iters.mp4'),
|
||
|
|
n_frames=max((2, len(trajs_iters))),
|
||
|
|
anim_time=5
|
||
|
|
)
|
||
|
|
|
||
|
|
if isinstance(robot, RobotPanda):
|
||
|
|
# visualize in Isaac Gym
|
||
|
|
# POSITION CONTROL
|
||
|
|
# add initial positions for better visualization
|
||
|
|
n_first_steps = 10
|
||
|
|
n_last_steps = 10
|
||
|
|
|
||
|
|
trajs_pos = robot.get_position(trajs_final_free).movedim(1, 0)
|
||
|
|
trajs_vel = robot.get_velocity(trajs_final_free).movedim(1, 0)
|
||
|
|
|
||
|
|
trajs_pos = interpolate_traj_via_points(trajs_pos.movedim(0, 1), 2).movedim(1, 0)
|
||
|
|
|
||
|
|
motion_planning_isaac_env = PandaMotionPlanningIsaacGymEnv(
|
||
|
|
env, robot, task,
|
||
|
|
asset_root="../../deps/isaacgym/assets",
|
||
|
|
controller_type='position',
|
||
|
|
num_envs=trajs_pos.shape[1],
|
||
|
|
all_robots_in_one_env=True,
|
||
|
|
color_robots=False,
|
||
|
|
show_goal_configuration=True,
|
||
|
|
sync_with_real_time=True,
|
||
|
|
show_collision_spheres=False,
|
||
|
|
dt=dt,
|
||
|
|
**results_data_dict,
|
||
|
|
# show_collision_spheres=True
|
||
|
|
)
|
||
|
|
|
||
|
|
motion_planning_controller = MotionPlanningController(motion_planning_isaac_env)
|
||
|
|
motion_planning_controller.run_trajectories(
|
||
|
|
trajs_pos,
|
||
|
|
start_states_joint_pos=trajs_pos[0], goal_state_joint_pos=trajs_pos[-1][0],
|
||
|
|
n_first_steps=n_first_steps,
|
||
|
|
n_last_steps=n_last_steps,
|
||
|
|
visualize=True,
|
||
|
|
render_viewer_camera=True,
|
||
|
|
make_video=True,
|
||
|
|
video_path=os.path.join(results_dir, f'{base_file_name}-isaac-controller-position.mp4'),
|
||
|
|
make_gif=False
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
# visualize in the planning environment
|
||
|
|
planner_visualizer.animate_opt_iters_robots(
|
||
|
|
trajs=pos_trajs_iters, start_state=start_state_pos, goal_state=goal_state_pos,
|
||
|
|
traj_best=traj_final_free_best,
|
||
|
|
video_filepath=os.path.join(results_dir, f'{base_file_name}-traj-opt-iters.mp4'),
|
||
|
|
n_frames=max((2, len(trajs_iters))),
|
||
|
|
anim_time=5
|
||
|
|
)
|
||
|
|
|
||
|
|
planner_visualizer.animate_robot_trajectories(
|
||
|
|
trajs=pos_trajs_iters[-1], start_state=start_state_pos, goal_state=goal_state_pos,
|
||
|
|
plot_trajs=True,
|
||
|
|
video_filepath=os.path.join(results_dir, f'{base_file_name}-robot-traj.mp4'),
|
||
|
|
# n_frames=max((2, pos_trajs_iters[-1].shape[1]//10)),
|
||
|
|
n_frames=pos_trajs_iters[-1].shape[1],
|
||
|
|
anim_time=trajectory_duration
|
||
|
|
)
|
||
|
|
|
||
|
|
plt.show()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
# Leave unchanged
|
||
|
|
run_experiment(experiment)
|