File size: 5,970 Bytes
0d96152 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """
Generate training demonstrations for SMD on the "empty" random-map scenario
with a 3-robot composite planar-disk robot.
This reimplements (for our composite/multi-robot robot) the same
generate_trajectories.py pipeline used upstream in mpd-public: sample a
random collision-free start/goal pair in the composite robot's joint
configuration space (this makes inter-robot collision avoidance and per-robot
obstacle avoidance just "collision avoidance for one big robot"), plan a
path with RRT-Connect, smooth/resample it to a fixed horizon with cubic
splines (matching torch_robotics' smoothen_trajectory, used identically in
mpd-public's data generation), and keep it if the resampled trajectory is
still collision-free.
Output layout matches what smd.datasets.trajectories.TrajectoryDatasetBase
expects to load:
data_trajectories/<dataset_subdir>/<task_id>/args.yaml
data_trajectories/<dataset_subdir>/<task_id>/metadata.yaml
data_trajectories/<dataset_subdir>/<task_id>/trajs-free.pt
"""
import os
import sys
import time
import yaml
import torch
REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, REPO_ROOT)
from torch_robotics import environments, robots
from torch_robotics.tasks.tasks import PlanningTask
from torch_robotics.trajectory.utils import smoothen_trajectory
from mp_baselines.planners.rrt_connect import RRTConnect
ENV_ID = 'EnvEmptyNoWait2D' # base name written to metadata; loader appends 'ExtraObjects'
ROBOT_ID = 'RobotCompositeThreePlanarDisk'
DATASET_SUBDIR = f'{ENV_ID}-{ROBOT_ID}'
MAP_NAME = 'empty_map'
INSTANCE_IDX = 0 # obstacle set to use while generating data (always empty here)
N_TASKS = 300 # number of distinct start/goal contexts
N_TRAJ_PER_TASK = 8 # RRT-Connect samples per context (RRT-Connect is randomized)
N_SUPPORT_POINTS = 64
DURATION = 5.0
THRESHOLD_START_GOAL_POS = 1.0
OBSTACLE_CUTOFF_MARGIN = 0.03
RRT_STEP_SIZE = 0.05
RRT_N_RADIUS = 0.3
RRT_N_ITERS = 5000
RRT_MAX_TIME = 15.0
RRT_N_PRE_SAMPLES = 5000
def build_task(tensor_args):
env_class = getattr(environments, ENV_ID + 'ExtraObjects')
env = env_class(tensor_args=tensor_args, instance_idx=INSTANCE_IDX, map_name=MAP_NAME)
robot_class = getattr(robots, ROBOT_ID)
robot = robot_class(tensor_args=tensor_args)
task = PlanningTask(
env=env, robot=robot, tensor_args=tensor_args,
obstacle_cutoff_margin=OBSTACLE_CUTOFF_MARGIN,
)
return env, robot, task
def generate_one_task(task, robot, tensor_args, results_dir, task_id):
# Sample a start/goal pair far enough apart in the composite config space.
start_state_pos, goal_state_pos = None, None
for _ in range(200):
q_free = task.random_coll_free_q(n_samples=2)
s, g = q_free[0], q_free[1]
if torch.linalg.norm(s - g) > THRESHOLD_START_GOAL_POS:
start_state_pos, goal_state_pos = s, g
break
if start_state_pos is None:
print(f"[task {task_id}] could not sample a far-enough start/goal pair, skipping")
return 0
trajs_free_l = []
for attempt in range(N_TRAJ_PER_TASK):
planner = RRTConnect(
task=task,
n_iters=RRT_N_ITERS,
start_state_pos=start_state_pos,
goal_state_pos=goal_state_pos,
step_size=RRT_STEP_SIZE,
n_radius=RRT_N_RADIUS,
max_time=RRT_MAX_TIME,
tensor_args=tensor_args,
n_pre_samples=RRT_N_PRE_SAMPLES,
)
path = planner.optimize(debug=False)
if path is None or len(path) < 2:
continue
path_t = torch.stack(list(path)) # (N, q_dim)
dt = DURATION / N_SUPPORT_POINTS
pos, vel = smoothen_trajectory(
path_t, n_support_points=N_SUPPORT_POINTS, dt=dt,
set_average_velocity=True, tensor_args=tensor_args,
)
traj = torch.cat((pos, vel), dim=-1).unsqueeze(0) # (1, H, 2*q_dim)
_, traj_free = task.get_trajs_collision_and_free(traj)
if traj_free is not None and traj_free.nelement() > 0:
trajs_free_l.append(traj_free)
if len(trajs_free_l) == 0:
print(f"[task {task_id}] no collision-free trajectories found")
return 0
trajs_free = torch.cat(trajs_free_l, dim=0)
task_dir = os.path.join(results_dir, str(task_id))
os.makedirs(task_dir, exist_ok=True)
torch.save(trajs_free.cpu(), os.path.join(task_dir, 'trajs-free.pt'))
with open(os.path.join(task_dir, 'args.yaml'), 'w') as f:
yaml.dump({
'threshold_start_goal_pos': THRESHOLD_START_GOAL_POS,
'obstacle_cutoff_margin': OBSTACLE_CUTOFF_MARGIN,
}, f)
with open(os.path.join(task_dir, 'metadata.yaml'), 'w') as f:
yaml.dump({
'env_id': ENV_ID,
'robot_id': ROBOT_ID,
'instance_idx': INSTANCE_IDX,
'map_name': MAP_NAME,
'num_trajectories_free': int(trajs_free.shape[0]),
}, f)
print(f"[task {task_id}] saved {trajs_free.shape[0]} free trajectories")
return trajs_free.shape[0]
def main():
device = 'cuda' if torch.cuda.is_available() else 'cpu'
tensor_args = {'device': device, 'dtype': torch.float32}
results_dir = os.path.join(REPO_ROOT, 'data_trajectories', DATASET_SUBDIR)
os.makedirs(results_dir, exist_ok=True)
env, robot, task = build_task(tensor_args)
t_start = time.time()
total = 0
for task_id in range(N_TASKS):
total += generate_one_task(task, robot, tensor_args, results_dir, task_id)
if task_id % 10 == 0:
elapsed = time.time() - t_start
print(f"--- progress: {task_id + 1}/{N_TASKS} tasks, {total} total trajs, {elapsed:.1f}s elapsed ---")
print(f"DONE. {total} total free trajectories across {N_TASKS} tasks in {time.time() - t_start:.1f}s")
if __name__ == '__main__':
main()
|