code/Diffusion-MRMP/scripts/repro/generate_singleagent_data.py

"""
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 <scenario>   # 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")