""" Generates single-robot (RobotPlanarDisk) demonstrations for each of the 6 scenario environments, to be used for training single-agent SMD models (smd.planners.single_agent.mpd.SMD). These, combined with PrioritizedPlanning/ CBS (smd.planners.multi_agent), let us generate *real* coordinated multi-robot solutions as training data for composite SMD -- mirroring what the paper's own Appendix A.2 says they did (train a single-robot model, run it through search-based multi-robot coordination, use those solutions as SMD's training data) instead of the from-scratch synthetic method used in the first pass. Single-robot avoidance is a much easier problem than the composite case (no inter-agent coordination at all), so the same fast batch straight-line-plus- perturbation method from generate_data_all.py applies directly and should have very high yield. Usage: python generate_singleagent_data.py # e.g. empty, basic, room python generate_singleagent_data.py --all """ import os import sys import math import time import argparse 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 N_SUPPORT_POINTS = 64 DURATION = 5.0 OBSTACLE_CUTOFF_MARGIN = 0.05 # matches mpd.py's SMD planner's own obstacle_cutoff_margin INSTANCE_IDX = 0 ROBOT_ID = 'RobotPlanarDisk' N_CONTEXTS = 600 K_VARIANTS = 16 THRESHOLD_START_GOAL_POS = 1.0 SCENARIOS = { 'empty': dict(env_id='EnvEmptyNoWait2D', map_name='empty_map'), 'basic': dict(env_id='EnvEmptyNoWait2D', map_name='basic_map'), 'dense': dict(env_id='EnvEmptyNoWait2D', map_name='dense_map'), 'corridor': dict(env_id='EnvConveyor2D', map_name='corridor_map', threshold=0.6), 'room': dict(env_id='EnvHighways2D', map_name='room_map', threshold=0.6), 'shelf': dict(env_id='EnvDropRegion2D', map_name='shelf_map', threshold=0.6), } def build_task(env_id, map_name, 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 sample_context(task, threshold, max_tries=300): for _ in range(max_tries): 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: return s, g return None, None def build_candidate_batch(start, goal, k_variants, n_support_points, dt, tensor_args): q_dim = start.shape[0] # 2 for a single planar disk alphas = torch.linspace(0, 1, n_support_points, **tensor_args).unsqueeze(-1) base_pos = start.unsqueeze(0) * (1 - alphas) + goal.unsqueeze(0) * alphas t_norm = torch.linspace(0, 1, n_support_points, **tensor_args) envelope = torch.sin(math.pi * t_norm).unsqueeze(-1) pos_variants = [base_pos] for _ in range(k_variants - 1): detour = torch.zeros(n_support_points, q_dim, **tensor_args) n_harmonics = int(torch.randint(1, 3, (1,)).item()) for _h in range(n_harmonics): freq = float(torch.randint(1, 3, (1,)).item()) phase = torch.rand(1, **tensor_args).item() * math.pi shape = (torch.sin(freq * math.pi * t_norm + phase).unsqueeze(-1) * envelope) noise_scale = 0.05 + 0.45 * torch.rand(1, **tensor_args).item() offset_dir = (torch.rand(q_dim, **tensor_args) * 2 - 1) * noise_scale detour = detour + shape * offset_dir.unsqueeze(0) pos_variants.append(base_pos + detour) pos_batch = torch.stack(pos_variants, dim=0) avg_vel = (goal - start) / (n_support_points * dt) vel_batch = torch.zeros_like(pos_batch) vel_batch[:, 1:-1, :] = avg_vel.unsqueeze(0).unsqueeze(0) return torch.cat([pos_batch, vel_batch], dim=-1) def run_scenario(scenario): cfg = SCENARIOS[scenario] threshold = cfg.get('threshold', THRESHOLD_START_GOAL_POS) device = 'cuda' if torch.cuda.is_available() else 'cpu' tensor_args = {'device': device, 'dtype': torch.float32} dataset_subdir = f'singleagent_{scenario}_{cfg["env_id"]}-{ROBOT_ID}' results_dir = os.path.join(REPO_ROOT, 'data_trajectories', dataset_subdir) os.makedirs(results_dir, exist_ok=True) existing = [d for d in os.listdir(results_dir) if os.path.isdir(os.path.join(results_dir, d))] if len(existing) >= 100: print(f"[singleagent-{scenario}] already has {len(existing)} task dirs, skipping") return env, robot, task = build_task(cfg['env_id'], cfg['map_name'], tensor_args) dt = DURATION / N_SUPPORT_POINTS t_start = time.time() total = 0 saved = 0 for attempt in range(N_CONTEXTS): start, goal = sample_context(task, threshold) if start is None: continue traj_batch = build_candidate_batch(start, goal, K_VARIANTS, N_SUPPORT_POINTS, dt, tensor_args) _, traj_free = task.get_trajs_collision_and_free(traj_batch) if traj_free is None or traj_free.nelement() == 0: continue task_dir = os.path.join(results_dir, str(saved)) os.makedirs(task_dir, exist_ok=True) torch.save(traj_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, 'obstacle_cutoff_margin': OBSTACLE_CUTOFF_MARGIN}, f) with open(os.path.join(task_dir, 'metadata.yaml'), 'w') as f: yaml.dump({'env_id': cfg['env_id'], 'robot_id': ROBOT_ID, 'instance_idx': INSTANCE_IDX, 'map_name': cfg['map_name'], 'num_trajectories_free': int(traj_free.shape[0])}, f) total += traj_free.shape[0] saved += 1 if attempt % 100 == 0: print(f"[singleagent-{scenario}] progress: {attempt + 1}/{N_CONTEXTS} contexts, " f"{saved} saved, {total} total trajs, {time.time() - t_start:.1f}s elapsed") print(f"[singleagent-{scenario}] DONE. {total} total free trajectories across {saved} contexts " f"in {time.time() - t_start:.1f}s") if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('scenario', nargs='?', default=None) parser.add_argument('--all', action='store_true') args = parser.parse_args() if args.all: for s in SCENARIOS: run_scenario(s) elif args.scenario: run_scenario(args.scenario) else: parser.error("provide a scenario or --all")