""" MIT License Copyright (c) 2024 Yorai Shaoul Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ # Standard imports. # from torch_robotics.isaac_gym_envs.motion_planning_envs import PandaMotionPlanningIsaacGymEnv, MotionPlanningController import os import pickle from datetime import datetime import time from math import ceil from pathlib import Path import einops import matplotlib.pyplot as plt import numpy as np import pandas as pd import torch from einops._torch_specific import allow_ops_in_compiled_graph # requires einops>=0.6.1 from typing import Tuple, List import concurrent.futures # Project imports. from experiment_launcher import single_experiment_yaml, run_experiment from mp_baselines.planners.costs.cost_functions import CostCollision, CostComposite, CostGPTrajectory, CostConstraint from smd.models import TemporalUnet, UNET_DIM_MULTS from smd.models.diffusion_models.guides import GuideManagerTrajectoriesWithVelocity from smd.models.diffusion_models.sample_functions import guide_gradient_steps, ddpm_sample_fn from smd.trainer import get_dataset, get_model from smd.utils.loading import load_params_from_yaml from torch_robotics.robots import * 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, \ compute_average_acceleration, compute_average_acceleration_from_pos_vel, compute_path_length_from_pos from torch_robotics.trajectory.utils import interpolate_traj_via_points from torch_robotics.visualizers.planning_visualizer import PlanningVisualizer from torch_robotics.robots.robot_planar_disk import RobotPlanarDisk from smd.planners.multi_agent import CBS, PrioritizedPlanning, SMDComposite from smd.planners.single_agent import SMD, SMDEnsemble from smd.common.constraints import MultiPointConstraint, VertexConstraint, EdgeConstraint from smd.common.conflicts import VertexConflict, PointConflict, EdgeConflict from smd.common.trajectory_utils import smooth_trajs, densify_trajs from smd.common import compute_collision_intensity, is_multi_agent_start_goal_states_valid, global_pad_paths, \ get_start_goal_pos_circle, get_state_pos_column, get_start_goal_pos_boundary, get_start_goal_pos_random_in_env from smd.common.pretty_print import * from smd.config.smd_params import SMDParams as params from smd.common.experiments import MultiAgentPlanningSingleTrialConfig, get_result_dir_from_trial_config from torch_robotics.environments import * allow_ops_in_compiled_graph() TRAINED_MODELS_DIR = '../../data_trained_models/' device = 'cuda' device = get_torch_device(device) tensor_args = {'device': device, 'dtype': torch.float32} def run_multi_agent_trial(test_config: MultiAgentPlanningSingleTrialConfig): # ============================ # Start time per agent. # ============================ start_time_l = [i * test_config.stagger_start_time_dt for i in range(test_config.num_agents)] # ============================ # Arguments for the high/low level planner. # ============================ low_level_planner_model_args = { 'planner_alg': 'smd', 'use_guide_on_extra_objects_only': params.use_guide_on_extra_objects_only, 'n_samples': params.n_samples, 'n_local_inference_noising_steps': params.n_local_inference_noising_steps, 'n_local_inference_denoising_steps': params.n_local_inference_denoising_steps, 'start_guide_steps_fraction': params.start_guide_steps_fraction, 'n_guide_steps': params.n_guide_steps, 'n_diffusion_steps_without_noise': params.n_diffusion_steps_without_noise, 'weight_grad_cost_collision': params.weight_grad_cost_collision, 'weight_grad_cost_smoothness': params.weight_grad_cost_smoothness, 'weight_grad_cost_constraints': params.weight_grad_cost_constraints, 'weight_grad_cost_soft_constraints': params.weight_grad_cost_soft_constraints, 'factor_num_interpolated_points_for_collision': params.factor_num_interpolated_points_for_collision, 'trajectory_duration': params.trajectory_duration, 'device': params.device, 'debug': params.debug, 'seed': params.seed, 'results_dir': params.results_dir, 'trained_models_dir': TRAINED_MODELS_DIR, 'instance_idx': test_config.instance_idx, 'map_name': test_config.map_name, 'init_traj4proj': test_config.init_traj4proj, 'proj_params': test_config.proj_params } high_level_planner_model_args = { 'is_xcbs': True if test_config.multi_agent_planner_class in ["XECBS", "XCBS"] else False, 'is_ecbs': True if test_config.multi_agent_planner_class in ["ECBS", "XECBS"] else False, 'start_time_l': start_time_l, 'runtime_limit': test_config.runtime_limit, 'conflict_type_to_constraint_types': {PointConflict: {MultiPointConstraint}}, 'instance_idx': test_config.instance_idx, 'map_name': test_config.map_name, 'init_traj4proj': test_config.init_traj4proj, 'proj_params': test_config.proj_params } # ============================ # Create a results directory. # ============================ results_dir = get_result_dir_from_trial_config(test_config, test_config.time_str, test_config.trial_number) os.makedirs(results_dir, exist_ok=True) num_agents = test_config.num_agents # ============================ # Get planning problem. # ============================ # If want to get random starts and goals, then must do that after creating the reference task and robot. start_l = test_config.start_state_pos_l goal_l = test_config.goal_state_pos_l global_model_ids = test_config.global_model_ids agent_skeleton_l = test_config.agent_skeleton_l # ============================ # Transforms and model tiles setup. # ============================ # Create a reference planner from which we'll use the task and robot as the reference on in CBS. # Those are used for collision checking and visualization. This has a skeleton of all tiles. reference_agent_skeleton = [[r, c] for r in range(len(global_model_ids)) for c in range(len(global_model_ids[0]))] # ============================ # Transforms from tiles to global frame. # ============================ tile_width = 2.0 tile_height = 2.0 global_model_transforms = [[torch.tensor([x * tile_width, -y * tile_height], **tensor_args) for x in range(len(global_model_ids[0]))] for y in range(len(global_model_ids))] # ============================ # Parse the single agent planner class name. # ============================ if test_config.single_agent_planner_class == "SMD": low_level_planner_class = SMD elif test_config.single_agent_planner_class == "SMDEnsemble": low_level_planner_class = SMDEnsemble else: raise ValueError(f'Unknown single agent planner class: {test_config.single_agent_planner_class}') # ============================ # Create reference agent planner. # ============================ # And for the reference skeleton. reference_task = None reference_robot = None # ============================ # Run trial. # ============================ exp_name = f'smd_single_trial' # Transform starts and goals to the global frame. Right now they are in the local tile frames. start_l = [start_l[i] + global_model_transforms[agent_skeleton_l[i][0][0]][agent_skeleton_l[i][0][1]] for i in range(num_agents)] goal_l = [goal_l[i] + global_model_transforms[agent_skeleton_l[i][-1][0]][agent_skeleton_l[i][-1][1]] for i in range(num_agents)] # ============================ # Create global transforms for each agent's skeleton. # ============================ # Each agent has a dict entry. Each entry is a dict with the skeleton steps (0, 1, 2, ...), mapping to the # model transform. agent_model_transforms_l = [] agent_model_ids_l = [] for agent_id in range(num_agents): agent_model_transforms = {} agent_model_ids = {} for skeleton_step in range(len(agent_skeleton_l[agent_id])): skeleton_model_coord = agent_skeleton_l[agent_id][skeleton_step] agent_model_transforms[skeleton_step] = global_model_transforms[skeleton_model_coord[0]][ skeleton_model_coord[1]] agent_model_ids[skeleton_step] = global_model_ids[skeleton_model_coord[0]][skeleton_model_coord[1]] agent_model_transforms_l.append(agent_model_transforms) agent_model_ids_l.append(agent_model_ids) # Change the dict of the model ids to a list sorted by the skeleton steps. agent_model_ids_l = [[agent_model_ids_l[i][j] for j in range(len(agent_model_ids_l[i]))] for i in range(num_agents)] # ============================ # Create the low level planners. # ============================ planners_creation_start_time = time.time() low_level_planner_l = [] if test_config.multi_agent_planner_class != "SMDComposite": for i in range(num_agents): low_level_planner_model_args_i = low_level_planner_model_args.copy() low_level_planner_model_args_i['start_state_pos'] = start_l[i] low_level_planner_model_args_i['goal_state_pos'] = goal_l[i] low_level_planner_model_args_i['model_ids'] = agent_model_ids_l[i] low_level_planner_model_args_i['transforms'] = agent_model_transforms_l[i] if test_config.single_agent_planner_class == "SMD": # Set the model_id to the first one. low_level_planner_model_args_i['model_id'] = agent_model_ids_l[i][0] low_level_planner_l.append(low_level_planner_class(**low_level_planner_model_args_i)) print('Planners creation time:', time.time() - planners_creation_start_time) print("\n\n\n\n") # ============================ # Create the multi agent planner. # ============================ if test_config.multi_agent_planner_class == "SMDComposite": assert len(global_model_ids) == 1 assert len(global_model_ids[0]) == 1 model_id = global_model_ids[0][0] multi_agent_planner_class = SMDComposite high_level_planner_model_args['model_id'] = model_id high_level_planner_model_args['results_dir'] = results_dir else: raise ValueError(f'Unknown multi agent planner class: {test_config.multi_agent_planner_class}') planner = multi_agent_planner_class(low_level_planner_l, start_l, goal_l, reference_task=reference_task, reference_robot=reference_robot, **high_level_planner_model_args) # ============================ # Plan. # ============================ startt = time.time() whole_trajs = planner.plan(runtime_limit=test_config.runtime_limit) planning_time = time.time() - startt # Print planning times. print(GREEN, 'Planning times:', planning_time, RESET) # ============================ # Save outputs. # ============================ results_dir_uri = f'file://{os.path.abspath(results_dir)}' print('Results dir:', results_dir_uri) paths_np = whole_trajs.cpu().numpy() with open(os.path.join(results_dir, 'paths.npy'), 'wb') as f: np.save(f, paths_np) # save map info res_map_name = test_config.map_name res_instance_idx = test_config.instance_idx with open(os.path.join(results_dir, 'map_info.pkl'), 'wb') as f: pickle.dump({'map_name': res_map_name, 'instance_idx': res_instance_idx}, f) if __name__ == '__main__': test_config_single_tile = MultiAgentPlanningSingleTrialConfig() test_config_single_tile.num_agents = 3 test_config_single_tile.instance_name = "test" test_config_single_tile.multi_agent_planner_class = "XECBS" # Or "ECBS" or "XCBS" or "CBS" or "PP". test_config_single_tile.single_agent_planner_class = "SMDEnsemble" # Or "SMD" test_config_single_tile.stagger_start_time_dt = 0 test_config_single_tile.runtime_limit = 60 * 3 # 3 minutes. test_config_single_tile.time_str = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") test_config_single_tile.render_animation = True # Change the `densify_trajs` call above to create nicer animations. example_type = "single_tile" # example_type = "multi_tile" # ============================ # Single tile. # ============================ if example_type == "single_tile": # Choose the model to use. A model is for a map/robot combination. # test_config_single_tile.global_model_ids = [['EnvEmpty2D-RobotPlanarDisk']] test_config_single_tile.global_model_ids = [['EnvEmptyNoWait2D-RobotPlanarDisk']] # test_config_single_tile.global_model_ids = [['EnvConveyor2D-RobotPlanarDisk']] # test_config_single_tile.global_model_ids = [['EnvHighways2D-RobotPlanarDisk']] # test_config_single_tile.global_model_ids = [['EnvDropRegion2D-RobotPlanarDisk']] # Choose starts and goals. test_config_single_tile.agent_skeleton_l = [[[0, 0]]] * test_config_single_tile.num_agents torch.random.manual_seed(10) test_config_single_tile.start_state_pos_l, test_config_single_tile.goal_state_pos_l = \ get_start_goal_pos_circle(test_config_single_tile.num_agents, 0.8) # Another option is to get random starts and goals. # get_start_goal_pos_random_in_env(test_config_single_tile.num_agents, # EnvDropRegion2D, # tensor_args, # margin=0.2, # obstacle_margin=0.11) # A third option is to get starts and goals in a "boundary" formation. # get_start_goal_pos_boundary(test_config_single_tile.num_agents, 0.85) # And a final option is to hard-code starts and goals. # (torch.tensor([[-0.8, 0], [0.8, -0]], **tensor_args), # torch.tensor([[0.8, -0], [-0.8, 0]], **tensor_args)) print("Starts:", test_config_single_tile.start_state_pos_l) print("Goals:", test_config_single_tile.goal_state_pos_l) run_multi_agent_trial(test_config_single_tile) print(GREEN, 'OK.', RESET) # ============================ # Multiple tiles example. # ============================ if example_type == "multi_tile": test_config_multiple_tiles = test_config_single_tile test_config_multiple_tiles.num_agents = 4 test_config_multiple_tiles.stagger_start_time_dt = 5 test_config_multiple_tiles.global_model_ids = \ [['EnvEmptyNoWait2D-RobotPlanarDisk', 'EnvEmptyNoWait2D-RobotPlanarDisk']] test_config_multiple_tiles.agent_skeleton_l = [[[0, 0], [0, 1]], [[0, 1], [0, 0]], [[0, 0], [0, 1]], [[0, 1], [0, 0]]] test_config_multiple_tiles.start_state_pos_l, test_config_multiple_tiles.goal_state_pos_l = \ (torch.tensor([[0, 0.8], [0, 0.3], [0, -0.3], [0, -0.8]], **tensor_args), torch.tensor([[0, -0.8], [0, -0.3], [0, 0.3], [0, 0.8]], **tensor_args)) print(test_config_multiple_tiles.start_state_pos_l) test_config_multiple_tiles.multi_agent_planner_class = "XECBS" run_multi_agent_trial(test_config_multiple_tiles) print(GREEN, 'OK.', RESET)