start_batch.py

import numpy as np
import torch
import os
import sys
import json
sys.path.append(os.getcwd())

from points_trajectory_prediction.trajectory_predictor import TrajectoryTransformer
from safetensors.torch import load_model,save_file
from points_trajectory_prediction.utils.llm_utils import description_matching
import math
import open3d as o3d





def closed_loop(env, model, cloth_config, cloth_name):
    device = torch.device("cuda")

    threshold = cloth_config["threshold"]  
    step = cloth_config["frames_per_step"]
    description_embed_list = cloth_config["description_embed_list"]
    env.world_pause()

    print("get_now_pcd_from_camera")
    now_real_pcd = env.get_now_pcd_from_camera()
    output_pcd=[]
    output_pcd.append(now_real_pcd)
    cloth_inital_pcd = now_real_pcd
    for i in range(cloth_config["stage_num"]):
        description_embed = description_embed_list[i].to(torch.float32).to(device)
        for _ in range(cloth_config["step_turn"][i]):
            times = 0
            stage_terminal = False
            first_contact = np.array([0, 0, 0])
            first_time = True
            begin = True
            print("DESCRIPTION",i)
   

            while True:
                if times >= math.ceil(cloth_config["frame_per_stage"][i]/(step)) or stage_terminal:
                    break
                

                target_num_points = 2048          
                pts = np.array(now_real_pcd)

                if pts.shape[0] > target_num_points:    # downsample
                    sampled_indices = np.random.choice(pts.shape[0], target_num_points, replace=False)
                    pts = pts[sampled_indices, :]
                elif pts.shape[0] < target_num_points:
                    sampled_indices = np.random.choice(pts.shape[0], target_num_points - pts.shape[0], replace=True)
                    pts = np.concatenate((pts, pts[sampled_indices, :]), axis=0)        

                if cloth_config["type"] in ["Pants"]:
                    pts[:,0]*=-1
                    pts[:,2]*=-1

                aligned_pcd = torch.from_numpy(pts).to(torch.float32)     
                aligned_pcd = aligned_pcd.to(device)

                pcd_for_predict, transform_scale, transform_offset = normalize_point_cloud_trajectory(aligned_pcd.unsqueeze(0).unsqueeze(2))

                cpu_transform_scale = transform_scale.cpu().numpy()
                cpu_transform_offset = transform_offset.cpu().numpy()
                
                outputs_traj = predict(model, pcd_for_predict[:,:,0,:], description_embed, "").detach().cpu().numpy()

                aligned_pcd = pcd_for_predict[0,:,0,:].detach().cpu().numpy()

                pred_next_pcds = outputs_traj[0, :, :(step), :]
                # if times >= math.ceil(cloth_config["frame_per_stage"][i]/(step)/2):
                #     pred_next_pcds = outputs_traj[0, :, :-1, :]


                masked_pred_next_pcds=pred_next_pcds[:, -1, :]
                hmask = None
                step_control = cloth_config["force_per_stage"][i]
                if (cloth_config["type"] in ["Short-sleeve","Long-sleeve","No-sleeve"] or (cloth_config["type"] in ["Pants"] and i==1)) and first_time is True and step_control==2:
                    extent = np.ptp(aligned_pcd[:, 2])  # Peak-to-peak range (max - min)
                    mask = aligned_pcd[:, 2] >=  (0.5 *extent+np.min(aligned_pcd[:, 2]))
                    hmask = np.zeros((aligned_pcd.shape[0]))
                    hmask[mask]=1
                    print(hmask.shape)
                if (cloth_config["type"] in ["Short-sleeve","Long-sleeve"] or (cloth_config["type"] in ["Pants"] and i==1)) and first_time is False and step_control==2:
                    extent = np.ptp(aligned_pcd[:, 2])  # Peak-to-peak range (max - min)
                    mask = first_contact
                    hmask = np.zeros((aligned_pcd.shape[0]))
                    hmask[mask]=1
                    print(hmask.shape)

                contacts, forces = ManiFM_model(aligned_pcd, masked_pred_next_pcds,hmask)
                

                if first_time:
                    first_time = False
                    first_contact = np.array([], dtype=int) 
                    for contact in contacts:
                        distances = np.linalg.norm(aligned_pcd - contact, axis=1)
                        indices = np.where(distances < 0.2)[0]
                        first_contact = np.concatenate((first_contact, indices))
                    print("LEN",len(first_contact))# Because real robot can only grasp the same point
                    
                contacts = np.array(contacts) * cpu_transform_scale[0, 0, 0] + cpu_transform_offset[0, 0]
                forces = np.array(forces) * cpu_transform_scale[0, 0, 0] 
                

                if cloth_config["type"] in ["Pants"]:
                    contacts[:,0]*=-1
                    contacts[:,2]*=-1
                    forces[:,0]*=-1
                    forces[:,2]*=-1

                env.control_to_next_step(begin,step_control,contacts, forces)  
                begin = False 
                last_pred_pcd = pred_next_pcds[:, -1, :]     
                last_real_pcd = now_real_pcd
                now_real_pcd = env.get_now_pcd_from_camera()    

                times += 1
               
            env.wait_step()
            now_real_pcd = env.get_now_pcd_from_camera() 
            output_pcd.append(now_real_pcd)
    
    # output meshs
    index = 0
    output_path = os.path.join("./isaac_sim/output", cloth_name)
    if not os.path.exists(output_path):
        os.makedirs(output_path)
        print(f"Created folder: {output_path}")
    base_name = "mesh"
    while True:
        folder_name = os.path.join(output_path, f"{base_name}{index}")
        if not os.path.exists(folder_name):
            os.makedirs(folder_name)
            print(f"Created folder: {folder_name}")
            break
        index += 1
  
    np.savetxt(folder_name+'/fin_pcd.txt', now_real_pcd)
    np.savetxt(folder_name+'/initial_pcd.txt', cloth_inital_pcd)
    for i,pcd in enumerate(output_pcd):
        np.savetxt(folder_name+'/step'+str(i)+'_pcd.txt', pcd)

def save_embeded_list(key, descriptions, embeded_list, cache_dir="./isaac_sim/output/cache"):
    
    os.makedirs(cache_dir, exist_ok=True)
    meta_path = os.path.join(cache_dir, f"{key}_meta.json")
    embed_path = os.path.join(cache_dir, f"{key}_embed.pt")
    
    with open(meta_path, 'w') as meta_file:
        json.dump({"descriptions": descriptions}, meta_file)
    
    torch.save(embeded_list, embed_path)


def load_embeded_list(key, cache_dir="./isaac_sim/output/cache"):
    
    meta_path = os.path.join(cache_dir, f"{key}_meta.json")
    embed_path = os.path.join(cache_dir, f"{key}_embed.pt")
    
    if not os.path.exists(meta_path) or not os.path.exists(embed_path):
        return None, None
    
    with open(meta_path, 'r') as meta_file:
        meta_data = json.load(meta_file)
        descriptions = meta_data.get("descriptions", [])
    
    embeded_list = torch.load(embed_path)
    return descriptions, embeded_list

def process_key_descriptions(key, descriptions, embed_dict, cache_dir="./isaac_sim/output/cache"):
   
    cached_descriptions, cached_embeded_list = load_embeded_list(key, cache_dir)
    
    if cached_descriptions == descriptions:
        return cached_embeded_list
    
    embeded_list = []
    for description in descriptions:
        _, description_embed = description_matching(description, embed_dict)
        embeded_list.append(description_embed)
    
    embeded_list = torch.stack(embeded_list) 
    save_embeded_list(key, descriptions, embeded_list, cache_dir)
    return embeded_list

if __name__=="__main__":


    device = torch.device('cuda:0')
    
    model = TrajectoryTransformer(
        input_dim=128,      # or the appropriate input dimension
        hidden_dim=256,     # make sure this matches the hidden dimension from the checkpoint
        output_dim=128,     # or the appropriate output dimension
        nhead=4,            # use the correct number of attention heads
        num_encoder_layers=4,  # set to 4 to match the checkpoint
        num_decoder_layers=4,  # adjust if necessary
        num_points=2048, 
        num_frames=21, 
        point_dim=3,
        device=device
    ).to(device)

    # input
    #--------------------------------------------------------------------------------------------
    embed_dict = torch.load('./data/description_embeddings_mirrored.pt')
    model.load_state_dict(torch.load('./data/model_1113_199.pth', map_location=device))
    batch_input_path = "./data/Cloth-Simulation/Batch_Input/batch_input.json"

    with open(batch_input_path, 'r') as file:
        batch_input = json.load(file)
    instructions_path = "./data/Cloth-Simulation/Batch_Input/Instructions.json"
    with open(instructions_path, 'r') as file:
        instructions = json.load(file)


    output_cache = os.path.join("./isaac_sim/output", "cache")

    if not os.path.exists(output_cache):
        os.makedirs(output_cache)
        print(f"Created folder: {output_cache}")    
    # preprocess instructions
    #---------------------------------------------------------------
    preprocess_instructions = {}
    for key, instruction in instructions.items():

        cloth_config_path = instruction["cloth_config_path"]
        descriptions = instruction["descriptions"]
        with open(cloth_config_path, 'r') as file:
            cloth_arg_config = json.load(file)

        description_embed_list = []

        description_embed_list = process_key_descriptions(key, descriptions, embed_dict, output_cache)

        preprocess_instructions[key] = {
            "type": cloth_arg_config["type"],
            "scale": cloth_arg_config["scale"],
            "stage_num": cloth_arg_config["stage_num"],                 # total stage(action) num
            "frame_per_stage": cloth_arg_config["frame_per_stage"],     # CONST_STAGE
            "force_per_stage": cloth_arg_config["force_per_stage"],     # Force number in each stage(action)
            "step_turn": cloth_arg_config["step_turn"],                 # Currently useless 
            "frames_per_step": cloth_arg_config["frames_per_step"],     # CONST_STEP
            "threshold": cloth_arg_config["threshold"],
            "description_embed_list": description_embed_list,  
        }


    # batch input
    #---------------------------------------------------------------    
    from isaac_sim.physxDemo.cloth import cloth_main,normalize_point_cloud_trajectory,predict,cloth_next
    from isaac_sim.physxDemo.cloth import ClothDemoEnv
    from ManiFM_clothing.close_loop_pred import ManiFM_model
    output_path = os.path.join("./isaac_sim/output", "log")
    index  = 0
    if not os.path.exists(output_path):
        os.makedirs(output_path)
        print(f"Created folder: {output_path}")
    base_name = "log"

    while True:
        file_name = os.path.join(output_path, f"{base_name}{index}.txt")
        if os.path.exists(file_name):
            index += 1
        else:
            break
    with open(file_name, 'w') as file:
        pass

    cloth_root = batch_input[0]["cloth_root"]   # batch_input.json 's first item used for init
    env = None
    for bi,input in enumerate(batch_input[1:]):
        
        cloth_type = input["cloth_type"]
        cloth_name = input["cloth_name"]
        if cloth_type not in ["Short-sleeve"]:
            continue
        cloth_config = preprocess_instructions[cloth_type]
        if bi % 1 == 0:
            if env is not None:

                env.world_delete_prim("/World")
                env.world_delete_prim("/Replicator")
            env = cloth_main(cloth_root, cloth_name, cloth_config["scale"])
        else :
            env = cloth_next(env,cloth_root, cloth_name, cloth_config["scale"])
        
        # env.stop()
        closed_loop(env, model, cloth_config, cloth_name)
        

        with open(file_name, 'a') as file:
            print("\n",cloth_type, cloth_name,"  Result:",file=file)
            env.check_success(cloth_type,file=file)

        for i in range(500):
            env.step()
        
        env.world_delete_prim("/World/Garment/garment")