#!/usr/bin/env python3

"""

Batch generate YAML config files for pick_and_place tasks.

Replaces asset names and USD paths based on available assets.

"""

import os

import shutil

from pathlib import Path

from typing import List, Dict, Optional

# Configuration

ASSETS_DIR = Path("workflows/simbox/assets/pick_and_place/pre-train-pick/assets")

FUNCTIONAL_ASSETS_DIR = Path("workflows/simbox/assets/pick_and_place/functional-pick-assets")

CONFIG_BASE_DIR = Path("workflows/simbox/core/configs/tasks/pick_and_place")

# Exclude these asset folders from pre-train-pick

EXCLUDE_ASSETS = ["google_scan-book", "google_scan-box", "omniobject3d-rubik_cube-old"]

# Template configs: (robot, task_type, arm_side, template_path, template_asset_name, assets_dir_type)

# task_type: "single_pick", "single_pnp", "single_func_pick", etc.

# arm_side: "left", "right", or None (for single-arm robots like franka)

# assets_dir_type: "pre-train" or "functional"

TEMPLATE_CONFIGS = [

# lift2 - single_pick (pre-train assets)

("lift2", "single_pick", "left", "lift2/single_pick/left/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

("lift2", "single_pick", "right", "lift2/single_pick/right/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

# split_aloha - single_pick (pre-train assets)

("split_aloha", "single_pick", "left", "split_aloha/single_pick/left/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

("split_aloha", "single_pick", "right", "split_aloha/single_pick/right/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

# franka - single_pick (pre-train assets, single arm, no left/right)

("franka", "single_pick", None, "franka/single_pick/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

# genie1 - single_pick (pre-train assets)

("genie1", "single_pick", "left", "genie1/single_pick/left/omniobject3d-lemon.yaml", "omniobject3d-lemon", "pre-train"),

("genie1", "single_pick", "right", "genie1/single_pick/right/omniobject3d-lemon.yaml", "omniobject3d-lemon", "pre-train"),

# genie1 - single_pnp (pre-train assets, pick and place)

("genie1", "single_pnp", "left", "genie1/single_pnp/left/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

("genie1", "single_pnp", "right", "genie1/single_pnp/right/omniobject3d-banana.yaml", "omniobject3d-banana", "pre-train"),

]

# Functional pick template configs (uses functional-pick-assets)

FUNCTIONAL_TEMPLATE_CONFIGS = [

# genie1 - single_func_pick (functional assets)

("genie1", "single_func_pick", "left", "genie1/single_func_pick/left/omniobject3d-hammer.yaml", "omniobject3d-hammer", "functional"),

("genie1", "single_func_pick", "right", "genie1/single_func_pick/right/omniobject3d-hammer.yaml", "omniobject3d-hammer", "functional"),

]

def get_all_assets(assets_dir: Path, exclude: List[str] = None) -> List[str]:

"""Get all asset folder names, excluding specified ones."""

if exclude is None:

exclude = []

assets = []

for item in sorted(assets_dir.iterdir()):

if item.is_dir() and item.name not in exclude:

assets.append(item.name)

return assets

def get_first_usd_path(assets_dir: Path, asset_name: str, assets_dir_type: str = "pre-train") -> Optional[str]:

"""Get the first USD file path for an asset."""

asset_folder = assets_dir / asset_name

if not asset_folder.exists():

return None

# Get all subfolders

subfolders = sorted([f for f in asset_folder.iterdir() if f.is_dir()])

if not subfolders:

return None

# Find the first subfolder with Aligned_obj.usd

for subfolder in subfolders:

usd_file = subfolder / "Aligned_obj.usd"

if usd_file.exists():

if assets_dir_type == "functional":

return f"pick_and_place/functional-pick-assets/{asset_name}/{subfolder.name}/Aligned_obj.usd"

else:

return f"pick_and_place/pre-train-pick/assets/{asset_name}/{subfolder.name}/Aligned_obj.usd"

return None

def replace_in_yaml(content: str, old_asset: str, new_asset: str, new_usd_path: str) -> str:

"""Replace asset name and USD path in YAML content.

Note: Does NOT modify distractors section - it should remain unchanged.

Only replaces paths that end with .usd (object paths, not directory paths).

"""

lines = content.split('\n')

new_lines = []

for line in lines:

# Replace path line ONLY if it's a USD file path (ends with .usd)

# This distinguishes object paths from distractor directory paths

# Support both pre-train-pick/assets and functional-pick-assets

if "path:" in line and "pick_and_place/" in line and ".usd" in line:

# Replace the USD path

indent = len(line) - len(line.lstrip())

new_lines.append(" " * indent + f"path: {new_usd_path}")

elif "category:" in line and old_asset in line:

# Replace category

indent = len(line) - len(line.lstrip())

new_lines.append(" " * indent + f'category: "{new_asset}"')

elif "task_dir:" in line and old_asset in line:

# Replace task_dir

indent = len(line) - len(line.lstrip())

# Extract the directory structure before the asset name

old_task_dir = line.split('"')[1]

new_task_dir = old_task_dir.replace(old_asset, new_asset)

new_lines.append(" " * indent + f'task_dir: "{new_task_dir}"')

else:

new_lines.append(line)

return '\n'.join(new_lines)

def generate_configs(

assets: List[str],

template_configs: List[tuple],

config_base_dir: Path,

assets_dir: Path,

dry_run: bool = True

):

"""Generate YAML config files for all assets."""

for robot, task_type, arm_side, template_rel_path, template_asset_name, assets_dir_type in template_configs:

template_path = config_base_dir / template_rel_path

if not template_path.exists():

print(f"[WARNING] Template not found: {template_path}")

continue

# Read template content

with open(template_path, 'r') as f:

template_content = f.read()

# Determine output directory based on task_type and arm_side

if arm_side is None:

# Single-arm robot (like franka)

output_dir = config_base_dir / robot / task_type

else:

# Dual-arm robot with left/right

output_dir = config_base_dir / robot / task_type / arm_side

output_dir.mkdir(parents=True, exist_ok=True)

for asset_name in assets:

# Get USD path for this asset

usd_path = get_first_usd_path(assets_dir, asset_name, assets_dir_type)

if usd_path is None:

print(f"[SKIP] No USD found for {asset_name}")

continue

# Generate new content

new_content = replace_in_yaml(

template_content,

template_asset_name,

asset_name,

usd_path

)

# Output file path

output_file = output_dir / f"{asset_name}.yaml"

if dry_run:

print(f"[DRY-RUN] Would create: {output_file}")

print(f" USD path: {usd_path}")

else:

with open(output_file, 'w') as f:

f.write(new_content)

print(f"[CREATED] {output_file}")

def main():

import argparse

parser = argparse.ArgumentParser(description="Batch generate YAML config files")

parser.add_argument("--dry-run", action="store_true", help="Show what would be done without creating files")

parser.add_argument("--assets", nargs="+", help="Specific assets to process (default: all)")

parser.add_argument("--robots", nargs="+", help="Specific robots to process (default: all)")

parser.add_argument("--task-types", nargs="+", help="Specific task types to process (default: all)")

parser.add_argument("--functional", action="store_true", help="Generate configs for functional pick assets")

args = parser.parse_args()

# Change to repo root

script_dir = Path(__file__).parent

os.chdir(script_dir.parent)

# Select template configs based on --functional flag

if args.functional:

template_configs = FUNCTIONAL_TEMPLATE_CONFIGS

assets_dir = FUNCTIONAL_ASSETS_DIR

exclude = [] # Don't exclude anything for functional assets

else:

template_configs = TEMPLATE_CONFIGS

assets_dir = ASSETS_DIR

exclude = EXCLUDE_ASSETS

# Get assets

if args.assets:

assets = args.assets

else:

assets = get_all_assets(assets_dir, exclude)

# Filter templates by robots and task_types if specified

if args.robots:

template_configs = [t for t in template_configs if t[0] in args.robots]

if args.task_types:

template_configs = [t for t in template_configs if t[1] in args.task_types]

print(f"Assets directory: {assets_dir}")

print(f"Found {len(assets)} assets to process")

print(f"Assets: {assets[:5]}..." if len(assets) > 5 else f"Assets: {assets}")

print(f"Templates: {[(t[0], t[1], t[2]) for t in template_configs]}")

print(f"Dry run: {args.dry_run}")

print("-" * 50)

generate_configs(

assets=assets,

template_configs=template_configs,

config_base_dir=CONFIG_BASE_DIR,

assets_dir=assets_dir,

dry_run=args.dry_run

)

if __name__ == "__main__":

main()