scripts_v2/tools/record_partial_assemblies.py

# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md).
# All Rights Reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Script to record partial assemblies using IsaacLab framework."""

from __future__ import annotations

"""Launch Isaac Sim Simulator first."""

import argparse
import os
import torch
from tqdm import tqdm
from typing import cast

from isaaclab.app import AppLauncher

# add argparse arguments
parser = argparse.ArgumentParser(description="Record partial assemblies for object pairs.")
parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default="UW-FBLeg-PartialAssemblies-v0", help="Name of the task.")
parser.add_argument(
    "--dataset_dir", type=str, default="./Datasets/OmniReset/", help="Root Datasets/OmniReset/ directory."
)
parser.add_argument(
    "--num_trajectories", type=int, default=1, help="Number of physics trajectories to run for pose discovery."
)
parser.add_argument("--pos_similarity_threshold", type=float, default=0.001, help="Threshold for pose similarity.")
parser.add_argument(
    "--ori_similarity_threshold", type=float, default=0.01, help="Threshold for orientation similarity."
)

AppLauncher.add_app_launcher_args(parser)
args_cli, remaining_args = parser.parse_known_args()

# Launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

"""Rest everything else."""

import gymnasium as gym
import time

import isaaclab_tasks  # noqa: F401
from isaaclab.envs import ManagerBasedRLEnv

import uwlab_tasks  # noqa: F401
from uwlab_tasks.manager_based.manipulation.omnireset.mdp.utils import compute_pair_dir
from uwlab_tasks.utils.hydra import hydra_task_compose

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False


@hydra_task_compose(args_cli.task, "env_cfg_entry_point", hydra_args=remaining_args)
def main(env_cfg, agent_cfg) -> None:
    """Main function to record partial assemblies."""
    # create directory if it does not exist
    if not os.path.exists(args_cli.dataset_dir):
        os.makedirs(args_cli.dataset_dir, exist_ok=True)

    # override configurations with non-hydra CLI arguments
    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
    env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device

    # make sure environment is non-deterministic for diverse pose discovery
    env_cfg.seed = None

    # Create environment
    env = cast(ManagerBasedRLEnv, gym.make(args_cli.task, cfg=env_cfg)).unwrapped

    # Derive pair directory for output path
    insertive_usd_path = env_cfg.scene.insertive_object.spawn.usd_path
    receptive_usd_path = env_cfg.scene.receptive_object.spawn.usd_path
    pair = compute_pair_dir(insertive_usd_path, receptive_usd_path)

    print(f"Recording partial assemblies for: {pair}")
    print(f"Insertive: {insertive_usd_path}")
    print(f"Receptive: {receptive_usd_path}")

    # Reset environment (this will position objects at assembled state)
    env.reset()

    # Initialize pose tracking
    recorded_poses = []
    all_recorded_poses = None  # Keep track of ALL recorded poses for uniqueness checking

    # Run pose discovery
    episode_count = 0
    actions = torch.zeros(env.action_space.shape, device=env.device, dtype=torch.float32)

    # Create progress bar
    pbar = tqdm(total=args_cli.num_trajectories, desc="Trajectories", unit="episodes")
    start_time = time.time()
    total_poses_collected = 0

    while episode_count < args_cli.num_trajectories:
        # Step environment (forces will be applied automatically by pose_discovery event)
        _, rewards, terminated, truncated, _ = env.step(actions)
        dones = terminated | truncated

        # Reset environments that are done and update progress
        if dones.any():
            episodes_completed = dones.sum().item()
            episode_count += episodes_completed
            pbar.update(episodes_completed)

        # Get all pose data and use rewards from step
        if "current_pose_data" in env.extras["log"]:
            valid_mask = rewards > 0

            if valid_mask.any():
                # Filter pose data to only valid environments
                all_poses_data = env.extras["log"]["current_pose_data"]
                valid_poses_data = {key: all_poses_data[key][valid_mask] for key in all_poses_data.keys()}

                # Calculate relative poses for similarity checking
                relative_poses = valid_poses_data["relative_pose"]

                # Check uniqueness against ALL previously recorded poses
                if all_recorded_poses is not None:
                    # Calculate distance matrices between all pairs
                    relative_pos = relative_poses[:, :3]  # (N, 3)
                    relative_quat = relative_poses[:, 3:]  # (N, 4)
                    all_recorded_pos = all_recorded_poses[:, :3]  # (M, 3)
                    all_recorded_quat = all_recorded_poses[:, 3:]  # (M, 4)

                    # Compute distance matrices: (N, M)
                    pos_dists = torch.cdist(relative_pos, all_recorded_pos, p=2)  # Euclidean distance
                    ori_dists = torch.cdist(relative_quat, all_recorded_quat, p=2)  # Euclidean distance

                    # Find minimum distance to any previously recorded pose
                    min_pos_dists = torch.min(pos_dists, dim=1)[0]  # (N,)
                    min_ori_dists = torch.min(ori_dists, dim=1)[0]  # (N,)
                    new_pose_mask = (min_pos_dists > args_cli.pos_similarity_threshold) & (
                        min_ori_dists > args_cli.ori_similarity_threshold
                    )
                else:
                    new_pose_mask = torch.ones(len(relative_poses), dtype=torch.bool, device=env.device)

                # Save new unique poses
                if new_pose_mask.any():
                    new_poses = {key: valid_poses_data[key][new_pose_mask] for key in valid_poses_data.keys()}
                    recorded_poses.append(new_poses)

                    # Update all recorded poses for comparison
                    if all_recorded_poses is None:
                        all_recorded_poses = relative_poses[new_pose_mask]
                    else:
                        all_recorded_poses = torch.cat([all_recorded_poses, relative_poses[new_pose_mask]], dim=0)

                    # Update total poses collected
                    new_count = sum(len(batch["relative_position"]) for batch in recorded_poses)
                    if new_count > total_poses_collected:
                        total_poses_collected = new_count

            else:
                # No valid poses this step, continue
                pass

        # Check if simulation should stop
        if env.sim.is_stopped():
            break

    # Save any remaining poses
    if recorded_poses:
        _save_poses_to_dataset(recorded_poses, args_cli.dataset_dir, pair)

    pbar.close()

    print("Partial assembly recording complete!")
    print(f"Trajectories completed: {episode_count}")
    print(f"Poses recorded: {total_poses_collected}")
    print(f"Time taken: {(time.time() - start_time) / 60:.2f} minutes")
    if episode_count > 0:
        print(f"Average poses per trajectory: {total_poses_collected / episode_count:.1f}")

    env.close()


def _save_poses_to_dataset(pose_batches: list, dataset_dir: str, pair_name: str) -> None:
    """Save pose batches to Torch dataset (.pt)."""
    if not pose_batches:
        return

    # Concatenate all batches into single arrays
    all_poses = {}
    for key in pose_batches[0].keys():
        all_poses[key] = torch.cat([batch[key] for batch in pose_batches], dim=0).cpu()

    output_dir = os.path.join(dataset_dir, "Resets", pair_name)
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, "partial_assemblies.pt")
    torch.save(all_poses, output_file)

    print(f"Saved {len(all_poses['relative_position'])} poses to {output_file}")


if __name__ == "__main__":
    main()
    simulation_app.close()