HACMan: Learning Hybrid Actor-Critic Maps for 6D Non-Prehensile Manipulation
from bin_env.base_env import BaseEnv
from bin_env.poke_env import PokeEnv
from robosuite.environments.base import register_env

register_env(BaseEnv)
register_env(PokeEnv)

bin_env/__init__.py
HACMan-2023/HACMan
from .housekeep_object import HousekeepSampler
bin_env/assets/__init__.py
HACMan-2023/HACMan
import open3d as o3d
import os
import h5py
from sklearn.decomposition import PCA
import numpy as np
import pandas as pd
import os
import json

import matplotlib.pyplot as plt
import seaborn as sns

root_dir = os.path.dirname(os.path.abspath(__file__))
dirname = 'bin_env/assets/housekeep/models'
processed_dirname = 'bin_env/assets/housekeep/models_processed/meshes'
info_dirname = 'bin_env/assets/housekeep/object_info'
v_hacd_path = '../v-hacd/app/build/TestVHACD'
manifold_path = '../Manifold/build/manifold'

def process_objs():
    for obj_file in os.listdir(dirname):
        if obj_file[-4:] != '.obj':
            continue
        
        print(obj_file)
        obj_path = os.path.join(dirname, obj_file)
        
        # v-hacd
        v_hacd_dir = os.path.dirname(v_hacd_path)
        os.system(f'{v_hacd_path} {obj_path}')
        
        # manifold
        os.system(f'{manifold_path} decomp.obj decomp.obj 20000')
        
        mesh = o3d.io.read_triangle_mesh("decomp.obj", True)
        # mesh = o3d.io.read_triangle_mesh(obj_path, True)
        mesh.compute_vertex_normals()
        # pcd = o3d.geometry.sample_points_uniformly(mesh, number_of_points=10000)
        # points = np.asarray(pcd.points)
        # pcd_scaled = o3d.geometry.PointCloud()
        # pcd_scaled.points = o3d.utility.Vector3dVector(points)
        vertices = np.asarray(mesh.vertices)
        print(vertices.shape)
        print("Water tight? ", mesh.is_watertight())

        max_points = np.max(vertices, axis=0)
        min_points = np.min(vertices, axis=0)
        print(max_points)
        print(min_points)
        size = (max_points-min_points)/2
        origin = (max_points+min_points)/2
        
        stl_file = obj_file[:-4] + '.stl'
        stl_path = os.path.join(processed_dirname, stl_file)
        # o3d.visualization.draw_geometries([mesh])
        o3d.io.write_triangle_mesh(stl_path, mesh)
        
        for f in ["decomp.mtl", "decomp.obj", "decomp.stl"]:
            if os.path.exists(f):
                os.remove(f)
        
        filename = stl_file[:-4] +".h5"
        filepath = os.path.join(info_dirname, filename)
        with h5py.File(filepath, 'w') as f:
            f.create_dataset("origin", data=origin)
            f.create_dataset("size", data=size)

def generate_metadata():
    metadata = {
        'name': [],
        'type': [],
    }
    
    for h5_file in os.listdir(info_dirname):
        if h5_file[-3:] != '.h5':
            continue
        
        obj_name = h5_file[:-3]
        obj_type = obj_name.split('_')[0]
        metadata['name'].append(obj_name)
        metadata['type'].append(obj_type)
        
    df = pd.DataFrame(metadata)
    df.to_csv(os.path.join(root_dir, 'metainfo.csv'))
    return df

def generate_dataset_split(df: pd.DataFrame):
    # Split into train, unseen instance, unseen class

    # Unseen instance: 2 objects per type from the top 5 classes:
    # Cup, mug, bowl, bottle, plate
    unseen_instance_classes = ['cup', 'mug', 'bottle', 'pillbottle']
    unseen_instances = []
    for obj_type in unseen_instance_classes:
        obj_names = df[df['type'] == obj_type]['name'].values
        unseen_instances.extend(np.random.choice(obj_names, 1, replace=False))

    # Unseen class: 15 objects from the rest of the classes
    # rest_classes = np.setdiff1d(df['type'].unique(), unseen_instance_classes)
    # unseen_class_objs = []
    # while len(unseen_class_objs) != 15:
    #     unseen_classes = np.random.choice(rest_classes, 7, replace=False)
    #     unseen_class_objs = df[df['type'].isin(unseen_classes)]['name'].values
    unseen_class = "cup"
    unseen_class_objs = df[df['type'] == unseen_class]['name'].values
    
    # Train: the rest
    test_objs = np.concatenate([unseen_instances, unseen_class_objs])
    train_objs = np.setdiff1d(df['name'].values, test_objs)

    data = {
        "train": train_objs.tolist(),
        "unseen_instance": unseen_instances,
        "unseen_class": unseen_class_objs.tolist()
    }
    print(f"Dataset split: {len(train_objs)} train, {len(unseen_instances)} unseen instance, {len(unseen_class_objs)} unseen class")

    with open(os.path.join(root_dir, 'dataset_split.json'), 'w') as f:
        json.dump(data, f, indent=4)

    return

def plot_type_histogram(df: pd.DataFrame):
    """
    Plot the histogram of number of objects per type
    """
    plt.figure(figsize=(10, 5))
    sns.countplot(x='type', data=df)
    plt.title('Number of objects per type')
    # plt.savefig(os.path.join(root_dir, 'type_histogram.png'))
    plt.show()
    return
    

def generate_unseen_instance_split(df: pd.DataFrame):
    # per type
    all_types = df['type'].unique()
    test_set = []
    train_set = []
    
    for obj_type in all_types:
        obj_names = df[df['type'] == obj_type]['name'].values
        test_objs = np.random.choice(obj_names, 1, replace=False)
        train_objs = np.setdiff1d(obj_names, test_objs)
        
        test_set.extend(test_objs)
        train_set.extend(train_objs)
    
    data = {
        "train": train_set,
        "test": test_set
    }
    print(f"Unseen instance split: {len(test_set)} test, {len(train_set)} train")
    with open(os.path.join(root_dir, 'unseen_instance_split.json'), 'w') as f:
        json.dump(data, f, indent=2)
    
    return

def generate_unseen_type_split(df: pd.DataFrame, types=['cup']):
    all_types = df['type'].unique()
    if types is not None:
        test_types = types
    else:
        num_test_type = int(np.ceil(len(all_types) * 0.2)) 
        test_types = np.random.choice(all_types, num_test_type, replace=False)
    
    test_set = df[df['type'].isin(test_types)]['name'].values
    train_set = np.setdiff1d(df['name'].values, test_set)
    
    data = {
        "train": train_set.tolist(),
        "test": test_set.tolist()
    }
    print(f"Unseen type split: {len(test_set)} test, {len(train_set)} train")
    with open(os.path.join(root_dir, 'unseen_type_split.json'), 'w') as f:
        json.dump(data, f, indent=2)
    return
        

def generated_mixed_split(df: pd.DataFrame):
    num_obj = len(df)
    num_test_obj = int(np.ceil(num_obj * 0.2))
    
    test_set = np.random.choice(df['name'].values, num_test_obj, replace=False)
    train_set = np.setdiff1d(df['name'].values, test_set)
    
    data = {
        "train": train_set.tolist(),
        "test": test_set.tolist()
    }
    print(f"Mixed split: {len(test_set)} test, {len(train_set)} train")
    with open(os.path.join(root_dir, 'mixed_split.json'), 'w') as f:
        json.dump(data, f, indent=2)
    return

def test_read():
    df = pd.read_csv(os.path.join(root_dir, 'metainfo.csv'))
    print(df['name'].values)

# pca = PCA()
# pca.fit(points)
# eign_vec = pca.components_
# pcd_extra = o3d.geometry.PointCloud()
# # origin = np.mean(points, axis=0).reshape(1,-1)
# # print(origin)
# # print(data['object/com'][()])
# extra_points = []
# extra_points.append(np.array([[max_points[0], 0, 0]]))
# extra_points.append(np.array([[min_points[0], 0, 0]]))
# # for i in range(10):
# #     # extra_points.append(origin+eign_vec*i*0.01)
#
# print(eign_vec)
#
# extra_points = np.concatenate(extra_points)
# pcd_extra.points = o3d.utility.Vector3dVector(extra_points)
#
# o3d.visualization.draw_geometries([pcd_scaled, pcd_extra])

if __name__ == "__main__":
    # process_objs()
    df = generate_metadata()
    plot_type_histogram(df)
    generate_dataset_split(df)
    # generate_unseen_instance_split(df)
    # generate_unseen_type_split(df)
    # generated_mixed_split(df)
    test_read()

bin_env/assets/housekeep/process.py
HACMan-2023/HACMan
import open3d as o3d
import os
import h5py
from sklearn.decomposition import PCA
import numpy as np
import pandas as pd
import os
import json
from tqdm import tqdm
from collections import defaultdict

import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme(style="white")

root_dir = os.path.dirname(os.path.abspath(__file__))
dirname = 'bin_env/assets/housekeep_all/models'
processed_dirname = 'bin_env/assets/housekeep_all/models_processed/meshes'
info_dirname = 'bin_env/assets/housekeep_all/object_info'
v_hacd_path = '../v-hacd/app/build/TestVHACD'
manifold_path = '../Manifold/build/manifold'

def process_objs():
    for obj_dir in tqdm(os.listdir(dirname)):
        if obj_dir[-2:] != '_M':
            continue
        
        obj_type = obj_dir.split('_')[0]
        if obj_type in ['fork', 'knife', 'scissors', 'spoon', 'screwdriver']:
            continue
        # if obj_file[-4:] != '.obj':
        #     continue
        
        print(obj_dir)
        obj_path = os.path.join(dirname, obj_dir, 'model.obj')
        
        # v-hacd
        # v_hacd_dir = os.path.dirname(v_hacd_path)
        os.system(f'{v_hacd_path} {obj_path}')
        
        # manifold
        os.system(f'{manifold_path} decomp.obj decomp.obj 20000')
        
        mesh = o3d.io.read_triangle_mesh("decomp.obj", True)
        # mesh = o3d.io.read_triangle_mesh(obj_path, True)
        mesh.compute_vertex_normals()
        # pcd = o3d.geometry.sample_points_uniformly(mesh, number_of_points=10000)
        # points = np.asarray(pcd.points)
        # pcd_scaled = o3d.geometry.PointCloud()
        # pcd_scaled.points = o3d.utility.Vector3dVector(points)
        vertices = np.asarray(mesh.vertices)
        print(vertices.shape)
        print("Water tight? ", mesh.is_watertight())

        max_points = np.max(vertices, axis=0)
        min_points = np.min(vertices, axis=0)
        print(max_points)
        print(min_points)
        size = (max_points-min_points)/2
        origin = (max_points+min_points)/2
        
        stl_file = obj_dir + '.stl'
        stl_path = os.path.join(processed_dirname, stl_file)
        # o3d.visualization.draw_geometries([mesh])
        o3d.io.write_triangle_mesh(stl_path, mesh)
        
        for f in ["decomp.mtl", "decomp.obj", "decomp.stl"]:
            if os.path.exists(f):
                os.remove(f)
        
        filename = stl_file[:-4] +".h5"
        filepath = os.path.join(info_dirname, filename)
        with h5py.File(filepath, 'w') as f:
            f.create_dataset("origin", data=origin)
            f.create_dataset("size", data=size)

def generate_metadata():
    unstable_objs = [
        "handheld_game_console_RedBlack_Nintendo_3DSXL_M",
        "hat_DPC_Handmade_Hat_Brown_M",
        "can_opener_OXO_Soft_Works_Can_Opener_SnapLock_M",
        "tray_Design_Ideas_Drawer_Store_Organizer_M",
        "lego_block_073-c_lego_duplo_M",
        "keyboard_Kanex_MultiSync_Wireless_Keyboard_M",
        "clock_frl_apartment_clock_M",
        "hard_drive_Seagate_1TB_Backup_Plus_portable_drive_Blue_M",
        "lego_block_073-a_lego_duplo_M",
        "mouse_Razer_Naga_MMO_Gaming_Mouse_M",
        "box_Little_Debbie_Donut_Sticks_6_cake_donuts_10_oz_total_M",
        "keyboard_Razer_Blackwidow_Tournament_Edition_Keyboard_M",
        "magnifying_glass_Magnifying_Glassassrt_M",
        "pitcher_019_pitcher_base_M",
        "book_frl_apartment_book_02_M",
        "bowl_kitchen_set_kitchen_set_soup_bowl_M",
        "hair_straightener_Remington_TStudio_Silk_Ceramic_Hair_Straightener_2_Inch_Floating_Plates_M",
        "tray_Ecoforms_Saucer_SQ3_Turquoise_M",
        ]
    unstable_types = [
        "plate"
    ]
    
    metadata = {
        'name': [],
        'type': [],
    }
    
    for h5_file in os.listdir(info_dirname):
        if h5_file[-3:] != '.h5':
            continue
        
        obj_name = h5_file[:-3]
        obj_type = obj_name.split('_')[0]
        
        if obj_name in unstable_objs or obj_type in unstable_types:
            continue

        metadata['name'].append(obj_name)
        metadata['type'].append(obj_type)
        
    df = pd.DataFrame(metadata)
    df.to_csv(os.path.join(root_dir, 'metainfo.csv'))
    return df


def plot_type_histogram(df: pd.DataFrame):
    """
    Plot the histogram of number of objects per type
    """
    plt.figure(figsize=(10, 5))
    sns.countplot(y='type', data=df,
                  order = df['type'].value_counts().index,
                  palette='rocket')
    plt.title('Number of objects per type')
    # plt.savefig(os.path.join(root_dir, 'type_histogram.png'))
    plt.show()
    return

def generate_dataset_split(df: pd.DataFrame):
    # Split into train, unseen instance, unseen class

    # Unseen instance: 10 instances, 1 instance from each classes with >= 3 instances
    unseen_instance_classes = np.random.choice(df['type'].value_counts()[df['type'].value_counts() >= 3].index, 10, replace=False)
    unseen_instances = []
    for obj_type in unseen_instance_classes:
        obj_names = df[df['type'] == obj_type]['name'].values
        unseen_instances.extend(np.random.choice(obj_names, 1, replace=False))

    # Unseen class: 15 objects from the rest of the classes
    rest_classes = np.setdiff1d(df['type'].unique(), unseen_instance_classes)
    unseen_class_objs = []
    while len(unseen_class_objs) != 15:
        unseen_classes = np.random.choice(rest_classes, 7, replace=False)
        unseen_class_objs = df[df['type'].isin(unseen_classes)]['name'].values
    
    # Train: the rest
    test_objs = np.concatenate([unseen_instances, unseen_class_objs])
    train_objs = np.setdiff1d(df['name'].values, test_objs)

    data = {
        "train": train_objs.tolist(),
        "unseen_instance": unseen_instances,
        "unseen_class": unseen_class_objs.tolist()
    }
    print(f"Dataset split: {len(train_objs)} train, {len(unseen_instances)} unseen instance, {len(unseen_class_objs)} unseen class")

    with open(os.path.join(root_dir, 'dataset_split.json'), 'w') as f:
        json.dump(data, f, indent=4)

    return

def generate_unseen_instance_split(df: pd.DataFrame):
    # per type
    all_types = df['type'].unique()
    test_set = []
    train_set = []
    
    for obj_type in all_types:
        obj_names = df[df['type'] == obj_type]['name'].values
        test_objs = np.random.choice(obj_names, 1, replace=False)
        train_objs = np.setdiff1d(obj_names, test_objs)
        
        test_set.extend(test_objs)
        train_set.extend(train_objs)
    
    data = {
        "train": train_set,
        "test": test_set
    }
    print(f"Unseen instance split: {len(test_set)} test, {len(train_set)} train")
    with open(os.path.join(root_dir, 'unseen_instance_split.json'), 'w') as f:
        json.dump(data, f, indent=2)
    
    return

def generate_unseen_type_split(df: pd.DataFrame, types=None):
    all_types = df['type'].unique()
    if types is not None:
        test_types = types
    else:
        num_test_type = int(np.ceil(len(all_types) * 0.2)) 
        test_types = np.random.choice(all_types, num_test_type, replace=False)
    
    test_set = df[df['type'].isin(test_types)]['name'].values
    train_set = np.setdiff1d(df['name'].values, test_set)
    
    data = {
        "train": train_set.tolist(),
        "test": test_set.tolist()
    }
    print(f"Unseen type split: {len(test_set)} test, {len(train_set)} train")
    with open(os.path.join(root_dir, 'unseen_type_split.json'), 'w') as f:
        json.dump(data, f, indent=2)
    return
        

def generated_mixed_split(df: pd.DataFrame):
    num_obj = len(df)
    num_test_obj = int(np.ceil(num_obj * 0.2))
    
    test_set = np.random.choice(df['name'].values, num_test_obj, replace=False)
    train_set = np.setdiff1d(df['name'].values, test_set)
    
    data = {
        "train": train_set.tolist(),
        "test": test_set.tolist()
    }
    print(f"Mixed split: {len(test_set)} test, {len(train_set)} train")
    with open(os.path.join(root_dir, 'mixed_split.json'), 'w') as f:
        json.dump(data, f, indent=2)
    return

def test_read():
    df = pd.read_csv(os.path.join(root_dir, 'metainfo.csv'))
    print(df['name'].values)

def read_unstable_objects():
    unstable_objects = defaultdict(dict)
    unstable_object_files = [
        "unstable_objects_fixed_size.json",
        "unstable_objects_variable_size.json",
        "unstable_objects_variable_sizes_with_wall.json"]

    root_dir = "data/housekeep_all"
    for file in unstable_object_files:
        with open(os.path.join(root_dir, file), 'r') as f:
            data = json.load(f)
        
        for key, value in data.items():
            if value >= 12:
                unstable_objects[key][file] = value
    
    # Calculate the percentage of unstable poses for each object
    for obj, unstable_poses in unstable_objects.items():
        total_count = 0
        for mode, count in unstable_poses.items():
            total_count += count
        
        perc = total_count / 300
        unstable_objects[obj]['percentage'] = perc
        if perc >= 0.1:
            print(f'"{obj}",')
    
    # Export to json
    with open(os.path.join(root_dir, "all_unstable_objects.json"), 'w') as f:
        json.dump(unstable_objects, f, indent=2)
    
    print(f"Total unstable objects: {len(unstable_objects)}")
    
    return

def add_thick_splits():
    # Read the flat objects csv
    thick_obj_lookup = pd.read_csv(os.path.join(root_dir, 'flat_objects.csv'))
    thick_obj_lookup = thick_obj_lookup.set_index('name')
    
    # Read the current splits
    with open(os.path.join(root_dir, 'dataset_split.json'), 'r') as f:
        data = json.load(f)
    
    # Add the flat objects to the splits
    thick_splits = {}
    for split in data.keys():
        if split.endswith('_thick'):
            continue
        objects = data[split]
        thick_objects = []

        for obj in objects:
            if thick_obj_lookup['thick'][obj] == 'True':
                thick_objects.append(obj)
        
        split_name = f"{split}_thick"
        thick_splits[split_name] = thick_objects

        print(f"{split_name}: {len(thick_objects)}")
    
    # Export the new splits
    data.update(thick_splits)
    with open(os.path.join(root_dir, 'dataset_split.json'), 'w') as f:
        json.dump(data, f, indent=2)


# pca = PCA()
# pca.fit(points)
# eign_vec = pca.components_
# pcd_extra = o3d.geometry.PointCloud()
# # origin = np.mean(points, axis=0).reshape(1,-1)
# # print(origin)
# # print(data['object/com'][()])
# extra_points = []
# extra_points.append(np.array([[max_points[0], 0, 0]]))
# extra_points.append(np.array([[min_points[0], 0, 0]]))
# # for i in range(10):
# #     # extra_points.append(origin+eign_vec*i*0.01)
#
# print(eign_vec)
#
# extra_points = np.concatenate(extra_points)
# pcd_extra.points = o3d.utility.Vector3dVector(extra_points)
#
# o3d.visualization.draw_geometries([pcd_scaled, pcd_extra])

if __name__ == "__main__":
    # process_objs()
    # df = generate_metadata()
    # plot_type_histogram(df)
    # generate_dataset_split(df)
    # generate_unseen_instance_split(df)
    # generate_unseen_type_split(df)
    # generated_mixed_split(df)
    add_thick_splits()
    # test_read()
    # read_unstable_objects()

bin_env/assets/housekeep_all/process.py
HACMan-2023/HACMan
import os
import open3d as o3d
import math
import json

from PIL import Image
from tqdm import tqdm

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = os.path.join(ROOT_DIR, 'models_processed', 'meshes') 
IMAGE_DIR = os.path.join(ROOT_DIR, 'images')

def vis_object(obj_name):
    """Visualize an object mesh and save it to a PNG file"""
    obj_file = os.path.join(MODEL_DIR, obj_name + '.stl')
    mesh = o3d.io.read_triangle_mesh(obj_file)
    mesh.compute_vertex_normals()

    # Create a visualization settings object and set the desired view angle
    vis = o3d.visualization.VisualizerWithEditing()
    vis.create_window(width=400, height=400, visible=True)
    vis.add_geometry(mesh)
    vis.get_view_control().rotate(45.0, 45.0, 45.0)

    # Start the visualizer and allow user to interactively adjust the view
    vis.run()

    vis.update_renderer()
    vis.capture_screen_image(os.path.join(IMAGE_DIR, obj_name + '.png'), do_render=True)
    vis.destroy_window()

def generate_images():
    """Generate PNG images for all STL files in the model folder"""
    print("Generating images...")
    for obj_name in tqdm(os.listdir(MODEL_DIR)):
        if obj_name.endswith('.stl'):
            obj_name = os.path.splitext(obj_name)[0]
            vis_object(obj_name)

def align_images(split=None):
    """Align all PNG images into a single picture"""
    images = []
    split_file = os.path.join(ROOT_DIR, "dataset_split.json")
    with open(split_file, 'r') as f:
        data = json.load(f)
        if split is not None:
            object_list = data[split]
        else:
            object_list = data["train"] + data["unseen_instance"] + data["unseen_class"]

    for obj_name in object_list:
        image_path = os.path.join(IMAGE_DIR, obj_name + '.png')
        if os.path.exists(image_path):
            image = Image.open(image_path)
            images.append(image)

    if images:
        # Calculate the number of rows and columns based on the aspect ratio
        num_images = len(images)
        aspect_ratio = 16 / 9
        num_columns = int(math.sqrt(num_images / aspect_ratio))
        num_rows = int(math.ceil(num_images / num_columns))
        
        # Resize all images to the same height and calculate the width based on the aspect ratio
        target_height = 300
        target_width = int(target_height * aspect_ratio)
        resized_images = []
        for image in images:
            width, height = image.size
            ratio = target_height / height
            resized_image = image.resize((int(width * ratio), target_height))
            resized_images.append(resized_image)
        
        # Create a new image with the appropriate size
        total_width = num_columns * target_width
        total_height = num_rows * target_height
        result = Image.new('RGB', (total_width, total_height), color='white')
        
        # Paste the images into the new image
        index = 0
        for row in range(num_rows):
            for col in range(num_columns):
                if index >= num_images:
                    break
                x_offset = col * target_width
                y_offset = row * target_height
                result.paste(resized_images[index], (x_offset, y_offset))
                index += 1
        
        if split:
            img_name = 'objects_' + split + '.png'
        else:
            img_name = 'objects_all.png'
        result.save(os.path.join(ROOT_DIR, img_name))
    else:
        print('No PNG images found')

# generate_images()
# align_images()
align_images('train')
align_images('unseen_instance')
align_images('unseen_class')

# vis_object('box_Just_For_Men_Mustache_Beard_Brushin_Hair_Color_Gel_Kit_Jet_Black_M60_M')
# vis_object("cup_BIA_Cordon_Bleu_White_Porcelain_Utensil_Holder_900028_M")
# vis_object("pill_bottle_Beta_Glucan_M")
# vis_object("rubiks_cube_077_rubiks_cube_M")
# vis_object("pencil_case_Wishbone_Pencil_Case_M")

# vis_object("book_frl_apartment_book_03_M")
# vis_object("clamp_Pony_C_Clamp_1440_M")
# vis_object("tray_Threshold_Tray_Rectangle_Porcelain_M")
# vis_object("bowl_kitchen_set_kitchen_set_bowl_M")
# vis_object("handheld_game_console_Nintendo_2DS_Crimson_Red_M")

bin_env/assets/housekeep_all/vis_object.py
HACMan-2023/HACMan
from collections import OrderedDict
import numpy as np
import pandas as pd

from robosuite.models.objects import MujocoXMLObject
from robosuite.utils.mjcf_utils import xml_path_completion, array_to_string

import h5py
import tempfile
import json
import shutil
import os, random
from lxml import etree

from bin_env.assets.objects import MujocoMeshObject

curr_path = os.path.dirname(os.path.abspath(__file__))

class HousekeepSampler:
    def __init__(self, use_full_set=False, object_scale_range=None, object_size_limit=[0.05, 0.05], **kwargs):
        if use_full_set:
            self.housekeep_root = os.path.join(curr_path, 'housekeep_all')
        else:
            self.housekeep_root = os.path.join(curr_path, 'housekeep')
        
        self.objects_root = os.path.join(self.housekeep_root, 'models_processed')
        self.object_info_root = os.path.join(self.housekeep_root, 'object_info')
        self.texture_root = os.path.join(curr_path, 'textures')
            
        self.object_set, self.all_object_info = self.load_all_objects(**kwargs)
        self.object_scale_range = object_scale_range
        self.object_size_limit = object_size_limit


    def load_all_objects(self, object_types=None, object_name=None, object_split=None, **kwargs):
        object_df = pd.read_csv(os.path.join(self.housekeep_root, 'metainfo.csv'))
        object_set = []
        all_object_info = {}
        
        # Found the set of specified objects
        if object_name is None:            
            # Sample from a split configuration
            if object_split is not None and (object_split != 'no_split'):
                with open(os.path.join(self.housekeep_root, 'dataset_split.json')) as f:
                    split = json.load(f)
                object_set = split[object_split]
            
            # Sample from specific types
            elif object_types is not None:
                object_set = object_df[object_df['type'].isin(object_types)]['name'].tolist()
            
            # Sample from all objects
            else:
                object_set = object_df['name'].values.tolist()
            
        else:
            assert object_name in object_df['name'].values.tolist(), f"Specified object {object_name} does not exist."
            object_set = [object_name]
            
        assert len(object_set) > 0, "No objects found in the specified set."
        
        # Load all object info
        for object_name in object_set:
            info_path = os.path.join(self.object_info_root, object_name+'.h5')
            assert os.path.exists(info_path), f"Specified object {object_name} does not exist."
            all_object_info[object_name] = h5py.File(info_path, "r")

        print(f"Loaded {len(object_set)} objects from {self.housekeep_root}.")
        
        return object_set, all_object_info
    

    def set_object(self, object_name):
        assert object_name in self.all_object_info.keys(), f"Specified object {object_name} does not exist."
        self.object_set = [object_name]
    

    def get_object_set(self):
        return self.object_set

    def generate_housekeep_object(self):
        object_info = self.sample_housekeep_object()

        obj = MujocoMeshObject(object_info=object_info, 
                            name="object", 
                            stl_dir=self.objects_root,
                            texture_dir=self.texture_root, 
                            joints=[dict(type="free", damping="0.0005")], 
                            obj_type="all", 
                            duplicate_collision_geoms=True, 
                            size_limit=self.object_size_limit,
                            object_scale_range=self.object_scale_range)
        
        obj_target = MujocoMeshObject(object_info=object_info, 
                                    name="object_target", 
                                    stl_dir=self.objects_root, 
                                    texture_dir=self.texture_root, 
                                    joints=None, 
                                    obj_type="visual",
                                    duplicate_collision_geoms=True, 
                                    size_limit=self.object_size_limit,
                                    object_scale_range=self.object_scale_range,
                                    object_scale=obj.scale)
        
        return obj, obj_target

    def sample_housekeep_object(self):        
        object_name = np.random.choice(self.object_set).item()
        object_info = self.all_object_info[object_name]

        return object_info

# class HousekeepObject(MujocoXMLObject):
#     """
#     Loading Acronym Objects and Grasps
#     """
#     def __init__(self, name=None, joints=None, object_info=None, size=None, density=None, obj_type='all', pos_offset=None):
#         # filename = object_info.filename.split('/')[-1].split('_')[1] + ".h5"
#         # additional_data = h5py.File(os.path.join(ACRONYM_OBJECT_INFO_ROOT, filename), "r")
#         self.temp_dir = tempfile.TemporaryDirectory(dir=HOUSEKEEP_OBJECTS_ROOT)
#         self.origin = object_info['origin'][()]
#         self.original_size = object_info['size'][()]
#         self.size = size # np.array([0.05, 0.05, 0.02])
#         self.density = density
#         self.pos_offset = pos_offset
#         self.mesh_name = object_info.filename.split('/')[-1].rstrip('.h5')
        
#         self.xml_path = self.parse_object_info(object_info)
#         super().__init__(self.xml_path, name=name, joints=joints, obj_type=obj_type, duplicate_collision_geoms=True)

#     def __del__(self):
#         self.temp_dir.cleanup()
#         try:
#             if os.path.exists(self.temp_dir.name):
#                 shutil.rmtree(self.temp_dir.name)
#         finally:
#             del self.temp_dir

#     def parse_object_info(self, data):
#         orig_size_norm = np.max(self.original_size)
#         # orig_size_norm = np.sqrt(np.sum(self.original_size ** 2))
#         # orig_size_norm = np.sum(self.original_size)
#         object_scale = np.ones(3) / orig_size_norm * self.size
#         object_com = self.origin * object_scale
#         if self.pos_offset is not None:
#             object_com += self.pos_offset

#         # Save them for other process to read
#         self.scale = object_scale
#         self.scaled_origin = object_com
#         self.scaled_size = self.original_size * object_scale

#         # # Using the scale from acronym
#         # object_scale = np.ones(3) * data['object/scale'][()]
#         # # object_com = self.origin * data['object/scale'][()]
#         # object_com = data["object/com"][()]

#         object_quat = np.array([1, 0, 0, 0])
#         # # Change the orientation according to the shape in the order of x>y>z # BUG
#         # size = self.size.copy()
#         # original_size = self.original_size.copy()
#         # print(original_size)
#         # if original_size[1] < original_size[2]:
#         #     object_quat = quat_mul(euler2quat(np.array([np.pi/2, 0, 0])), object_quat)
#         #     original_size[2], original_size[1] = original_size[1], original_size[2]
#         #     size[2], size[1] = size[1], size[2]
#         # if original_size[0] < original_size[2]:
#         #     object_quat = quat_mul(euler2quat(np.array([0, np.pi/2, 0])), object_quat)
#         #     object_quat = quat_mul(euler2quat(np.array([0, 0, np.pi/2])), object_quat)
#         #     original_size[0], original_size[1], original_size[2] = original_size[1], original_size[2], original_size[0]
#         #     size[0], size[1], size[2] = size[1], size[2], size[0]
#         # elif original_size[0] < original_size[1]:
#         #     object_quat = quat_mul(euler2quat(np.array([0, 0, np.pi/2])), object_quat)
#         #     original_size[0], original_size[1] = original_size[1], original_size[0]
#         #     size[0], size[1] = size[1], size[0]
#         # print(original_size)
#         # object_scale = np.ones(3) / self.original_size * size
#         # object_com = self.origin / self.original_size * size
#         # object_com = quat_rot_vec_arr(object_quat, object_com)

#         filename = data.filename.split('/')[-1]
#         object_type = filename.split('_')[0]
#         object_id = filename.split('.')[0]
#         object_texture = os.path.join(TEXTURE_ROOT, "ceramic.png")
#         # object_mass = str(data["object/mass"][()])
#         object_density = str(self.density) #str(data["object/density"][()])
#         # object_friction = str(data["object/friction"][()])
#         object_friction = "1.0"
#         xml_path = self.generate_xml(object_id, object_type, object_scale, object_texture, 
#                                      object_density, object_friction, object_com, object_quat)
#         return xml_path

#     def generate_xml(self, object_id, object_type, object_scale, object_texture, object_density,
#                      object_friction, object_com, object_quat):
#         object_xml = object_id + '.xml'
#         # object_xml_path = os.path.join(HOUSEKEEP_OBJECTS_ROOT, object_xml)
#         object_xml_path = os.path.join(self.temp_dir.name, object_xml)
#         if True: # not os.path.exists(object_xml_path): # Always generate new files because scale and mass might change
#             root = etree.Element("mujoco")
#             root.set("model", object_id)

#             asset = etree.SubElement(root, "asset")
#             mesh = etree.SubElement(asset, "mesh")
#             mesh.set("file", "../meshes/" + object_id + ".stl")
#             mesh.set("name", object_type + "_mesh")
#             mesh.set("scale", str(object_scale[0]) + " " + str(object_scale[1]) + " " + str(object_scale[2]))
            
#             texture = etree.SubElement(asset, "texture")
#             texture.set("type", "cube")
#             texture.set("file", object_texture)
#             texture.set("rgb1", "1 1 1")
#             texture.set("name", "tex-" + object_type)
#             material = etree.SubElement(asset, "material")
#             material.set("name", object_type)
#             material.set("reflectance", "0.5")
#             material.set("texrepeat", "3 3")
#             material.set("texture", "tex-" + object_type)

#             # TODO: check XML (such as mass values)
#             worldbody = etree.SubElement(root, "worldbody")
#             body = etree.SubElement(worldbody, "body")
#             collision = etree.SubElement(body, "body")
#             collision.set("name", "object")
#             geom = etree.SubElement(collision, "geom")
#             geom.set("pos", str(-object_com[0])+" "+str(-object_com[1])+" "+str(-object_com[2]))
#             # geom.set("pos", str(object_com[0])+" "+str(object_com[1])+" "+str(-object_com[2]))
#             # geom.set("pos", str(object_com[0])+" "+str(object_com[1])+" "+str(0))
#             geom.set("quat", str(object_quat[0])+" "+str(object_quat[1])+" "+str(object_quat[2])+" "+str(object_quat[3]))
#             # geom.set("pos", "0 0 0")
#             geom.set("mesh", object_type + "_mesh")
#             geom.set("type", "mesh")
#             geom.set("density", object_density)

#             # Mujoco default
#             # geom.set("solimp", "0.9 0.95 0.001")
#             # geom.set("solref", "0.02 1")
#             # geom.set("friction", object_friction + " " + object_friction + " " + object_friction)
            
#             # Robosuite xml object settings
#             # geom.set("solimp", "0.998 0.998 0.001")
#             # geom.set("solref", "0.001 1")
#             # geom.set("density", "50")
#             # geom.set("friction", "0.95 0.3 0.1")
            
#             geom.set("group", "0")
#             geom.set("condim", "4")
#             geom.set("material", object_type)

#             # visual = etree.SubElement(body, "body")
#             # visual.set("name", "visual")
#             # geom1 = etree.SubElement(visual, "geom")
#             # geom1.set("pos", "0 0 0")
#             # geom1.set("mesh", object_type + "_mesh")
#             # geom1.set("type", "mesh")
#             # geom1.set("material", object_type)
#             # geom1.set("conaffinity", "0")
#             # geom1.set("contype", "0")
#             # geom1.set("group", "0")
#             # geom1.set("mass", "0.1")

#             # geom2 = etree.SubElement(visual, "geom")
#             # geom2.set("pos", "0 0 0")
#             # geom2.set("mesh", object_type + "_mesh")
#             # geom2.set("type", "mesh")
#             # geom2.set("material", object_type)
#             # geom2.set("conaffinity", "0")
#             # geom2.set("contype", "0")
#             # geom2.set("group", "1")
#             # geom2.set("mass", "0.1")

#             # TODO: set the sites properly
#             bottom_site = etree.SubElement(body, "site")
#             bottom_site.set("rgba", "0 0 0 0")
#             bottom_site.set("size", "0.005")
#             bottom_site.set("pos", "0 0 -0.05")
#             bottom_site.set("name", "bottom_site")
#             top_site = etree.SubElement(body, "site")
#             top_site.set("rgba", "0 0 0 0")
#             top_site.set("size", "0.005")
#             top_site.set("pos", "0 0 0.03")
#             top_site.set("name", "top_site")
#             horizontal_radius_site = etree.SubElement(body, "site")
#             horizontal_radius_site.set("rgba", "0 0 0 0")
#             horizontal_radius_site.set("size", "0.005")
#             horizontal_radius_site.set("pos", "0.02 0.015 0")
#             horizontal_radius_site.set("name", "horizontal_radius_site")

#             et = etree.ElementTree(root)
#             # if object_xml_path is not None and os.path.exists(object_xml_path):
#             #     os.remove(object_xml_path)
#             et.write(object_xml_path, pretty_print=True)
#         return object_xml_path
    

if __name__ == '__main__':
    sampler = HousekeepSampler(object_types='all', object_name=None, object_split=None, object_eval_set=False)
    for _ in range(15):
        obj, target_obj = sampler.generate_housekeep_object()
        print(obj.mesh_name)
bin_env/assets/housekeep_object.py
HACMan-2023/HACMan
import os
import numpy as np
import random
import logging
import xml.etree.ElementTree as ET
from robosuite.models.objects.objects import MujocoXMLObject, GEOM_GROUPS
    
    
class MujocoTreeObject(MujocoXMLObject):
    """
    Initialize a mujoco object without creating the xml file
    """
    def __init__(self, tree, name, root_dir=None, joints="default", obj_type="all", duplicate_collision_geoms=True):
        self.init_mujocoxml(tree, root_dir)
        self.init_mujocoxmlobject(name, joints, obj_type, duplicate_collision_geoms)

    def init_mujocoxml(self, tree, folder=None):
        """
        __init__ from MujocoXML
        """
        # Root dir to be used for any relative path in the xml file
        self.folder = folder
        
        self.tree = tree
        
        self.root = self.tree.getroot()
        self.worldbody = self.create_default_element("worldbody")
        self.actuator = self.create_default_element("actuator")
        self.sensor = self.create_default_element("sensor")
        self.asset = self.create_default_element("asset")
        self.tendon = self.create_default_element("tendon")
        self.equality = self.create_default_element("equality")
        self.contact = self.create_default_element("contact")

        # Parse any default classes and replace them inline
        default = self.create_default_element("default")
        default_classes = self._get_default_classes(default)
        self._replace_defaults_inline(default_dic=default_classes)

        # Remove original default classes
        self.root.remove(default)

        self.resolve_asset_dependency()
        
    def init_mujocoxmlobject(self, name, joints="default", obj_type="all", duplicate_collision_geoms=True):
        """
        __init__ from MujocoXMLObject
        """
        
        # Set obj type and duplicate args
        assert obj_type in GEOM_GROUPS, "object type must be one in {}, got: {} instead.".format(GEOM_GROUPS, obj_type)
        self.obj_type = obj_type
        self.duplicate_collision_geoms = duplicate_collision_geoms

        # Set name
        self._name = name

        # joints for this object
        if joints == "default":
            self.joint_specs = [self.get_joint_attrib_template()]  # default free joint
        elif joints is None:
            self.joint_specs = []
        else:
            self.joint_specs = joints

        # Make sure all joints have names!
        for i, joint_spec in enumerate(self.joint_specs):
            if "name" not in joint_spec:
                joint_spec["name"] = "joint{}".format(i)

        # Lastly, parse XML tree appropriately
        self._obj = self._get_object_subtree()

        # Extract the appropriate private attributes for this
        self._get_object_properties()


class MujocoMeshObject(MujocoTreeObject):
    def __init__(self, object_info, name, stl_dir=None, texture_dir=None, joints="default", obj_type="all", duplicate_collision_geoms=True, 
                 size_limit=None, xml_callback=None, object_scale_range=(0.8, 1.2), object_scale=None):
        self.size_limit = np.array(size_limit) if size_limit is not None else None
        self.obj_type = obj_type
        self.stl_dir = stl_dir
        if object_scale_range is None:
            self.scales = np.array([1.0])
        else:
            self.scales = np.arange(object_scale_range[0], object_scale_range[1]+0.001, 0.1).round(2)

        self.mesh_name = object_info.filename.split('/')[-1].rstrip('.h5')
        self.texture_dir = texture_dir
        self.xml_callback = xml_callback # Additional changes to the xml file

        et = self.parse_object_info(object_info, object_scale=object_scale)
        super().__init__(et, name=name, root_dir=stl_dir, joints=joints, obj_type=obj_type, duplicate_collision_geoms=duplicate_collision_geoms)

    def parse_object_info(self, object_info, object_scale=None):
        origin = object_info['origin'][()]
        original_size = object_info['size'][()]
        filename = object_info.filename.split('/')[-1]
        object_type = filename.split('_')[0]
        object_id = filename.split('.')[0]
        object_texture = os.path.join(self.texture_dir, "color.png")
        
        if object_scale is not None:
            pass

        elif self.size_limit is None:
            random_scale = random.choice(self.scales)
            object_scale = np.ones(3) * random_scale

        elif self.size_limit.shape == (2,):
            # Limit the size of the object while keeping the shape
            orig_size_norm = np.max(original_size)
            object_scale = np.ones(3)
            if orig_size_norm > self.size_limit[1]:
                object_scale *= self.size_limit[1] / orig_size_norm
            elif orig_size_norm < self.size_limit[0]:
                object_scale *= self.size_limit[0] / orig_size_norm

            # Multiplies the scale by a random factor
            random_scale = random.choice(self.scales)
            object_scale *= random_scale

        # elif self.size_limit.shape == (3,):
        #     # Scale the object in x,y,z dimension separately
        #     object_scale = np.ones(3) / original_size * self.size_limit
        #     object_com = origin * object_scale
        #     object_dimension = original_size * object_scale

        else:
            raise NotImplementedError
        
        object_com = origin * object_scale
        object_dimension = original_size * object_scale

        # Save them for other process to read
        self.scale = object_scale
        self.scaled_origin = object_com
        self.scaled_size = object_dimension

        # Create Element Tree
        root = ET.Element("mujoco")
        root.set("model", object_id)

        asset = ET.SubElement(root, "asset")
        mesh = ET.SubElement(asset, "mesh")
        mesh.set("file", self.stl_dir + "/meshes/" + object_id + ".stl")
        mesh.set("name", object_type + "_mesh")
        mesh.set("scale", str(object_scale[0]) + " " + str(object_scale[1]) + " " + str(object_scale[2]))
        
        texture = ET.SubElement(asset, "texture")
        texture.set("file", object_texture)
        texture.set("type", "2d")
        texture.set("name", "tex-" + object_type)
        
        material = ET.SubElement(asset, "material")
        material.set("name", object_type)
        material.set("reflectance", "0.7")
        material.set("texrepeat", "15 15")
        material.set("texture", "tex-" + object_type)
        material.set("texuniform", "true")

        worldbody = ET.SubElement(root, "worldbody")
        body = ET.SubElement(worldbody, "body")
        collision = ET.SubElement(body, "body")
        collision.set("name", "object")
        geom = ET.SubElement(collision, "geom")
        geom.set("pos", str(-object_com[0])+" "+str(-object_com[1])+" "+str(-object_com[2]))
        geom.set("mesh", object_type + "_mesh")
        geom.set("type", "mesh")

        if self.obj_type == 'visual':
            geom.set("rgba", "0.8 0.8 0.8 0.3")
            geom.set("conaffinity", "0")
            geom.set("contype", "0")
            geom.set("group", "1")
            geom.set("mass", "0.0001")
            geom.set("material", object_type)
        else:
            # Robosuite xml object settings
            geom.set("solimp", "0.998 0.998 0.001")
            geom.set("solref", "0.001 1")
            geom.set("density", "50")
            geom.set("friction", "0.95 0.3 0.1") # Currently overwritten by table friction            
            geom.set("group", "0")
            geom.set("condim", "4")
            geom.set("material", object_type)

        bottom_site = ET.SubElement(body, "site")
        bottom_site.set("rgba", "1 0 0 1")
        bottom_site.set("size", "0.005")
        bottom_site.set("pos", "0 0 " + str(object_dimension[2]/2-0.1))
        bottom_site.set("name", "bottom_site")
        top_site = ET.SubElement(body, "site")
        top_site.set("rgba", "1 0 0 1")
        top_site.set("size", "0.005")
        top_site.set("pos", "0 0 " + str(-object_dimension[2]/2-0.1))
        top_site.set("name", "top_site")
        horizontal_radius_site = ET.SubElement(body, "site")
        horizontal_radius_site.set("rgba", "1 1 0 1")
        horizontal_radius_site.set("size", "0.005")
        horizontal_radius_site.set("pos", str(object_dimension[0]) + " " + str(object_dimension[1])+" 0")
        horizontal_radius_site.set("name", "horizontal_radius_site")
        
        if self.xml_callback is not None:
            # Additional changes to the xml file
            self.xml_callback(globals(), locals())

        et = ET.ElementTree(root)
        return et

bin_env/assets/objects.py
HACMan-2023/HACMan
from collections import OrderedDict
import numpy as np

from bin_env.assets import HousekeepSampler

from robosuite.utils.mjcf_utils import CustomMaterial
from robosuite.models.tasks import ManipulationTask
from robosuite.environments.manipulation.manipulation_env import ManipulationEnv
from robosuite.utils.placement_samplers import UniformRandomSampler
from robosuite.utils.observables import Observable, sensor
from bin_env.bin_arena import BinArena
from bin_env.single_arm import SingleArm
from robosuite.utils import transform_utils

import copy

class BaseEnv(ManipulationEnv):
    """
    This class corresponds to the lifting task for a single robot arm.

    Args:
        robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
            (e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
            Note: Must be a single single-arm robot!

        env_configuration (str): Specifies how to position the robots within the environment (default is "default").
            For most single arm environments, this argument has no impact on the robot setup.

        controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
            custom controller. Else, uses the default controller for this specific task. Should either be single
            dict if same controller is to be used for all robots or else it should be a list of the same length as
            "robots" param

        gripper_types (str or list of str): type of gripper, used to instantiate
            gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
            with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
            overrides the default gripper. Should either be single str if same gripper type is to be used for all
            robots or else it should be a list of the same length as "robots" param

        initialization_noise (dict or list of dict): Dict containing the  initialization noise parameters.
            The expected keys and corresponding value types are specified below:

            :`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
                joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
                If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
                If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
            :`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"

            Should either be single dict if same noise value is to be used for all robots or else it should be a
            list of the same length as "robots" param

            :Note: Specifying "default" will automatically use the default noise settings.
                Specifying None will automatically create the required dict with "magnitude" set to 0.0.

        table_full_size (3-tuple): x, y, and z dimensions of the table.

        table_friction (3-tuple): the three mujoco friction parameters for
            the table.

        table_offset (3-tuple): x, y, and z location of the table top with respect to the robot base.

        use_camera_obs (bool): if True, every observation includes rendered image(s)

        use_object_obs (bool): if True, include object (cube) information in
            the observation.

        reward_scale (None or float): Scales the normalized reward function by the amount specified.
            If None, environment reward remains unnormalized

        reward_shaping (bool): if True, use dense rewards.

        placement_initializer (ObjectPositionSampler): if provided, will
            be used to place objects on every reset, else a UniformRandomSampler
            is used by default.

        has_renderer (bool): If true, render the simulation state in
            a viewer instead of headless mode.

        has_offscreen_renderer (bool): True if using off-screen rendering

        render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
            will result in the default angle being applied, which is useful as it can be dragged / panned by
            the user using the mouse

        render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.

        render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.

        render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
            Defaults to -1, in which case the device will be inferred from environment variables
            (GPUS or CUDA_VISIBLE_DEVICES).

        control_freq (float): how many control signals to receive in every second. This sets the amount of
            simulation time that passes between every action input.

        horizon (int): Every episode lasts for exactly @horizon timesteps.

        ignore_done (bool): True if never terminating the environment (ignore @horizon).

        hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
            only calls sim.reset and resets all robosuite-internal variables

        camera_names (str or list of str): name of camera to be rendered. Should either be single str if
            same name is to be used for all cameras' rendering or else it should be a list of cameras to render.

            :Note: At least one camera must be specified if @use_camera_obs is True.

            :Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
                convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
                robot's camera list).

        camera_heights (int or list of int): height of camera frame. Should either be single int if
            same height is to be used for all cameras' frames or else it should be a list of the same length as
            "camera names" param.

        camera_widths (int or list of int): width of camera frame. Should either be single int if
            same width is to be used for all cameras' frames or else it should be a list of the same length as
            "camera names" param.

        camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
            bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
            "camera names" param.

        camera_segmentations (None or str or list of str or list of list of str): Camera segmentation(s) to use
            for each camera. Valid options are:

                `None`: no segmentation sensor used
                `'instance'`: segmentation at the class-instance level
                `'class'`: segmentation at the class level
                `'element'`: segmentation at the per-geom level

            If not None, multiple types of segmentations can be specified. A [list of str / str or None] specifies
            [multiple / a single] segmentation(s) to use for all cameras. A list of list of str specifies per-camera
            segmentation setting(s) to use.

    Raises:
        AssertionError: [Invalid number of robots specified]
    """

    def __init__(
        self,
        robots,
        initial_qpos=None,
        env_configuration="default",
        controller_configs=None,
        gripper_types="default",
        initialization_noise="default",
        table_full_size=(0.45, 0.54, 0.107),
        table_size_rand=(0., 0., 0.),
        table_friction=(0.3, 5e-3, 1e-4),
        table_solref=(0.02, 1),
        table_solimp=(0.9, 0.95, 0.001),
        table_offset=(0.5, 0, 0.065),
        table_xml="bin_arena.xml",
        use_camera_obs=True,
        use_object_obs=True,
        reward_scale=1.0,
        reward_shaping=False,
        placement_initializer=None,
        has_renderer=False,
        has_offscreen_renderer=True,
        render_camera="frontview",
        render_collision_mesh=False,
        render_visual_mesh=True,
        render_gpu_device_id=-1,
        control_freq=20,
        horizon=1000,
        ignore_done=False,
        hard_reset=True,
        camera_names="agentview",
        camera_heights=256,
        camera_widths=256,
        camera_depths=False,
        camera_segmentations=None,      # {None, instance, class, element}
        adaptive=False,
        additional_obs_keys=None,
        close_gripper=False,
        remove_walls=False,
        object_dataset="cube",     # {"cube", "acronym", "shapenet", "housekeep"}
        object_types=None,
        object_name=None,
        object_split=None,
        object_scale_range=None,
        object_size_limit=[0.05, 0.05],
        object_init_pose="planar", #{"planar", "fixed", "any"}
        **kwargs
    ):
        if initial_qpos is None:
            self.initial_qpos = np.array([0, 0.15, 0, -2.44, 0, 2.62, -7.84e-01])
        else:
            self.initial_qpos = initial_qpos

        # settings for table top
        self.table_full_size = np.array(table_full_size)
        self.table_size_rand = np.array(table_size_rand)
        self.table_size_curr = np.array(table_full_size) # table size after randomization
        self.table_friction = np.array(table_friction)
        self.table_solref = np.array(table_solref)
        self.table_solimp = np.array(table_solimp)
        self.table_offset = np.array(table_offset)
        self.table_xml = table_xml

        # reward configuration
        self.reward_scale = reward_scale
        self.reward_shaping = reward_shaping

        # whether to use ground-truth object states
        self.use_object_obs = use_object_obs

        # object placement initializer 
        # placement_initializer will be overwritten in _load_model
        self.placement_initializer = placement_initializer
        self.object_init_pose = object_init_pose
        self.object_scale = None
        
        # specify object
        self.object_dataset = object_dataset
        # self.object_types = object_types
        # self.object_name = object_name
        # self.object_split = object_split
        # self.object_eval_set = object_eval_set
        if object_dataset == 'housekeep' or object_dataset == 'housekeep_all':
            self.object_sampler = HousekeepSampler(
                object_types=object_types, object_name=object_name, 
                object_scale_range=object_scale_range,
                object_size_limit=object_size_limit,
                object_split=object_split, use_full_set=(object_dataset=='housekeep_all'))

        if not (kwargs is None or len(kwargs.keys()) == 0):
            print('Unused kwargs in DexEnv:', kwargs)

        # Additional Parameters for ADR
        self.object_to_wall_dist_min = 0
        self.object_to_wall_dist_max = 0
        self.object_pos_noise = 0
        self.object_ori_noise = 0
        self.table_offset_x_min = self.table_offset[0]
        self.table_offset_x_max = self.table_offset[0]
        self.table_offset_z_min = self.table_offset[2]
        self.table_offset_z_max = self.table_offset[2]

        self.object_size_x_min = 0.087
        self.object_size_x_max = 0.087
        self.object_size_x_val = 0.087
        self.object_size_y_min = 0.18
        self.object_size_y_max = 0.18
        self.object_size_y_val = 0.18
        self.object_size_z_min = 0.043
        self.object_size_z_max = 0.043
        self.object_size_z_val = 0.043
        self.object_density_min = 86.
        self.object_density_max = 86.
        self.object_density_val = 86.

        self.table_friction_min = self.table_friction[0]
        self.table_friction_max = self.table_friction[0]
        self.table_friction_val = self.table_friction[0]
        self.gripper_friction_min = 3
        self.gripper_friction_max = 3
        self.gripper_friction_val = 3
        self.close_gripper = close_gripper

        self.controller_max_translation_max = 0.03
        self.controller_max_translation_min = 0.03
        self.controller_max_rotation_max = 0.2
        self.controller_max_rotation_min = 0.2

        self.additional_obs_keys = additional_obs_keys if adaptive and additional_obs_keys else None
        self.remove_walls = remove_walls
        super().__init__(
            robots=robots,
            env_configuration=env_configuration,
            controller_configs=controller_configs,
            mount_types="default",
            gripper_types=gripper_types,
            initialization_noise=initialization_noise,
            use_camera_obs=use_camera_obs,
            has_renderer=has_renderer,
            has_offscreen_renderer=has_offscreen_renderer,
            render_camera=render_camera,
            render_collision_mesh=render_collision_mesh,
            render_visual_mesh=render_visual_mesh,
            render_gpu_device_id=render_gpu_device_id,
            control_freq=control_freq,
            horizon=horizon,
            ignore_done=ignore_done,
            hard_reset=hard_reset,
            camera_names=camera_names,
            camera_heights=camera_heights,
            camera_widths=camera_widths,
            camera_depths=camera_depths,
            camera_segmentations=camera_segmentations,
        )

    def _load_model(self):
        """
        Loads an xml model, puts it in self.model
        """
        super()._load_model()

        if 'max_translation' in self.robots[0].controller_config:
            self.robots[0].controller_config['max_translation'] = np.random.uniform(self.controller_max_translation_min,
                                                                                    self.controller_max_translation_max)
        if 'max_rotation' in self.robots[0].controller_config:
            self.robots[0].controller_config['max_rotation'] = np.random.uniform(self.controller_max_rotation_min,
                                                                                 self.controller_max_rotation_max)

        # load model for table top workspace
        self.table_friction_val = np.random.uniform(self.table_friction_min, self.table_friction_max)
        self.table_friction[0] = self.table_friction_val
        self.table_offset[0] = np.random.uniform(low=self.table_offset_x_min,
                                                 high=self.table_offset_x_max)
        self.table_offset[2] = np.random.uniform(low=self.table_offset_z_min,
                                                 high=self.table_offset_z_max)
        if self.remove_walls:
            bin_full_size = np.array([0.8, 0.8, 0.107])
            mujoco_arena = BinArena(bin_pos=self.table_offset,
                                    bin_full_size=bin_full_size,
                                    bin_friction=self.table_friction,
                                    bin_solref=self.table_solref,
                                    bin_solimp=self.table_solimp,
                                    hidden_walls="FBLR",
                                    xml_filename=self.table_xml)
        else:
            self.table_size_curr = self.table_full_size + (np.random.rand(3)*2 - 1)*self.table_size_rand
            mujoco_arena = BinArena(bin_pos=self.table_offset,
                                    bin_full_size=self.table_size_curr,
                                    bin_friction=self.table_friction,
                                    bin_solref=self.table_solref,
                                    bin_solimp=self.table_solimp,
                                    hidden_walls="",
                                    xml_filename=self.table_xml)

        # initialize objects of interest
        tex_attrib = {
            "type": "cube",
        }
        mat_attrib = {
            "texrepeat": "1 1",
            "specular": "0.4",
            "shininess": "0.1",
        }
        redwood = CustomMaterial(
            texture="WoodRed",
            tex_name="redwood",
            mat_name="redwood_mat",
            tex_attrib=tex_attrib,
            mat_attrib=mat_attrib,
        )

        self.object_size_x_val = np.random.uniform(self.object_size_x_min, self.object_size_x_max)
        self.object_size_y_val = np.random.uniform(self.object_size_y_min, self.object_size_y_max)
        self.object_size_z_val = np.random.uniform(self.object_size_z_min, self.object_size_z_max)
        self.object_density_val = np.random.randint(self.object_density_min, self.object_density_max+1)
        
        if self.object_dataset == "cube":
            self.cube = BoxObject(
                name="cube",
                size_min=[self.object_size_x_val/2, self.object_size_y_val/2, self.object_size_z_val/2],
                size_max=[self.object_size_x_val/2, self.object_size_y_val/2, self.object_size_z_val/2],
                rgba=[1, 0, 0, 1],
                material=redwood,
                density=self.object_density_val,
            )
            
            self.cube_target = BoxObject(
                name="cube_target",
                size_min=[self.object_size_x_val/2, self.object_size_y_val/2, self.object_size_z_val/2],
                size_max=[self.object_size_x_val/2, self.object_size_y_val/2, self.object_size_z_val/2],
                rgba=[0, 1, 0, 0.1],
                joints=None,
                obj_type="visual",
            )
        
        elif self.object_dataset == "housekeep" or self.object_dataset == "housekeep_all":
            obj, target_obj = self.object_sampler.generate_housekeep_object()
            self.cube = obj
            self.cube_target = target_obj
        
        else:
            raise ValueError(f"Unknown object type: {self.object_dataset}")

        # Create placement initializer
        if self.object_init_pose == "planar":
            max_dimension = max(self.object_size_x_val, self.object_size_y_val, self.object_size_z_val)
            self.placement_initializer = UniformRandomSampler(
                name="ObjectSampler",
                mujoco_objects=self.cube,
                x_range=[-self.table_size_curr[0]/2+max_dimension/2, self.table_size_curr[0]/2-max_dimension/2],
                y_range=[-self.table_size_curr[1]/2+max_dimension/2, self.table_size_curr[1]/2-max_dimension/2],
                rotation=None,
                rotation_axis='z',
                ensure_object_boundary_in_range=False,  # Problematic
                ensure_valid_placement=False,
                reference_pos=self.table_offset,
                z_offset=0.05,
            )
        elif self.object_init_pose == "fixed":
            self.placement_initializer = UniformRandomSampler(
                name="ObjectSampler",
                mujoco_objects=self.cube,
                x_range=[0, 0],
                y_range=[0, 0],
                rotation=0,
                rotation_axis='z',
                ensure_object_boundary_in_range=False,  # Problematic
                ensure_valid_placement=False,
                reference_pos=self.table_offset,
                z_offset=0.05,
            )
        elif self.object_init_pose == "any-ori":
            max_dimension = max(self.object_size_x_val, self.object_size_y_val, self.object_size_z_val)
            self.placement_initializer = CustomizedSampler(
                name="ObjectSampler",
                mujoco_objects=self.cube,
                x_range=[0, 0],
                y_range=[0, 0],
                rotation=None,
                rotation_axis='all',
                ensure_object_boundary_in_range=False,  # Problematic
                ensure_valid_placement=False,
                reference_pos=self.table_offset,
                z_offset=0.05,
            )
        elif self.object_init_pose == "any":
            max_dimension = max(self.object_size_x_val, self.object_size_y_val, self.object_size_z_val)
            self.placement_initializer = CustomizedSampler(
                name="ObjectSampler",
                mujoco_objects=self.cube,
                x_range=[-self.table_size_curr[0]/2+max_dimension/2, self.table_size_curr[0]/2-max_dimension/2],
                y_range=[-self.table_size_curr[1]/2+max_dimension/2, self.table_size_curr[1]/2-max_dimension/2],
                rotation=None,
                rotation_axis='all',
                ensure_object_boundary_in_range=False,  # Problematic
                ensure_valid_placement=False,
                reference_pos=self.table_offset,
                z_offset=0.05,
            )
        else:
            raise NotImplementedError

        # task includes arena, robot, and objects of interest
        self.model = ManipulationTask(
            mujoco_arena=mujoco_arena,
            mujoco_robots=[robot.robot_model for robot in self.robots], 
            mujoco_objects=[self.cube, self.cube_target],
        )

    def _setup_references(self):
        """
        Sets up references to important components. A reference is typically an
        index or a list of indices that point to the corresponding elements
        in a flatten array, which is how MuJoCo stores physical simulation data.
        """
        super()._setup_references()

        # Additional object references from this env
        self.cube_body_id = self.sim.model.body_name2id(self.cube.root_body)

    def _setup_observables(self):
        """
        Sets up observables to be used for this environment. Creates object-based observables if enabled

        Returns:
            OrderedDict: Dictionary mapping observable names to its corresponding Observable object
        """
        observables = super()._setup_observables()

        # low-level object information
        if self.use_object_obs:
            # Get robot prefix and define observables modality
            pf = self.robots[0].robot_model.naming_prefix
            modality = "object"

            # cube-related observables
            @sensor(modality=modality)
            def cube_pos(obs_cache):
                return np.array(self.sim.data.body_xpos[self.cube_body_id])

            @sensor(modality=modality)
            def cube_quat(obs_cache):
                return np.array(self.sim.data.body_xquat[self.cube_body_id])
            
            @sensor(modality=modality)
            def cube_size(obs_cache):
                return np.array([self.object_size_x_val, self.object_size_y_val, self.object_size_z_val])

            sensors = [cube_pos, cube_quat, cube_size]
            names = [s.__name__ for s in sensors]

            # Create observables
            for name, s in zip(names, sensors):
                observables[name] = Observable(
                    name=name,
                    sensor=s,
                    sampling_rate=self.control_freq,
                )

        return observables
        
    def _reset_internal(self):
        """
        Resets simulation internal configurations.
        """
        super()._reset_internal()

        # Reset all object positions using initializer sampler if we're not directly loading from an xml
        if not self.deterministic_reset:

            # Sample from the placement initializer for all objects
            object_placements = self.placement_initializer.sample()

            # Loop through all objects and reset their positions
            for obj_pos, obj_quat, obj in object_placements.values():
                self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)]))

    def _load_robots(self):
        gripper_friction = np.array([3, 0.015, 0.0003])
        self.gripper_friction_val = np.random.uniform(self.gripper_friction_min, self.gripper_friction_max)
        gripper_friction[0] = self.gripper_friction_val
        if self.robot_names[0] == "Virtual":
            self.robot_configs[0].update({'mount_type':None})
        self.robots[0] = SingleArm(robot_type=self.robot_names[0], idn=0, control_gripper=False,
                                   initial_qpos=self.initial_qpos, gripper_friction=gripper_friction,
                                   **self.robot_configs[0], close_gripper=self.close_gripper)
        self.robots[0].load_model()


class CustomizedSampler(UniformRandomSampler):
    def _sample_quat(self):
        if self.rotation_axis == 'all':
            return transform_utils.random_quat()
        else:
            super()._sample_quat()
bin_env/base_env.py
HACMan-2023/HACMan
import numpy as np
from robosuite.models.arenas import Arena
from robosuite.utils.mjcf_utils import xml_path_completion
from robosuite.utils.mjcf_utils import array_to_string
import os


class BinArena(Arena):
    """
    Workspace that contains one bin. The table underneath the bin is just for visualization purpose.

    Args:
        bin_pos (3-tuple): (x,y,z) position to place bin (center of bottom plate, upper surface)
        bin_full_size (3-tuple): (L,W,H) full dimensions of the bin (inner dimension)
        bin_thickness (float): thickness of the walls
        bin_friction (3-tuple): (sliding, torsional, rolling) friction parameters of the table
        bin_height (float): height of the bin relative to the ground
        hidden_walls (str): hide some of the walls if neccessary. i.e.: "FL" will hide front and left walls.
    """

    def __init__(
            self,
            xml_filename="bin_arena.xml",
            bin_pos=(0, 0, 0),
            bin_full_size=(0.45, 0.54, 0.107),
            bin_friction=(0.3, 0.005, 0.0001),
            bin_solref=(0.02, 1.),
            bin_solimp=(0.9, 0.95, 0.001), 
            bin_thickness=0.01,
            bin_height=0.8,
            hidden_walls="",
    ):
        xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets/arenas', xml_filename)
        super().__init__(xml_path)

        self.bin_full_size = np.array(bin_full_size)
        self.bin_half_size = self.bin_full_size / 2
        self.bin_friction = bin_friction
        self.bin_solref = bin_solref
        self.bin_solimp = bin_solimp
        self.bin_thickness = bin_thickness
        self.hidden_walls = str.upper(hidden_walls)

        self.bin_body = self.worldbody.find("./body[@name='bin']")
        self.bin_height = bin_height
        self.bin_visuals = {}
        self.bin_collisions = {}
        self.bin_planes = {}
        for name in ['bottom', 'left', 'right', 'front', 'back']:
            self.bin_visuals[name] = self.bin_body.find("./geom[@name='bin_" + name + "_visual']")
            self.bin_collisions[name] = self.bin_body.find("./geom[@name='bin_" + name + "_collision']")
            self.bin_planes[name] = self.bin_body.find("./geom[@name='bin_" + name + "_plane']")

        self.configure_location()

        self.set_origin((bin_pos[0], bin_pos[1], bin_pos[2] - bin_height))
        self.table_offset = np.array(bin_pos)  # For compatibility

    def configure_location(self):
        """Configures correct locations for this arena"""
        self.floor.set("pos", array_to_string(self.bottom_pos))

        # Set bin locations
        self.bin_body.set("pos", array_to_string(np.array([0, 0, self.bin_height])))
        # Bottom
        size = np.array([self.bin_half_size[0], self.bin_half_size[1], self.bin_thickness])
        size += np.array([self.bin_thickness*2, self.bin_thickness*2, 0])  # To cover the edges
        pos = np.array([0, 0, -self.bin_thickness])
        self.bin_visuals["bottom"].set("size", array_to_string(size))
        self.bin_visuals["bottom"].set("pos", array_to_string(pos))
        self.bin_collisions["bottom"].set("size", array_to_string(size))
        self.bin_collisions["bottom"].set("pos", array_to_string(pos))
        if self.bin_planes["bottom"] is not None:
            self.bin_planes["bottom"].set("pos", array_to_string(np.array([0, 0, 0])))
            self.bin_collisions["bottom"].set("pos", array_to_string(np.array([0, 0, -1]))) # Hide

        # Left
        size = np.array([self.bin_half_size[0], self.bin_thickness, self.bin_half_size[2]])
        size += np.array([self.bin_thickness*2, 0, 0])  # To cover the edges
        pos = np.array([0, -self.bin_half_size[1] - self.bin_thickness, self.bin_half_size[2]])
        if "L" in self.hidden_walls: pos = np.array([0, 0, -1])
        self.bin_visuals["left"].set("size", array_to_string(size))
        self.bin_visuals["left"].set("pos", array_to_string(pos))
        self.bin_collisions["left"].set("size", array_to_string(size))
        self.bin_collisions["left"].set("pos", array_to_string(pos))
        if self.bin_planes["left"] is not None:
            self.bin_planes["left"].set("pos", array_to_string(np.array([0, -self.bin_half_size[1], self.bin_half_size[2]])))
            self.bin_collisions["left"].set("pos", array_to_string(np.array([0, 0, -1]))) # Hide
                    
        # Right
        size = np.array([self.bin_half_size[0], self.bin_thickness, self.bin_half_size[2]])
        size += np.array([self.bin_thickness*2, 0, 0])  # To cover the edges
        pos = np.array([0, self.bin_half_size[1] + self.bin_thickness, self.bin_half_size[2]])
        if "R" in self.hidden_walls: pos = np.array([0, 0, -1])
        self.bin_visuals["right"].set("size", array_to_string(size))
        self.bin_visuals["right"].set("pos", array_to_string(pos))
        self.bin_collisions["right"].set("size", array_to_string(size))
        self.bin_collisions["right"].set("pos", array_to_string(pos))
        if self.bin_planes["right"] is not None:
            self.bin_planes["right"].set("pos", array_to_string(np.array([0, self.bin_half_size[1], self.bin_half_size[2]])))
            self.bin_collisions["right"].set("pos", array_to_string(np.array([0, 0, -1]))) # Hide
                    
        # Front
        size = np.array([self.bin_thickness, self.bin_half_size[1], self.bin_half_size[2]])
        size += np.array([0, self.bin_thickness*2, 0])  # To cover the edges
        pos = np.array([self.bin_half_size[0] + self.bin_thickness, 0, self.bin_half_size[2]])
        if "F" in self.hidden_walls: pos = np.array([0, 0, -1])
        self.bin_visuals["front"].set("size", array_to_string(size))
        self.bin_visuals["front"].set("pos", array_to_string(pos))
        self.bin_collisions["front"].set("size", array_to_string(size))
        self.bin_collisions["front"].set("pos", array_to_string(pos))
        if self.bin_planes["front"] is not None:
            self.bin_planes["front"].set("pos", array_to_string(np.array([self.bin_half_size[0], 0, self.bin_half_size[2]])))
            self.bin_collisions["front"].set("pos", array_to_string(np.array([0, 0, -1]))) # Hide
                    
        # Back
        size = np.array([self.bin_thickness, self.bin_half_size[1], self.bin_half_size[2]])
        size += np.array([0, self.bin_thickness*2, 0])  # To cover the edges
        pos = np.array([-self.bin_half_size[0] - self.bin_thickness, 0, self.bin_half_size[2]])
        if "B" in self.hidden_walls: pos = np.array([0, 0, -1])
        self.bin_visuals["back"].set("size", array_to_string(size))
        self.bin_visuals["back"].set("pos", array_to_string(pos))
        self.bin_collisions["back"].set("size", array_to_string(size))
        self.bin_collisions["back"].set("pos", array_to_string(pos))
        if self.bin_planes["back"] is not None:
            self.bin_planes["back"].set("pos", array_to_string(np.array([-self.bin_half_size[0], 0, self.bin_half_size[2]])))
            self.bin_collisions["back"].set("pos", array_to_string(np.array([0, 0, -1]))) # Hide
            
        # Set friction
        for key, value in self.bin_collisions.items():
            value.set("friction", array_to_string(self.bin_friction))
            value.set("solref", array_to_string(self.bin_solref))
            value.set("solimp", array_to_string(self.bin_solimp))
            value.set("priority", str(1))
            value.set("condim", str(4))
            
        # Set friction
        for key, value in self.bin_planes.items():
            # Note that planes are all infinite.
            if value is not None:
                value.set("friction", array_to_string(self.bin_friction))
                value.set("solref", array_to_string(self.bin_solref))
                value.set("solimp", array_to_string(self.bin_solimp))
                value.set("priority", str(1))
                value.set("condim", str(4))
                value.set("rgba", array_to_string(np.array([0, 0, 0, 0]))) # hide

        # Set table locations
        table_full_size = np.array([self.bin_full_size[0]+0.2, self.bin_full_size[1]+0.2, 0.05])
        table_offset = self.bin_height
        table_offset -= self.bin_thickness*2
        table_half_size = table_full_size / 2
        center_pos = self.bottom_pos + np.array([0, 0, -table_half_size[2] + table_offset])

        table_body = self.worldbody.find("./body[@name='table']")
        table_collision = table_body.find("./geom[@name='table_collision']")
        table_visual = table_body.find("./geom[@name='table_visual']")
        table_top = table_body.find("./site[@name='table_top']")

        table_legs_visual = [
            table_body.find("./geom[@name='table_leg1_visual']"),
            table_body.find("./geom[@name='table_leg2_visual']"),
            table_body.find("./geom[@name='table_leg3_visual']"),
            table_body.find("./geom[@name='table_leg4_visual']"),
        ]

        table_body.set("pos", array_to_string(center_pos))
        table_collision.set("size", array_to_string(table_half_size))
        table_visual.set("size", array_to_string(table_half_size))
        table_top.set("pos", array_to_string(np.array([0, 0, table_half_size[2]])))

        # Set leg locations
        delta_x = [0.1, -0.1, -0.1, 0.1]
        delta_y = [0.1, 0.1, -0.1, -0.1]
        for leg, dx, dy in zip(table_legs_visual, delta_x, delta_y):
            # If x-length of table is less than a certain length, place leg in the middle between ends
            # Otherwise we place it near the edge
            x = 0
            if table_half_size[0] > abs(dx * 2.0):
                x += np.sign(dx) * table_half_size[0] - dx
            # Repeat the same process for y
            y = 0
            if table_half_size[1] > abs(dy * 2.0):
                y += np.sign(dy) * table_half_size[1] - dy
            # Get z value
            z = (table_offset - table_half_size[2]) / 2.0
            # Set leg position
            leg.set("pos", array_to_string([x, y, -z]))
            # Set leg size
            leg.set("size", array_to_string([0.025, z]))

bin_env/bin_arena.py
HACMan-2023/HACMan
import os
from robosuite.models.grippers.gripper_model import GripperModel


class CustomizedRobotiq85Gripper(GripperModel):
    """
    6-DoF Robotiq gripper with festo fingertip.

    Args:
        idn (int or str): Number or some other unique identification string for this gripper instance
    """
    def __init__(self, idn=0):
        xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets/grippers', 'robotiq_gripper_85.xml')
        super().__init__(xml_path, idn=idn)

    def format_action(self, action):
        return action

    @property
    def _important_geoms(self):
        return {
            "left_finger": [
                "left_outer_finger_collision",
                "left_inner_finger_collision",
                "left_fingertip_collision",
                "left_fingerpad_collision"
            ],
            "right_finger": [
                "right_outer_finger_collision",
                "right_inner_finger_collision",
                "right_fingertip_collision",
                "right_fingerpad_collision"
            ],
            "left_fingerpad": [
                "left_fingerpad_collision"
            ],
            "right_fingerpad": [
                "right_fingerpad_collision"
            ],
        }
    
    @property
    def dof(self):
        return 0
bin_env/gripper.py
HACMan-2023/HACMan
from robosuite.controllers.base_controller import Controller
from robosuite.utils.control_utils import *
import robosuite.utils.transform_utils as T
import numpy as np
from bin_env.util import get_local_pose

# Set VERBOSE to True to debug
np.set_printoptions(precision=3)
VERBOSE = False

# Supported impedance modes
IMPEDANCE_MODES = {"fixed", "variable", "variable_kp"}


def angle_diff(vec1, vec2, degree=True):
    vec1 = vec1 / np.linalg.norm(vec1)
    vec2 = vec2 / np.linalg.norm(vec2)
    prod = np.dot(vec1, vec2)
    prod = np.clip(prod, -1, 1)
    angle = np.arccos(prod)
    if degree:
        angle = angle / np.pi * 180
    return angle


class OSC(Controller):
    """
    Controller for controlling robot arm via operational space control. Allows position and / or orientation control
    of the robot's end effector. For detailed information as to the mathematical foundation for this controller, please
    reference http://khatib.stanford.edu/publications/pdfs/Khatib_1987_RA.pdf

    NOTE: Control input actions can either be taken to be relative to the current position / orientation of the
    end effector or absolute values. In either case, a given action to this controller is assumed to be of the form:
    (x, y, z, ax, ay, az) if controlling pos and ori or simply (x, y, z) if only controlling pos

    Args:
        sim (MjSim): Simulator instance this controller will pull robot state updates from

        eef_name (str): Name of controlled robot arm's end effector (from robot XML)

        joint_indexes (dict): Each key contains sim reference indexes to relevant robot joint information, namely:

            :`'joints'`: list of indexes to relevant robot joints
            :`'qpos'`: list of indexes to relevant robot joint positions
            :`'qvel'`: list of indexes to relevant robot joint velocities

        actuator_range (2-tuple of array of float): 2-Tuple (low, high) representing the robot joint actuator range

        input_max (float or Iterable of float): Maximum above which an inputted action will be clipped. Can be either be
            a scalar (same value for all action dimensions), or a list (specific values for each dimension). If the
            latter, dimension should be the same as the control dimension for this controller

        input_min (float or Iterable of float): Minimum below which an inputted action will be clipped. Can be either be
            a scalar (same value for all action dimensions), or a list (specific values for each dimension). If the
            latter, dimension should be the same as the control dimension for this controller

        output_max (float or Iterable of float): Maximum which defines upper end of scaling range when scaling an input
            action. Can be either be a scalar (same value for all action dimensions), or a list (specific values for
            each dimension). If the latter, dimension should be the same as the control dimension for this controller

        output_min (float or Iterable of float): Minimum which defines upper end of scaling range when scaling an input
            action. Can be either be a scalar (same value for all action dimensions), or a list (specific values for
            each dimension). If the latter, dimension should be the same as the control dimension for this controller

        kp (float or Iterable of float): positional gain for determining desired torques based upon the pos / ori error.
            Can be either be a scalar (same value for all action dims), or a list (specific values for each dim)

        damping_ratio (float or Iterable of float): used in conjunction with kp to determine the velocity gain for
            determining desired torques based upon the joint pos errors. Can be either be a scalar (same value for all
            action dims), or a list (specific values for each dim)

        impedance_mode (str): Impedance mode with which to run this controller. Options are {"fixed", "variable",
            "variable_kp"}. If "fixed", the controller will have fixed kp and damping_ratio values as specified by the
            @kp and @damping_ratio arguments. If "variable", both kp and damping_ratio will now be part of the
            controller action space, resulting in a total action space of (6 or 3) + 6 * 2. If "variable_kp", only kp
            will become variable, with damping_ratio fixed at 1 (critically damped). The resulting action space will
            then be (6 or 3) + 6.

        kp_limits (2-list of float or 2-list of Iterable of floats): Only applicable if @impedance_mode is set to either
            "variable" or "variable_kp". This sets the corresponding min / max ranges of the controller action space
            for the varying kp values. Can be either be a 2-list (same min / max for all kp action dims), or a 2-list
            of list (specific min / max for each kp dim)

        damping_ratio_limits (2-list of float or 2-list of Iterable of floats): Only applicable if @impedance_mode is
            set to "variable". This sets the corresponding min / max ranges of the controller action space for the
            varying damping_ratio values. Can be either be a 2-list (same min / max for all damping_ratio action dims),
            or a 2-list of list (specific min / max for each damping_ratio dim)

        policy_freq (int): Frequency at which actions from the robot policy are fed into this controller

        position_limits (2-list of float or 2-list of Iterable of floats): Limits (m) below and above which the
            magnitude of a calculated goal eef position will be clipped. Can be either be a 2-list (same min/max value
            for all cartesian dims), or a 2-list of list (specific min/max values for each dim)

        orientation_limits (2-list of float or 2-list of Iterable of floats): Limits (rad) below and above which the
            magnitude of a calculated goal eef orientation will be clipped. Can be either be a 2-list
            (same min/max value for all joint dims), or a 2-list of list (specific min/mx values for each dim)

        interpolator_pos (Interpolator): Interpolator object to be used for interpolating from the current position to
            the goal position during each timestep between inputted actions

        interpolator_ori (Interpolator): Interpolator object to be used for interpolating from the current orientation
            to the goal orientation during each timestep between inputted actions

        control_ori (bool): Whether inputted actions will control both pos and ori or exclusively pos

        uncouple_pos_ori (bool): Whether to decouple torques meant to control pos and torques meant to control ori

        **kwargs: Does nothing; placeholder to "sink" any additional arguments so that instantiating this controller
            via an argument dict that has additional extraneous arguments won't raise an error

    Raises:
        AssertionError: [Invalid impedance mode]
    """

    def __init__(self,
                 sim,
                 eef_name,
                 joint_indexes,
                 actuator_range,
                 input_max=1,
                 input_min=-1,
                 output_max=1,
                 output_min=-1,
                 max_translation=0.03,
                 max_rotation=0.1,
                 kp=150,
                 damping_ratio=1,
                 impedance_mode="fixed",
                 kp_limits=(0, 300),
                 damping_ratio_limits=(0, 100),
                 policy_freq=20,
                 position_limits=None,
                 orientation_limits=None,
                 interpolator_pos=None,
                 interpolator_ori=None,
                 control_ori=True,
                 control_delta=True,
                 uncouple_pos_ori=True,
                 use_lambda=False,
                 use_nullspace=False,
                 type=None,
                 control_axis=(1, 1, 1, 1, 1, 1),
                 **kwargs  # does nothing; used so no error raised when dict is passed with extra terms used previously
                 ):

        super().__init__(
            sim,
            eef_name,
            joint_indexes,
            actuator_range,
        )
        # Determine whether this is pos ori or just pos
        self.controller_mode = 'relative' if control_delta is True else 'accumulate'

        # Control dimension
        self.control_axis = np.array(control_axis, dtype=np.float32)  # Corresponds to translation in x,y,z and rotation along x,y,z

        # Backward compatibility
        if control_ori is False:
            self.control_axis[3:] = 0

        self.controller_type = type
        self.control_dim = int(np.sum(self.control_axis))

        # input and output max and min (allow for either explicit lists or single numbers)
        self.input_max = self.nums2array(input_max, self.control_dim)
        self.input_min = self.nums2array(input_min, self.control_dim)
        self.output_max = self.nums2array(output_max, self.control_dim)
        self.output_min = self.nums2array(output_min, self.control_dim)
        self.max_translation = max_translation
        self.max_rotation = max_rotation

        # kp kd
        self.kp = self.nums2array(kp, 6)
        self.kd = 2 * np.sqrt(self.kp) * damping_ratio

        # kp and kd limits
        self.kp_min = self.nums2array(kp_limits[0], 6)
        self.kp_max = self.nums2array(kp_limits[1], 6)
        self.damping_ratio_min = self.nums2array(damping_ratio_limits[0], 6)
        self.damping_ratio_max = self.nums2array(damping_ratio_limits[1], 6)

        # Verify the proposed impedance mode is supported
        assert impedance_mode in IMPEDANCE_MODES, "Error: Tried to instantiate OSC controller for unsupported " \
                                                  "impedance mode! Inputted impedance mode: {}, Supported modes: {}". \
            format(impedance_mode, IMPEDANCE_MODES)

        # Impedance mode
        self.impedance_mode = impedance_mode

        # Add to control dim based on impedance_mode
        if self.impedance_mode == "variable":
            self.control_dim += 12
        elif self.impedance_mode == "variable_kp":
            self.control_dim += 6

        # limits
        self.position_limits = np.array(position_limits) if position_limits is not None else position_limits
        self.orientation_limits = np.array(orientation_limits) if orientation_limits is not None else orientation_limits

        # control frequency
        self.control_freq = policy_freq

        # interpolator
        self.interpolator_pos = interpolator_pos
        self.interpolator_ori = interpolator_ori

        # whether or not pos and ori want to be uncoupled
        self.uncoupling = uncouple_pos_ori
        self.use_lambda = use_lambda
        self.use_nullspace = use_nullspace

        # initialize goals based on initial pos / ori
        self.initial_ee_ori_mat = np.round(self.initial_ee_ori_mat)
        self.goal_ori = np.array(self.initial_ee_ori_mat)
        self.goal_pos = np.array(self.initial_ee_pos)

        self.relative_ori = np.zeros(3)
        self.ori_ref = None

    def set_goal(self, action, set_pos=None, set_ori=None):
        """
        Sets goal based on input @action. If self.impedance_mode is not "fixed", then the input will be parsed into the
        delta values to update the goal position / pose and the kp and/or damping_ratio values to be immediately updated
        internally before executing the proceeding control loop.

        Note that @action expected to be in the following format, based on impedance mode!

            :Mode `'fixed'`: [joint pos command]
            :Mode `'variable'`: [damping_ratio values, kp values, joint pos command]
            :Mode `'variable_kp'`: [kp values, joint pos command]

        Args:
            action (Iterable): Desired relative joint position goal state
            set_pos (Iterable): If set, overrides @action and sets the desired absolute eef position goal state
            set_ori (Iterable): IF set, overrides @action and sets the desired absolute eef orientation goal state
        """
        # Update state
        self.update()

        # # TODO: parse action for variable impedance mode
        # # Parse action based on the impedance mode, and update kp / kd as necessary
        # if self.impedance_mode == "variable":
        #     damping_ratio, kp, delta = action[:6], action[6:12], action[12:]
        #     self.kp = np.clip(kp, self.kp_min, self.kp_max)
        #     self.kd = 2 * np.sqrt(self.kp) * np.clip(damping_ratio, self.damping_ratio_min, self.damping_ratio_max)
        # elif self.impedance_mode == "variable_kp":
        #     kp, delta = action[:6], action[6:]
        #     self.kp = np.clip(kp, self.kp_min, self.kp_max)
        #     self.kd = 2 * np.sqrt(self.kp)  # critically damped
        # else:   # This is case "fixed"
        #     delta = action

        # Align actions to corresponding the enabled axis.
        # For example, 3D action space in XZPlane will be assigned to x-translation, z-translation and y-rotation.
        delta = np.zeros_like(self.control_axis)
        dim_count = 0
        for i in range(len(self.control_axis)):
            if self.control_axis[i] == 1:
                delta[i] = action[dim_count]
                dim_count += 1
        assert dim_count == len(action)

        pos, ori = self._get_desired_pose(delta)

        if set_ori is not None:
            ori = set_ori
        if set_pos is not None:
            pos = set_pos
        # If the desired pose is reaching orientation limit or joint limit,
        # it will use the previous goal
        # if self.orientation_feasible(ori) and self.joint_limit_feasible(pos, ori):
        self.goal_ori = ori
        self.goal_pos = pos

        if self.interpolator_pos is not None:
            self.interpolator_pos.set_goal(self.goal_pos)

        if self.interpolator_ori is not None:
            self.ori_ref = np.array(self.ee_ori_mat)  # reference is the current orientation at start
            self.interpolator_ori.set_goal(orientation_error(self.goal_ori, self.ori_ref))  # goal is the total orientation error
            self.relative_ori = np.zeros(3)  # relative orientation always starts at 0

    def _get_desired_pose(self, action):
        # Get base pose
        if self.controller_mode == 'accumulate':
            base_pose = T.make_pose(self.goal_pos, self.goal_ori).copy()
        elif self.controller_mode == 'relative':
            base_pose = T.make_pose(self.ee_pos, self.ee_ori_mat).copy()
        else:
            assert False, "Error: unsupported controller mode: {}, Available options: accumulate, relative"\
                .format(self.controller_mode)
        pos = base_pose[:3, 3]
        orn = base_pose[:3, :3]

        # Apply translation
        pos += action[:3] * self.max_translation
        # Use initial ee_pos if the axis is not enabled
        pos = np.where(self.control_axis[:3] == 1, pos, self.initial_ee_pos)

        # # Limit desired pose within the bounding box
        if self.position_limits is not None:
            if np.any(pos > self.position_limits[1]) or np.any(pos < self.position_limits[0]):
                if VERBOSE:
                    print('Reaching translation limit.')
                pos = np.minimum(self.position_limits[1], pos)
                pos = np.maximum(self.position_limits[0], pos)

        if np.all(self.control_axis[3:] == np.array([0, 1, 0])):
            # Rotation along Y axis only
            z_axis = orn[:, 2]  # global frame
            initial_z_axis = self.initial_ee_ori_mat[:, 2]  # global frame
            # TODO: remove angle_diff
            rotation_angle = angle_diff(z_axis, initial_z_axis)/180*np.pi
            rotation_angle *= np.sign(np.cross(initial_z_axis, z_axis)[1])
            # rotation axis given by positive y-direction in the global frame
            rotation_angle += action[4] * self.max_rotation
            orn = (T.quat2mat(T.axisangle2quat(np.array([0, 1, 0])*rotation_angle))).dot(self.initial_ee_ori_mat)
        elif np.all(self.control_axis[3:] == np.array([1, 1, 1])):
            # divide by sqrt(3) because the norm of the max action is larger than dim=1.
            orn = (T.quat2mat(T.axisangle2quat(action[3:] * self.max_rotation/np.sqrt(3)))).dot(orn)
        elif np.all(self.control_axis[3:] == np.array([0, 0, 0])):
            orn = self.initial_ee_ori_mat
        else:
            raise NotImplementedError

        return pos, orn

    def run_controller(self):
        """
        Calculates the torques required to reach the desired setpoint.

        Executes Operational Space Control (OSC) -- either position only or position and orientation.

        A detailed overview of derivation of OSC equations can be seen at:
        http://khatib.stanford.edu/publications/pdfs/Khatib_1987_RA.pdf

        Returns:
             np.array: Command torques
        """
        # Update state
        self.update()

        desired_pos = None
        # Only linear interpolator is currently supported
        if self.interpolator_pos is not None:
            # Linear case
            if self.interpolator_pos.order == 1:
                desired_pos = self.interpolator_pos.get_interpolated_goal()
            else:
                # Nonlinear case not currently supported
                pass
        else:
            desired_pos = np.array(self.goal_pos)

        if self.interpolator_ori is not None:
            # relative orientation based on difference between current ori and ref
            self.relative_ori = orientation_error(self.ee_ori_mat, self.ori_ref)

            ori_error = self.interpolator_ori.get_interpolated_goal()
        else:
            desired_ori = np.array(self.goal_ori)
            ori_error = orientation_error(desired_ori, self.ee_ori_mat)

        # Compute desired force and torque based on errors
        position_error = desired_pos - self.ee_pos
        vel_pos_error = -self.ee_pos_vel

        # F_r = kp * pos_err + kd * vel_err
        desired_force = (np.multiply(np.array(position_error), np.array(self.kp[0:3]))
                         + np.multiply(vel_pos_error, self.kd[0:3]))

        vel_ori_error = -self.ee_ori_vel

        # Tau_r = kp * ori_err + kd * vel_err
        desired_torque = (np.multiply(np.array(ori_error), np.array(self.kp[3:6]))
                          + np.multiply(vel_ori_error, self.kd[3:6]))

        # Compute nullspace matrix (I - Jbar * J) and lambda matrices ((J * M^-1 * J^T)^-1)
        lambda_full, lambda_pos, lambda_ori, nullspace_matrix = opspace_matrices(self.mass_matrix,
                                                                                 self.J_full,
                                                                                 self.J_pos,
                                                                                 self.J_ori)

        # Decouples desired positional control from orientation control
        if self.uncoupling and self.use_lambda:
            decoupled_force = np.dot(lambda_pos, desired_force)
            decoupled_torque = np.dot(lambda_ori, desired_torque)
            wrench = np.concatenate([decoupled_force, decoupled_torque])
        elif self.use_lambda:
            desired_wrench = np.concatenate([desired_force, desired_torque])
            wrench = np.dot(lambda_full, desired_wrench)
        else:
            wrench = np.concatenate([desired_force, desired_torque])

        # Gamma (without null torques) = J^T * F + gravity compensations
        self.torques = np.dot(self.J_full.T, wrench) + self.torque_compensation

        # Calculate and add nullspace torques (nullspace_matrix^T * Gamma_null) to final torques
        # Note: Gamma_null = desired nullspace pose torques, assumed to be positional joint control relative
        #                     to the initial joint positions
        if self.use_nullspace:
            self.torques += nullspace_torques(self.mass_matrix, nullspace_matrix,
                                              self.initial_joint, self.joint_pos, self.joint_vel)

        # Always run superclass call for any cleanups at the end
        super().run_controller()

        return self.torques

    def update_initial_joints(self, initial_joints):
        # First, update from the superclass method
        super().update_initial_joints(initial_joints)

        # We also need to reset the goal in case the old goals were set to the initial confguration
        self.reset_goal()

    def reset_goal(self):
        """
        Resets the goal to the current state of the robot
        """
        self.goal_ori = np.array(self.ee_ori_mat)
        self.goal_pos = np.array(self.ee_pos)

        # Also reset interpolators if required

        if self.interpolator_pos is not None:
            self.interpolator_pos.set_goal(self.goal_pos)

        if self.interpolator_ori is not None:
            self.ori_ref = np.array(self.ee_ori_mat)  # reference is the current orientation at start
            self.interpolator_ori.set_goal(orientation_error(self.goal_ori, self.ori_ref))  # goal is the total orientation error
            self.relative_ori = np.zeros(3)  # relative orientation always starts at 0

    @property
    def control_limits(self):
        """
        Returns the limits over this controller's action space, overrides the superclass property
        Returns the following (generalized for both high and low limits), based on the impedance mode:

            :Mode `'fixed'`: [joint pos command]
            :Mode `'variable'`: [damping_ratio values, kp values, joint pos command]
            :Mode `'variable_kp'`: [kp values, joint pos command]

        Returns:
            2-tuple:

                - (np.array) minimum action values
                - (np.array) maximum action values
        """
        if self.impedance_mode == "variable":
            low = np.concatenate([self.damping_ratio_min, self.kp_min, self.input_min])
            high = np.concatenate([self.damping_ratio_max, self.kp_max, self.input_max])
        elif self.impedance_mode == "variable_kp":
            low = np.concatenate([self.kp_min, self.input_min])
            high = np.concatenate([self.kp_max, self.input_max])
        else:  # This is case "fixed"
            low, high = self.input_min, self.input_max
        return low, high

    @property
    def name(self):
        return self.controller_type

    def orientation_feasible(self, orn):
        quat = T.convert_quat(T.mat2quat(orn.dot(self.initial_ee_ori_mat)), to='wxyz')
        angle = np.arccos(min(abs(quat[0]), 1))/np.pi*180*2
        assert self.orientation_limits.shape == (), 'Only taking a scalar limit for now.'
        if angle > self.orientation_limits:
            if VERBOSE:
                print('Desired_pose reaching orientation limit')
            return False

        return True

    def joint_limit_feasible(self, desired_pos, desired_ori):
        # joint limit from frankapy
        JOINT_LIMITS_MIN = np.array([-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973]) + 0.1
        JOINT_LIMITS_MAX = np.array([2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973]) - 0.1

        delta_pos = desired_pos - self.ee_pos
        delta_ori = orientation_error(desired_ori, self.ee_ori_mat)

        J_inv = np.linalg.pinv(self.J_full)
        new_joint_pos = self.joint_pos + np.dot(J_inv, np.concatenate([delta_pos, delta_ori]))

        if np.any(new_joint_pos < JOINT_LIMITS_MIN) or np.any(new_joint_pos > JOINT_LIMITS_MAX):
            if VERBOSE:
                print('Desired_pose reaching joint limit:')
                print('min:\t', JOINT_LIMITS_MIN)
                print('new:\t', new_joint_pos)
                print('max:\t', JOINT_LIMITS_MAX)
            return False

        return True

bin_env/osc.py
HACMan-2023/HACMan
import os, sys
import pickle
import numpy as np
import robosuite
import copy
from bin_env.base_env import BaseEnv
import open3d as o3d
from bin_env.utils.point_cloud_utils import convert_depth, get_point_cloud, add_additive_noise_to_xyz, dropout_random_ellipses
from robosuite.utils.control_utils import orientation_error
from bin_env.utils.transformations import to_pose_mat, decompose_pose_mat
from robosuite.controllers import load_controller_config

VERBOSE = False
robosuite.utils.macros.SIMULATION_TIMESTEP = 0.002

class PokeEnv(BaseEnv):
    def __init__(
            self,
            # Specific to PokeEnv
            ik_precontact=True,
            transparent_gripper_base=True,
            planar_action=False,
            record_from_cam=None,
            resting_time=2, 
            goal_mode='fixed',
            controller="osc_no_ori.json",  # "OSC_POSE"
            ignore_position_limits=True,
            renderer='offscreen',
            # Will be passed to BaseEnv
            robots='Virtual',
            close_gripper=True,
            object_dataset='cube',
            object_name=None,
            control_freq=2,
            ignore_done=True,
            render_camera='agentview',
            location_noise=0.,
            pcd_noise=None,
            object_voxel_down_sample=0.01,
            friction_config='default',
            **kwargs
    ):

        """
        Setup parameters for PokeEnv
        """
        self.planar_action = planar_action
        self.ik_precontact = ik_precontact
        self.transparent_gripper_base = transparent_gripper_base
        self.record_from_cam = record_from_cam 
        self.cam_frames = []
        self.goal = None
        self.resting_time = resting_time
        self.location_noise = location_noise
        self.pcd_noise = pcd_noise # [2, 30, 2, 0.005]
        self.object_voxel_down_sample = object_voxel_down_sample
        
        current_folder = os.path.dirname(os.path.abspath(__file__))
        
        # Setup goal sampling
        self.goal_mode = goal_mode
        self.goal_list = None
        if object_dataset == 'cube':
            if goal_mode == 'any-old':
                pose_file = os.path.join(current_folder, 'data/Dataset0004/states.npy')
            elif goal_mode == 'any-old-eval':
                pose_file = os.path.join(current_folder, 'data/Dataset0004_eval/states.npy')
            elif goal_mode == 'oop-old':
                pose_file = os.path.join(current_folder, 'data/Dataset0004/states_out_of_plane.npy')
            elif goal_mode == 'oop-old-eval':
                pose_file = os.path.join(current_folder, 'data/Dataset0004_eval/states_out_of_plane.npy')
            elif goal_mode in ['fixed', 'translation', 'upright', 'fixed_off_center', 'off_center']:
                pose_file = None
            else:
                raise ValueError("Unknown goal_mode: {}".format(goal_mode))
            
            if pose_file is not None:
                self.goal_list = {'cube': {'poses':np.load(pose_file)}}
        
        elif object_dataset in {'housekeep', 'housekeep_all'}:
            if goal_mode == 'any':
                if 'object_scale_range' in kwargs.keys() and kwargs['object_scale_range'] is not None:
                    assert kwargs['object_scale_range'][0] == kwargs['object_scale_range'][1], "object_scale_range should be a single value!"
                pose_file = os.path.join(current_folder, f'data/{object_dataset}/poses_fixed_size.pk')

            elif goal_mode == 'any_var_size':
                pose_file = os.path.join(current_folder, f'data/{object_dataset}/poses_variable_size_v1.pk')
            
            elif goal_mode == 'any_var_size_with_wall':
                pose_file = os.path.join(current_folder, f'data/{object_dataset}/poses_variable_size_with_wall.pk')

            elif goal_mode in ['fixed', 'translation', 'upright', 'fixed_off_center', 'off_center']:
                pose_file = None

            else:
                raise ValueError("Unknown goal_mode: {}".format(goal_mode))
            
            if pose_file is not None:
                assert os.path.exists(pose_file), "Pose file does not exist: {}".format(pose_file)
                with open(pose_file, 'rb') as f:
                    self.goal_list = pickle.load(f)
        
        else:
            ValueError("Unknown object_dataset: {}".format('object_dataset'))

        """
        Setup configurations for BaseEnv
        """
        if robots == 'Virtual':
            kwargs['initial_qpos'] = np.array([0.45, 0., 0.4, 0., np.pi, np.pi/2])
        else:
            kwargs['initial_qpos'] = np.array([0, -0.2743, 0, -2.263, 0, 2.008, -0.7946])
        kwargs['initialization_noise'] = None
        
        # if 'object_dataset' in kwargs and 'housekeep' in kwargs['object_dataset']:
        # kwargs['table_friction'] = (0., 0.3, 0.1)
        if friction_config == 'default':
            kwargs['table_friction']=(0.5, 0.005, 0.0001)
        elif friction_config == 'low':
            kwargs['table_friction']=(0.3, 0.005, 0.0001)
        elif friction_config == 'high':
            kwargs['table_friction']=(0.95, 0.05, 0.001)
        elif friction_config == 'robosuite':
            kwargs['table_friction']=(0.95, 0.3, 0.1)
        else:
            raise NotImplementedError    
        
        kwargs['table_solref'] = (0.01, 1)
        kwargs['table_xml'] = 'bin_arena_all_side_planes.xml'
        kwargs['hard_reset'] = object_name is None and object_dataset != 'cube'

        # Setup Controller
        if '.json' in controller:
            controller_configs = load_controller_config(custom_fpath=os.path.join(current_folder, 'controller_config', controller))
        else:
            controller_configs = load_controller_config(default_controller='OSC_POSE')
            
        if ignore_position_limits:
            controller_configs['position_limits'] = None
        kwargs.update(controller_configs=controller_configs)
        
        # Setup Renderer
        if renderer == 'onscreen':
            # Use onscreen renderer
            render_config = dict(has_renderer=True,
                                has_offscreen_renderer=False,
                                use_camera_obs=False)
            kwargs.update(render_config)
        elif renderer == 'offscreen':
            # Use offscreen renderer
            render_config = dict(has_renderer=False,
                                has_offscreen_renderer=True,
                                use_camera_obs=True,
                                camera_names=['leftview', 'rightview', 'agentview'],
                                camera_depths=True,
                                camera_segmentations='instance'
                                )
            kwargs.update(render_config)

        super().__init__(robots=robots,
            close_gripper=close_gripper,
            object_dataset=object_dataset,
            object_name=object_name,
            control_freq=control_freq,
            ignore_done=ignore_done,
            render_camera=render_camera,
            **kwargs)
                
        # Action location range (after BaseEnv init)
        self.location_scale = None
        self.location_center = None
        return
    
    def _reset_internal(self):
        super()._reset_internal()
        
        # Update location scale
        self.location_scale = self.table_size_curr/2 + np.array([-0.03, -0.03, 0])
        # make the range smaller for x and y, but not for z
        self.location_center = self.table_offset + np.array([0, 0, self.table_size_curr[2]/2])
        # self.table_offset is at the center of the bottom of the bin.
        # self.location_center is at the center of the entire free space of the bin.
        
        if self.transparent_gripper_base:
            gid = self.sim.model.geom_name2id('gripper0_hand_collision')
            self.sim.model.geom_conaffinity[gid] = 0
            self.sim.model.geom_contype[gid] = 0    
            gid = self.sim.model.geom_name2id('gripper0_hand_visual')
            self.sim.model.geom_rgba[gid] = np.array([1, 1, 1, 0.3])
        
        if self.robots[0].gripper_type == 'default':
            eef_body_id = self.sim.model.body_name2id('gripper0_eef')
            if self.close_gripper:
                self.sim.model.body_pos[eef_body_id] = np.array([0, 0, 0.105])
            else:
                self.sim.model.body_pos[eef_body_id] = np.array([-0.053, 0, 0.105])
        
        # Hide global coordinate markers
        for site_name in ['xaxis', 'yaxis', 'zaxis']:
            sid = self.sim.model.site_name2id(site_name)
            self.sim.model.site_rgba[sid] = np.array([1, 1, 1, 0])

        return
    
    def set_goal(self, goal):
        # Convert matrix into vector
        self.goal = decompose_pose_mat(goal, cat=True, wxyz=True)
        self.maybe_render()
        return
    
    def show_goal(self):
        goal = self.goal
        if goal is not None:
            body_id = self.sim.model.body_name2id(self.cube_target.root_body)
            self.sim.model.body_pos[body_id] = goal[:3]
            self.sim.model.body_quat[body_id] = goal[3:]
            self.sim.forward()
        return
    
    def hide_goal(self):
        body_id = self.sim.model.body_name2id(self.cube_target.root_body)
        self.sim.model.body_pos[body_id][2] = -0.2
        self.sim.forward()
        return

    def reset(self, object_pose=None, object_size=None, goal=None, hard_reset=None, object_name=None, start_gripper_above_obj=False,
              attempt=0):
        if object_size is not None:
            self.object_size_x_min = object_size[0]
            self.object_size_x_max = object_size[0]
            self.object_size_y_min = object_size[1]
            self.object_size_y_max = object_size[1]
            self.object_size_z_min = object_size[2]
            self.object_size_z_max = object_size[2]
        
        if object_name is not None:
            self.object_sampler.set_object(object_name)

        previous_hard_reset = self.hard_reset
        # Override hard_reset if specified
        self.hard_reset = hard_reset if hard_reset is not None else self.hard_reset
        super().reset()
        self.hard_reset = previous_hard_reset
        
        self.goal = goal
        # Update mujoco visual of the goal
        if self.goal is not None:
            body_id = self.sim.model.body_name2id(self.cube_target.root_body)
            self.sim.model.body_pos[body_id] = self.goal[:3]
            self.sim.model.body_quat[body_id] = self.goal[3:]
            self.sim.forward()
        else:
            body_id = self.sim.model.body_name2id(self.cube_target.root_body)
            self.sim.model.body_pos[body_id] = np.array([0,0,-1]) # hide it
            self.sim.forward()
                
        if object_pose is not None:
            self.sim.data.set_joint_qpos(self.cube.joints[0], object_pose)
            self.sim.forward()
            self.robots[0].controller.run_controller()
        
        if start_gripper_above_obj:
            obj_size = np.max(self.get_cube_size())
            
            gripper_pos = self.sim.data.get_joint_qpos(self.cube.joints[0])[:3]
            gripper_pos[2] = obj_size + self.table_offset[2] + 0.05
            self.move_to(gripper_pos)

        # Run simulation to rest the object.
        self.run_simulation(total_execution_time=self.resting_time)

        self.visualize(vis_settings=dict(env=False, grippers=True, robots=False))
        obs = self._get_observations(force_update=True)
        
        if 'object_pcd_points' in obs.keys() and obs['object_pcd_points'] is None:
            assert attempt < 5, "Failed to get object point cloud after 5 attempts."
            print(f'no object point cloud during reset. try to reset again. attempt {attempt}')
            obs = self.reset(hard_reset=False, attempt=attempt+1)

        return obs

    def get_cube_size(self):
        if hasattr(self.cube, 'scaled_size'):
            cube_size = self.cube.scaled_size
        elif hasattr(self.cube, 'size'):
            cube_size = np.array(self.cube.size) * 2
        else:
            raise ValueError("Unknown cube size.")
        
        return cube_size

    """
    Movement related
    """
    def sample_random_poke(self, obs=None, points=None, normals=None):
        # Take a dictionary of obs as input, or directly take points and normals as input
        if obs is not None:
            assert points is None and normals is None
            points = obs['object_pcd_points']
            normals = obs['object_pcd_normals']
        idx = np.random.randint(len(points))
        location = points[idx]
        normal = normals[idx]
        action = np.random.rand(self.action_dim)*2-1
        if self.planar_action:
            action[-1] = 0
        if action.dot(normal) > 0:
            action *= -1
        return location, normal, action, idx

    def poke(self, location, normal, action, action_repeat=10, start_location=None):
        '''
        Set normal to None to directly move to locations.
        '''
        self.cam_frames = []
        box_in_the_view = True
        # Save original object pose
        obj_pose = self.sim.data.get_joint_qpos(self.cube.joints[0]).copy()

        if location is not None and normal is None:
            assert start_location is not None, "Need to provide start location in regress_location_and_action mode."
        
        try:
            # Contiguous movements
            if location is None:
                pass    # No additional movement for contiguous movements

            # Regress locations and actions
            elif normal is None:
                # assert self.collision_check(start_location)
                self.robots[0].reset(deterministic=True)
                self.sim.forward()

                # Approach from start location (should be top)
                assert self.move_to_from_top(location=location, start_location=start_location)

            # Per-point actions
            else:
                # assert self.collision_check(location)
                self.robots[0].reset(deterministic=True)
                self.sim.forward()

                noise = (np.random.rand(3)*2-1)*self.location_noise
                location += noise
                assert self.move_to_contact(location=location, normal=normal)

        except:
            success = False

        else:
            # Execute action parameters
            success = True
            for _ in range(action_repeat):
                self.run_simulation(action)
        
        # Reset robot if using location
        if location is not None:
            self.robots[0].reset(deterministic=True)
            self.sim.forward()
            self.run_simulation(total_execution_time=self.resting_time)
        
        obs = self._get_observations(force_update=True)        
        if 'object_pcd_points' in obs.keys():
            if obs['object_pcd_points'] is None:
                box_in_the_view = False
                print(f'no object point cloud. try to reset.')
                obs = self.reset(hard_reset=False)
        else:
            obs['object_pcd_points'] = None

        # Add object name to the info
        if hasattr(self.cube, 'mesh_name'):
            obj_name = self.cube.mesh_name
        else:
            obj_name = "cube"

        info = {
            'poke_success': success and box_in_the_view,
            'box_in_the_view':box_in_the_view,
            'cam_frames': copy.copy(self.cam_frames),
            'object_name': obj_name}
        self.cam_frames.clear()
        # Following env.step() output format
        return obs, 0, False, info
    
    def collision_check(self, location):
        scaled_location = (location - self.location_center)/self.location_scale
        return np.all(np.abs(scaled_location) <= 1.01)

    def move_to_contact(self, location, normal):
        ee_pos = self.robots[0].controller.ee_pos
        precontact = location + normal * 0.02

        if self.robots[0].name == 'Virtual' and self.ik_precontact:
            # Set the floating gripper to the location directly (ik)
            preprecontact = location + normal * 0.04
            s1, s2 = self.move_to(pos=preprecontact, set_qpos=True)
            if not (s1 and s2):
                return False
            self.move_to(pos=precontact)
        else:
            # Gradually move to the contact point using the low-level controller
            midpoint = (precontact - ee_pos)*2/3 + ee_pos
            
            waypoints = 5
            for i in range(waypoints):
                new_pos = (midpoint - ee_pos)/waypoints*(i+1) + ee_pos
                s1, s2 = self.move_to(pos=new_pos)
                if not (s1 and s2):
                    return False
            
            waypoints = 5
            ee_pos = self.robots[0].controller.ee_pos
            for i in range(waypoints):
                new_pos = (precontact - ee_pos)/waypoints*(i+1) + ee_pos
                s1, s2 = self.move_to(pos=new_pos)
                if not (s1 and s2):
                    return False

        s1, s2 = self.move_to(pos=location, suppress_collision_warning=True) # it's ok if the object moves a little during contact
        return s1

    def move_to(self, pos, quat=None, set_qpos=False, run_simulation_time=None,
                suppress_collision_warning=False):
        # Designed to move in free space

        # Save original object pose
        obj_pos = np.array(self.sim.data.body_xpos[self.cube_body_id])
        obj_mat = np.array(self.sim.data.body_xmat[self.cube_body_id].reshape([3, 3]))

        # Visualize target eef pose
        target_eef_id = self.sim.model.body_name2id('target_eef')
        self.sim.model.body_pos[target_eef_id] = pos
        if quat is not None:
            self.sim.model.body_quat[target_eef_id] = quat
        self.sim.forward()
        
        site_id = self.sim.model.site_name2id('target_eef:grip_site')
        ee_pos = np.array(self.sim.data.site_xpos[site_id])
        ee_ori_mat = np.array(self.sim.data.site_xmat[site_id].reshape([3, 3]))
        
        if set_qpos:
            assert self.robots[0].name == 'Virtual'
            # Set the floating gripper to the location directly (ik)
            eef_body_id = self.sim.model.body_name2id('gripper0_eef')
            offset = self.sim.model.body_pos[eef_body_id]
            qpos = ee_pos + offset
            
            self.sim.data.set_joint_qpos('robot0_slide_joint0', qpos[0])
            self.sim.data.set_joint_qpos('robot0_slide_joint1', qpos[1])
            self.sim.data.set_joint_qpos('robot0_slide_joint2', qpos[2])
            
            self.sim.forward()
            self.sim.step()
            self.robots[0].controller.update(force=True)
            self.maybe_render()
        else:
            self.run_simulation(ee_pos=ee_pos, ee_ori_mat=ee_ori_mat,
                                total_execution_time=run_simulation_time)

        # Check if the gripper reaches the pose
        reachable = True
        ee_pos_after = self.robots[0].controller.ee_pos
        ee_ori_mat_after = self.robots[0].controller.ee_ori_mat.reshape(3,3)
        pos_diff = np.linalg.norm(ee_pos_after - ee_pos)
        ori_diff = np.abs(orientation_error(ee_ori_mat_after, ee_ori_mat)/np.pi*180)
        if pos_diff > 0.02 or np.any(ori_diff > 5):
            if VERBOSE:
                print(f"fail: cannot reach the ee pose.\t Current pose: {ee_pos_after} \t"
                      f"Desired Pose: {ee_pos}\t Diff:{pos_diff}\t{ori_diff}")
            reachable = False

        # Check if the object has moved
        collision_free = True
        obj_pos_after = np.array(self.sim.data.body_xpos[self.cube_body_id])
        obj_mat_after = np.array(self.sim.data.body_xmat[self.cube_body_id].reshape(3,3))
        pos_diff = np.linalg.norm(obj_pos_after - obj_pos)
        ori_diff = np.abs(orientation_error(obj_mat_after, obj_mat)/np.pi*180)
        if pos_diff > 0.02 or np.any(ori_diff > 5):
            if VERBOSE and not suppress_collision_warning:
                print(f"fail: object moved during initialization.\t {pos_diff}\t{ori_diff}")
            collision_free = False

        # Return success: the location is reachable and the object didn't move
        return reachable, collision_free
    

    def move_to_from_top(self, location, start_location=None):
        # First move to the start location
        s1, s2 = self.move_to(start_location, set_qpos=True)
        self.maybe_render()
        if not (s1 and s2):
            return False
        
        # Generating several waypoints to move the gripper from the top
        # Calls move_to to move between waypoints later
        site_id = self.sim.model.site_name2id('target_eef:grip_site')
        ee_pos = np.array(self.sim.data.site_xpos[site_id])
        ee_ori_mat = np.array(self.sim.data.site_xmat[site_id].reshape([3, 3]))
        ee_z = ee_pos[2]
        loc_z = location[2]

        # Create the waypoints
        assert np.isclose(ee_pos[:2], location[:2], atol=0.01).all() , "Warning: the gripper is not right above the target location."
        waypoints_z = np.arange(loc_z, ee_z, 0.02)[::-1]
        waypoints = [[location[0], location[1], z] for z in waypoints_z]
        
        # Move to each waypoint
        for waypoint in waypoints:
            reachable, collision_free = self.move_to(waypoint, set_qpos=False)
            self.maybe_render()

            if not reachable:
                break
        
        return True


    def run_simulation(self, action=None, ee_pos=None, ee_ori_mat=None, total_execution_time=None):
        # Run simulation without taking a env.step
        action = np.zeros(self.action_dim) if action is None else action
        total_execution_time = self.control_timestep if total_execution_time is None else total_execution_time
        for i in range(int(total_execution_time/self.control_timestep)):
            self.robots[0].controller.set_goal(action, set_pos=ee_pos, set_ori=ee_ori_mat)
            for _ in range(int(self.control_timestep / self.model_timestep)):
                self.sim.forward()
                torques = self.robots[0].controller.run_controller()
                low, high = self.robots[0].torque_limits
                torques = np.clip(torques, low, high)
                if self.close_gripper:
                    gripper_action = 1
                else:
                    gripper_action = -1
                self.robots[0].grip_action(gripper=self.robots[0].gripper, gripper_action=[gripper_action])
                self.sim.data.ctrl[self.robots[0]._ref_joint_actuator_indexes] = torques
                self.sim.step()
            self.maybe_render()

    def maybe_render(self):
        if self.has_offscreen_renderer:
            if self.record_from_cam is not None:
                self.show_goal()
                img = self.sim.render(
                    camera_name=self.record_from_cam,
                    width=1280,
                    height=720)
                self.cam_frames.append(img)
                # import matplotlib.pyplot as plt
                # import time
                # plt.imsave(f'agent_{self.cube.mesh_name}_{int(time.time())}.png', img[-1::-1, :, :])
                self.hide_goal()
            
        elif self.has_renderer:
            # Always render intermediate steps if there is an onscreen renderer.
            # Turn this off by setting renderer to None.
            self.show_goal()
            super().render()
            self.hide_goal()
            

    """
    Point cloud related
    """
    def _get_observations(self, force_update=False):
        obs = super()._get_observations(force_update=force_update)

        if self.has_offscreen_renderer:
            pcd_dict = self.get_point_cloud(obs)

            # Try 5 times:
            if self.pcd_noise is not None and pcd_dict['object_pcd_points'] is None:
                for i in range(5):
                    print(f'resample pcd: {i}')
                    pcd_dict = self.get_point_cloud(obs)
                    if pcd_dict['object_pcd_points'] is not None:
                        break
            obs.update(pcd_dict)
        if self.goal is not None:
            obs['goal'] = self.goal

        # Add oject size to the observation
        if hasattr(self.cube, 'scaled_size'):
            obs['object_size'] = self.cube.scaled_size
        elif hasattr(self.cube, 'size'):
            obs['object_size'] = self.cube.size
        else:
            raise ValueError("No object size found.")
            
        return obs

    def get_point_cloud(self, obs):
        box_in_the_view = True
        pcd_dict = {}
        pc_list = []
        pc_color_list = []
        mask_list = []

        for cam in self.camera_names:
            color = obs[cam + '_image'][-1::-1]
            depth = obs[cam + '_depth'][-1::-1][:, :, 0]
            if self.pcd_noise is not None:
                depth = dropout_random_ellipses(depth_img=depth, dropout_mean=self.pcd_noise[0], 
                                                gamma_shape=self.pcd_noise[1], gamma_scale=self.pcd_noise[2])

            # Segmentation from MuJoCo
            seg = obs[cam + '_segmentation_instance'][-1::-1]
            obj_mask = (seg == 1).reshape(color.shape[0], color.shape[1])

            depth = convert_depth(self, depth)
            pc = get_point_cloud(self, depth, camera_name=cam)
            if self.pcd_noise is not None:
                pc = add_additive_noise_to_xyz(pc.reshape(color.shape), gaussian_scale_range=[0.0, self.pcd_noise[3]]).reshape(-1, 3)
            pc_color = color.reshape(-1, 3).astype(np.float64)/256
            
            pc_list.append(pc)
            pc_color_list.append(pc_color)
            mask_list.append(obj_mask.reshape(-1))

        pc = np.concatenate(pc_list)
        pc_color = np.concatenate(pc_color_list)
        pc_mask = np.concatenate(mask_list)
        
        # Remove unnecessary points outside of the bin
        scene_mask = self.within_arena(pc, margin=np.array([-0.01, -0.01, -0.1]), check_z=True)
        pc = pc[scene_mask]
        pc_color = pc_color[scene_mask]
        pc_mask = pc_mask[scene_mask]
        if pc_mask.sum() <= 10:
            # import matplotlib.pyplot as plt
            # import time
            # plt.imsave(f'agent_{self.cube.mesh_name}_{int(time.time())}.png', obs['agentview_image'][-1::-1, :, :])
            # plt.imsave(f'left_{self.cube.mesh_name}_{int(time.time())}.png', obs['leftview_image'][-1::-1, :, :])
            # plt.imsave(f'right_{self.cube.mesh_name}_{int(time.time())}.png', obs['rightview_image'][-1::-1, :, :])
            box_in_the_view = False
        
        if box_in_the_view:
            try:
                # Downsample object pcd
                object_pcd = o3d.geometry.PointCloud()
                object_pcd.points = o3d.utility.Vector3dVector(pc[pc_mask])
                object_pcd.colors = o3d.utility.Vector3dVector(pc_color[pc_mask])
                # print(len(object_pcd.points))
                # o3d.visualization.draw_geometries([object_pcd])
                object_pcd = object_pcd.voxel_down_sample(self.object_voxel_down_sample)
                object_pcd.estimate_normals()
                object_pcd.orient_normals_consistent_tangent_plane(10)
                # print(len(object_pcd.points))
                # o3d.visualization.draw_geometries([object_pcd])
            
            except:
                box_in_the_view = False
        
        # Downsample background pcd
        bg_pcd = o3d.geometry.PointCloud()
        bg_pcd.points = o3d.utility.Vector3dVector(pc[~pc_mask])
        bg_pcd.colors = o3d.utility.Vector3dVector(pc_color[~pc_mask])
        # print(len(bg_pcd.points))
        # o3d.visualization.draw_geometries([bg_pcd])
        bg_pcd = bg_pcd.voxel_down_sample(0.02)
        # print(len(bg_pcd.points))
        # o3d.visualization.draw_geometries([bg_pcd])
        
        if box_in_the_view:
            pcd_dict['object_pcd_points'] = np.asarray(object_pcd.points)
            pcd_dict['object_pcd_colors'] = np.asarray(object_pcd.colors)
            pcd_dict['object_pcd_normals'] = np.asarray(object_pcd.normals)
        else:
            pcd_dict['object_pcd_points'] = None
            pcd_dict['object_pcd_colors'] = None
            pcd_dict['object_pcd_normals'] = None
        pcd_dict['background_pcd_points'] = np.asarray(bg_pcd.points)
        pcd_dict['background_pcd_colors'] = np.asarray(bg_pcd.colors)
        return pcd_dict

    def within_arena(self, points, margin=None, check_z=False):
        # Check if the points are within the box on the table
        # Input: N x 3
        if margin is None:
            margin = np.zeros(3)
        x_min = self.table_offset[0] - self.table_size_curr[0]/2 + margin[0]
        x_max = self.table_offset[0] + self.table_size_curr[0]/2 - margin[0]
        y_min = self.table_offset[1] - self.table_size_curr[1]/2 + margin[1]
        y_max = self.table_offset[1] + self.table_size_curr[1]/2 - margin[1]
        valid = (points[:, 0] >= x_min) * (points[:, 0] <= x_max) * \
                (points[:, 1] >= y_min) * (points[:, 1] <= y_max)
        if check_z:
            z_min = self.table_offset[2] + margin[2]
            z_max = self.table_offset[2] + self.table_size_curr[2] - margin[2]
            valid *= (points[:, 2] >= z_min) * (points[:, 2] <= z_max)
        return valid
    
    """
    Others
    """
    def reward(self, action=None):
        return 0
    
    def sample_goal(self, obs):
        scale = None
        if self.goal_list is not None:
            if hasattr(self, 'cube') and hasattr(self.cube, 'mesh_name'):
                mesh_name = self.cube.mesh_name
                scale = self.cube.scale

                # Find the goals with the same scale
                goal_scales = self.goal_list[mesh_name]['scales']
                scale_mask = np.isclose(np.linalg.norm(goal_scales - scale, axis=1), 0, rtol=1e-2)
                if not np.any(scale_mask):
                    raise ValueError('No goals with the same scale!')
                goals = self.goal_list[mesh_name]['poses'][scale_mask]
            else:
                mesh_name = 'cube'
                goals = self.goal_list[mesh_name]['poses']
            
            goal_idx = np.random.choice(len(goals))
            goal = goals[goal_idx]
            
            
            # default bin size when generating the dataset
            original_scale = np.array([0.45, 0.54, 0.107])
            new_scale = self.table_full_size
            if mesh_name == 'cube' and not np.all(np.isclose(original_scale, new_scale)):
                goal[:2] = ((goal[:3] - self.table_offset)*new_scale/original_scale*0.5 + self.table_offset)[:2]

            # add variation to the goals when there is no wall interactions
            if self.goal_mode in {"any", "any_var_size"}:
                pos_range_lb = self.location_center - self.location_scale * 0.75
                pos_range_ub = self.location_center + self.location_scale * 0.75
                goal_xy = np.random.uniform(pos_range_lb, pos_range_ub)[:2]
                goal[:2] = goal_xy
                # goal[6] = np.random.uniform(-1, 1)
                # goal_pos[2] = obs['cube_pos'][2]

            goal = to_pose_mat(goal[:3], goal[3:])

        elif self.goal_mode == 'fixed':
            goal_pos = self.location_center
            goal_pos[2] = self.get_cube_size()[2] + self.table_offset[2]
            goal_ori = np.array([1., 0., 0., 0.])
            goal = to_pose_mat(goal_pos, goal_ori)
        
        elif self.goal_mode == 'fixed_off_center':
            goal = to_pose_mat(np.array([0.5, 0.1, obs['cube_pos'][2]]), np.array([1., 0., 0., 0.]))
        
        elif self.goal_mode == 'off_center':
            y = np.random.uniform(0.0, 0.2)
            goal = to_pose_mat(np.array([0.5, y, obs['cube_pos'][2]]), np.array([1., 0., 0., 0.]))
            
        elif self.goal_mode == 'translation':
            # Same pose as the object, with x, y randomized 
            # in the range of the bin (scaled by 0.6)
            pos_range_lb = self.location_center - self.location_scale * 0.6
            pos_range_ub = self.location_center + self.location_scale * 0.6
            goal_pos = np.random.uniform(pos_range_lb, pos_range_ub)
            goal_pos[2] = obs['cube_pos'][2]
            goal_ori = obs['cube_quat']
            goal = to_pose_mat(goal_pos, goal_ori)
        
        elif self.goal_mode == 'upright':
            # Sample a x,y location
            pos_range_lb = self.location_center - self.location_scale * 0.6
            pos_range_ub = self.location_center + self.location_scale * 0.6
            goal_pos = np.random.uniform(pos_range_lb, pos_range_ub)
            goal_pos[2] = self.get_cube_size()[2] + self.table_offset[2]
            
            # Sample a planar rotation
            # quat=cos(a/2),sin(a/2)⋅(x,y,z)
            z_angle = np.random.uniform(360)
            goal_ori = np.array([np.cos(z_angle/2), 0, 0, np.sin(z_angle/2)])
            goal = to_pose_mat(goal_pos, goal_ori)
            
        else:
            raise NotImplementedError
             
        return goal
bin_env/poke_env.py
HACMan-2023/HACMan
import numpy as np

from collections import OrderedDict

import robosuite.utils.transform_utils as T

from robosuite.models.grippers import gripper_factory
from robosuite.controllers import controller_factory, load_controller_config

from robosuite.robots.manipulator import Manipulator
from robosuite.utils.buffers import DeltaBuffer, RingBuffer
from robosuite.utils.observables import Observable, sensor
from robosuite.utils.mjcf_utils import array_to_string, find_elements
import robosuite
import os
import copy
from bin_env.osc import OSC
from bin_env.gripper import CustomizedRobotiq85Gripper


class SingleArm(Manipulator):
    """
    Initializes a single-armed robot simulation object.

    Args:
        robot_type (str): Specification for specific robot arm to be instantiated within this env (e.g: "Panda")

        idn (int or str): Unique ID of this robot. Should be different from others

        controller_config (dict): If set, contains relevant controller parameters for creating a custom controller.
            Else, uses the default controller for this specific task

        initial_qpos (sequence of float): If set, determines the initial joint positions of the robot to be
            instantiated for the task

        initialization_noise (dict): Dict containing the initialization noise parameters. The expected keys and
            corresponding value types are specified below:

            :`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
                joint positions. Setting this value to "None" or 0.0 results in no noise being applied.
                If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
                If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
            :`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"

            :Note: Specifying None will automatically create the required dict with "magnitude" set to 0.0

        mount_type (str): type of mount, used to instantiate mount models from mount factory.
            Default is "default", which is the default mount associated with this robot's corresponding model.
            None results in no mount, and any other (valid) model overrides the default mount.

        gripper_type (str): type of gripper, used to instantiate
            gripper models from gripper factory. Default is "default", which is the default gripper associated
            within the 'robot' specification. None removes the gripper, and any other (valid) model overrides the
            default gripper

        control_freq (float): how many control signals to receive
            in every second. This sets the amount of simulation time
            that passes between every action input.
    """

    def __init__(
        self,
        robot_type: str,
        idn=0,
        controller_config=None,
        initial_qpos=None,
        initialization_noise=None,
        mount_type="default",
        gripper_type="default",
        control_freq=20,
        control_gripper=True,
        gripper_friction=(3, 0.015, 0.0003),
        close_gripper=False,
    ):
        self.robot_type = robot_type
        self.controller = None
        self.controller_config = copy.deepcopy(controller_config)
        self.gripper_type = gripper_type
        self.has_gripper = self.gripper_type is not None
        self.control_gripper = control_gripper and self.has_gripper
        self.gripper_friction = np.array(gripper_friction)

        self.gripper = None                                 # Gripper class
        self.gripper_joints = None                          # xml joint names for gripper
        self.close_gripper = close_gripper                  # keep the gripper close during the manipulation
        self._ref_gripper_joint_pos_indexes = None          # xml gripper joint position indexes in mjsim
        self._ref_gripper_joint_vel_indexes = None          # xml gripper joint velocity indexes in mjsim
        self._ref_joint_gripper_actuator_indexes = None     # xml gripper (pos) actuator indexes for robot in mjsim
        self.torques = None                                 # Current torques being applied

        self.recent_ee_forcetorques = None                  # Current and last forces / torques sensed at eef
        self.recent_ee_pose = None                          # Current and last eef pose (pos + ori (quat))
        self.recent_ee_vel = None                           # Current and last eef velocity
        self.recent_ee_vel_buffer = None                    # RingBuffer holding prior 10 values of velocity values
        self.recent_ee_acc = None                           # Current and last eef acceleration

        super().__init__(
            robot_type=robot_type,
            idn=idn,
            initial_qpos=initial_qpos,
            initialization_noise=initialization_noise,
            mount_type=mount_type,
            control_freq=control_freq,
        )

    def _load_controller(self):
        """
        Loads controller to be used for dynamic trajectories
        """
        # First, load the default controller if none is specified
        if not self.controller_config:
            # Need to update default for a single agent
            controller_path = os.path.join(os.path.dirname(robosuite.__file__),
                                           'controllers/config/{}.json'.format(
                                               self.robot_model.default_controller_config))
            self.controller_config = load_controller_config(custom_fpath=controller_path)

        # Assert that the controller config is a dict file:
        #             NOTE: "type" must be one of: {JOINT_POSITION, JOINT_TORQUE, JOINT_VELOCITY,
        #                                           OSC_POSITION, OSC_POSE, IK_POSE}
        assert type(self.controller_config) == dict, \
            "Inputted controller config must be a dict! Instead, got type: {}".format(type(self.controller_config))

        # Add to the controller dict additional relevant params:
        #   the robot name, mujoco sim, eef_name, joint_indexes, timestep (model) freq,
        #   policy (control) freq, and ndim (# joints)
        self.controller_config["robot_name"] = self.name
        self.controller_config["sim"] = self.sim
        self.controller_config["eef_name"] = self.gripper.important_sites["grip_site"]
        self.controller_config["joint_indexes"] = {
            "joints": self.joint_indexes,
            "qpos": self._ref_joint_pos_indexes,
            "qvel": self._ref_joint_vel_indexes
                                              }
        self.controller_config["actuator_range"] = self.torque_limits
        self.controller_config["policy_freq"] = self.control_freq
        self.controller_config["ndim"] = len(self.robot_joints)

        # Instantiate the relevant controller
        if "OSC_XZPLANE" in self.controller_config["type"]:
            self.controller = OscXZPlane(**self.controller_config)
        elif "OSC_customized" in self.controller_config["type"]:
            self.controller = OSC(**self.controller_config)
        else:
            self.controller = controller_factory(self.controller_config["type"], self.controller_config)

    def load_model(self):
        """
        Loads robot and optionally add grippers.
        """
        # First, run the superclass method to load the relevant model
        super().load_model()

        # Verify that the loaded model is of the correct type for this robot
        if self.robot_model.arm_type != "single":
            raise TypeError("Error loading robot model: Incompatible arm type specified for this robot. "
                            "Requested model arm type: {}, robot arm type: {}"
                            .format(self.robot_model.arm_type, type(self)))

        # Now, load the gripper if necessary
        if self.has_gripper:
            if self.gripper_type == 'default':
                # Load the default gripper from the robot file
                self.gripper = gripper_factory(self.robot_model.default_gripper, idn=self.idn)
            elif self.gripper_type == 'robotiq':
                self.gripper = CustomizedRobotiq85Gripper(idn=self.idn)
            else:
                # Load user-specified gripper
                self.gripper = gripper_factory(self.gripper_type, idn=self.idn)
        else:
            # Load null gripper
            self.gripper = gripper_factory(None, idn=self.idn)

        # Grab eef rotation offset
        # self.eef_rot_offset = T.quat_multiply(self.robot_model.hand_rotation_offset, self.gripper.rotation_offset)

        # Make the gripper stiff so that it won't move after collision
        if not self.control_gripper and self.gripper_type == 'default':
            self.gripper.actuator[0].set("kp", str(1e6))
            self.gripper.actuator[1].set("kp", str(1e6))

        # Increase default friction
        gripper_geoms = find_elements(self.gripper.root, 'geom', return_first=False)
        for geom in gripper_geoms:
            if "collision" in geom.get("name"):
                geom.set("friction", array_to_string(self.gripper_friction))

        # https://github.com/marcocognetti/FrankaEmikaPandaDynModel/blob/master/pdf/RA-L_2019_PandaDynIdent_SUPPLEMENTARY_MATERIAL.pdf
        if self.robot_type == 'Panda':
            damping = np.array([0.0628, 0.2088, 0.0361, 0.2174, 0.1021, 1.6128e-04, 0.0632])
            frictionloss = np.array([5.4615e-01, 0.87224, 6.4068e-01, 1.2794e+00, 8.3904e-01, 3.0301e-01, 5.6489e-01])/2
            self.robot_model.set_joint_attribute(attrib="damping", values=damping, force=True)
            self.robot_model.set_joint_attribute(attrib="frictionloss", values=frictionloss, force=True)
        # TODO: https://github.com/frankaemika/franka_ros/blob/develop/franka_description/robots/panda_gazebo.xacro

        # Add gripper to this robot model
        self.robot_model.add_gripper(self.gripper)

    def reset(self, deterministic=False):
        """
        Sets initial pose of arm and grippers. Overrides gripper joint configuration if we're using a
        deterministic reset (e.g.: hard reset from xml file)

        Args:
            deterministic (bool): If true, will not randomize initializations within the sim
        """
        # First, run the superclass method to reset the position and controller
        # deterministic = True
        super().reset(deterministic)

        if not deterministic and self.gripper_type=='default':
            if self.close_gripper:
                gripper_action = 0
            else:
                gripper_action = 0.04
            # Now, reset the gripper if necessary
            if self.has_gripper:
                self.sim.data.qpos[
                    self._ref_gripper_joint_pos_indexes
                ] = np.array([gripper_action, -gripper_action])  # start with open gripper instead of self.gripper.init_qpos

        # Update base pos / ori references in controller
        self.controller.update_base_pose(self.base_pos, self.base_ori)

        # # Setup buffers to hold recent values
        self.recent_ee_forcetorques = DeltaBuffer(dim=6)
        self.recent_ee_pose = DeltaBuffer(dim=7)
        self.recent_ee_vel = DeltaBuffer(dim=6)
        self.recent_ee_vel_buffer = RingBuffer(dim=6, length=10)
        self.recent_ee_acc = DeltaBuffer(dim=6)

    def setup_references(self):
        """
        Sets up necessary reference for robots, grippers, and objects.

        Note that this should get called during every reset from the environment
        """
        # First, run the superclass method to setup references for joint-related values / indexes
        super().setup_references()

        # Now, add references to gripper if necessary
        # indices for grippers in qpos, qvel
        if self.has_gripper:
            self.gripper_joints = list(self.gripper.joints)
            self._ref_gripper_joint_pos_indexes = [
                self.sim.model.get_joint_qpos_addr(x) for x in self.gripper_joints
            ]
            self._ref_gripper_joint_vel_indexes = [
                self.sim.model.get_joint_qvel_addr(x) for x in self.gripper_joints
            ]
            self._ref_joint_gripper_actuator_indexes = [
                self.sim.model.actuator_name2id(actuator)
                for actuator in self.gripper.actuators
            ]

    def control(self, action, policy_step=False):
        """
        Actuate the robot with the
        passed joint velocities and gripper control.

        Args:
            action (np.array): The control to apply to the robot. The first @self.robot_model.dof dimensions should be
                the desired normalized joint velocities and if the robot has a gripper, the next @self.gripper.dof
                dimensions should be actuation controls for the gripper.
            policy_step (bool): Whether a new policy step (action) is being taken

        Raises:
            AssertionError: [Invalid action dimension]
        """
        # clip actions into valid range
        assert len(action) == self.action_dim, \
            "environment got invalid action dimension -- expected {}, got {}".format(
                self.action_dim, len(action))

        gripper_action = None
        if self.control_gripper:
            gripper_action = action[self.controller.control_dim:]  # all indexes past controller dimension indexes
            arm_action = action[:self.controller.control_dim]
        else:
            arm_action = action

        # Update the controller goal if this is a new policy step
        if policy_step:
            self.controller.set_goal(arm_action)

        # Now run the controller for a step
        torques = self.controller.run_controller()

        # Clip the torques
        low, high = self.torque_limits
        self.torques = np.clip(torques, low, high)

        # Get gripper action, if applicable
        if self.control_gripper:
            self.grip_action(gripper=self.gripper, gripper_action=gripper_action)
        else:
            if self.close_gripper:
                gripper_action = 1
            else:
                gripper_action = -1
            self.grip_action(gripper=self.gripper, gripper_action=[gripper_action])

        # Apply joint torque control
        self.sim.data.ctrl[self._ref_joint_actuator_indexes] = self.torques

        # If this is a policy step, also update buffers holding recent values of interest
        if policy_step:
            # Update proprioceptive values
            self.recent_qpos.push(self._joint_positions)
            self.recent_actions.push(action)
            self.recent_torques.push(self.torques)
            self.recent_ee_forcetorques.push(np.concatenate((self.ee_force, self.ee_torque)))  # gripper base
            self.recent_ee_pose.push(np.concatenate((self.controller.ee_pos, T.mat2quat(self.controller.ee_ori_mat))))  # grip_site
            self.recent_ee_vel.push(np.concatenate((self.controller.ee_pos_vel, self.controller.ee_ori_vel)))  # grip_site

            # Estimation of eef acceleration (averaged derivative of recent velocities)
            self.recent_ee_vel_buffer.push(np.concatenate((self.controller.ee_pos_vel, self.controller.ee_ori_vel)))
            diffs = np.vstack([self.recent_ee_acc.current,
                               self.control_freq * np.diff(self.recent_ee_vel_buffer.buf, axis=0)])
            ee_acc = np.array([np.convolve(col, np.ones(10) / 10., mode='valid')[0] for col in diffs.transpose()])
            self.recent_ee_acc.push(ee_acc)

    def _visualize_grippers(self, visible):
        """
        Visualizes the gripper site(s) if applicable.

        Args:
            visible (bool): True if visualizing the gripper for this arm.
        """
        self.gripper.set_sites_visibility(sim=self.sim, visible=visible)

    def setup_observables(self):
        """
        Sets up observables to be used for this robot

        Returns:
            OrderedDict: Dictionary mapping observable names to its corresponding Observable object
        """
        # Get general robot observables first
        observables = super().setup_observables()

        # Get prefix from robot model to avoid naming clashes for multiple robots and define observables modality
        pf = self.robot_model.naming_prefix
        modality = f"{pf}proprio"

        # finger-tip
        @sensor(modality=modality)
        def eef_pos(obs_cache):
            return self.controller.ee_pos

        @sensor(modality=modality)
        def eef_quat(obs_cache):
            return T.convert_quat(T.mat2quat(self.controller.ee_ori_mat), to='wxyz')

        @sensor(modality=modality)
        def eef_pos_vel(obs_cache):
            return self.controller.ee_pos_vel

        @sensor(modality=modality)
        def eef_ori_vel(obs_cache):
            return self.controller.ee_ori_vel

        @sensor(modality=modality)
        def joint_torque(obs_cache):
            return self.controller.torques

        sensors = [eef_pos, eef_quat, eef_pos_vel, eef_ori_vel, joint_torque]
        names = [f"{pf}eef_pos", f"{pf}eef_quat", f"{pf}eef_pos_vel", f"{pf}eef_ori_vel", f"{pf}joint_torque"]

        if 'OSC' in self.controller.name:
            @sensor(modality=modality)
            def osc_desired_pos(obs_cache):
                return self.controller.goal_pos

            @sensor(modality=modality)
            def osc_desired_quat(obs_cache):
                return T.convert_quat(T.mat2quat(self.controller.goal_ori), to='wxyz')

            sensors += [osc_desired_pos, osc_desired_quat]
            names += [f"{pf}osc_desired_pos", f"{pf}osc_desired_quat"]

        # add in gripper sensors of the finger movement if this robot has a gripper
        if self.has_gripper:
            @sensor(modality=modality)
            def gripper_qpos(obs_cache):
                return np.array([self.sim.data.qpos[x] for x in self._ref_gripper_joint_pos_indexes])

            @sensor(modality=modality)
            def gripper_qvel(obs_cache):
                return np.array([self.sim.data.qvel[x] for x in self._ref_gripper_joint_vel_indexes])

            sensors += [gripper_qpos, gripper_qvel]
            names += [f"{pf}gripper_qpos", f"{pf}gripper_qvel"]

        # Create observables for this robot
        for name, s in zip(names, sensors):
            observables[name] = Observable(
                name=name,
                sensor=s,
                sampling_rate=self.control_freq,
            )

        return observables

    @property
    def dof(self):
        """
        Returns:
            int: degrees of freedom of the robot (with grippers).
        """
        # Get the dof of the base robot model
        dof = super().dof
        if self.control_gripper:
            for gripper in self.robot_model.grippers.values():
                dof += gripper.dof
        return dof

    @property
    def action_limits(self):
        """
        Action lower/upper limits per dimension.

        Returns:
            2-tuple:

                - (np.array) minimum (low) action values
                - (np.array) maximum (high) action values
        """
        # Action limits based on controller limits
        low, high = ([-1] * self.gripper.dof, [1] * self.gripper.dof) if self.control_gripper else ([], [])
        low_c, high_c = self.controller.control_limits
        low = np.concatenate([low_c, low])
        high = np.concatenate([high_c, high])

        return low, high

    @property
    def ee_ft_integral(self):
        """
        Returns:
            np.array: the integral over time of the applied ee force-torque
        """
        return np.abs((1.0 / self.control_freq) * self.recent_ee_forcetorques.average)

    @property
    def ee_force(self):
        """
        Returns:
            np.array: force applied at the force sensor at the robot arm's eef
        """
        return self.get_sensor_measurement(self.gripper.important_sensors["force_ee"])

    @property
    def ee_torque(self):
        """
        Returns torque applied at the torque sensor at the robot arm's eef
        """
        return self.get_sensor_measurement(self.gripper.important_sensors["torque_ee"])

    @property
    def _hand_pose(self):
        """
        Returns:
            np.array: (4,4) array corresponding to the eef pose in base frame of robot.
        """
        return self.pose_in_base_from_name(self.robot_model.eef_name)

    @property
    def _hand_quat(self):
        """
        Returns:
            np.array: (x,y,z,w) eef quaternion in base frame of robot.
        """
        return T.mat2quat(self._hand_orn)

    @property
    def _hand_total_velocity(self):
        """
        Returns:
            np.array: 6-array representing the total eef velocity (linear + angular) in the base frame
        """

        # Use jacobian to translate joint velocities to end effector velocities.
        Jp = self.sim.data.get_body_jacp(self.robot_model.eef_name).reshape((3, -1))
        Jp_joint = Jp[:, self._ref_joint_vel_indexes]

        Jr = self.sim.data.get_body_jacr(self.robot_model.eef_name).reshape((3, -1))
        Jr_joint = Jr[:, self._ref_joint_vel_indexes]

        eef_lin_vel = Jp_joint.dot(self._joint_velocities)
        eef_rot_vel = Jr_joint.dot(self._joint_velocities)
        return np.concatenate([eef_lin_vel, eef_rot_vel])

    @property
    def _hand_pos(self):
        """
        Returns:
            np.array: 3-array representing the position of eef in base frame of robot.
        """
        eef_pose_in_base = self._hand_pose
        return eef_pose_in_base[:3, 3]

    @property
    def _hand_orn(self):
        """
        Returns:
            np.array: (3,3) array representing the orientation of eef in base frame of robot as a rotation matrix.
        """
        eef_pose_in_base = self._hand_pose
        return eef_pose_in_base[:3, :3]

    @property
    def _hand_vel(self):
        """
        Returns:
            np.array: (x,y,z) velocity of eef in base frame of robot.
        """
        return self._hand_total_velocity[:3]

    @property
    def _hand_ang_vel(self):
        """
        Returns:
            np.array: (ax,ay,az) angular velocity of eef in base frame of robot.
        """
        return self._hand_total_velocity[3:]

bin_env/single_arm.py
HACMan-2023/HACMan
from bin_env.utils.rotations import *
import numpy as np
import robosuite.utils.transform_utils as T

def angle_diff(quat_a, quat_b):
    # Subtract quaternions and extract angle between them.
    quat_diff = quat_mul(quat_a, quat_conjugate(quat_b))
    a_diff = 2 * np.arccos(np.clip(quat_diff[..., 0], -1., 1.))
    return min(a_diff, np.pi*2 - a_diff)


def convert_to_batch(arr):
    if len(arr.shape) == 1:
        arr = arr[np.newaxis]
    return arr


def quat_rot_vec_arr(q, v0):
    q_v0 = np.array([np.zeros_like(v0[...,0]), v0[...,0], v0[...,1], v0[...,2]]).transpose()
    q_v = quat_mul(q, quat_mul(q_v0, quat_conjugate(q)))
    v = q_v[..., 1:]
    return v


def get_global_pose(frame_pos, frame_quat, pos_in_frame, quat_in_frame):
    # Convert the grasp pose in the object coordinate to the global coordinate
    quat = quat_mul(frame_quat, quat_in_frame)
    pos = quat_rot_vec_arr(frame_quat, pos_in_frame) + frame_pos
    return pos, quat


def get_local_pose(frame_pos, frame_quat, global_pos, global_quat):
    # Convert the grasp pose in the global coordinate to the object coordinate
    local_quat = quat_mul(quat_conjugate(frame_quat), global_quat)
    local_pos = quat_rot_vec_arr(quat_conjugate(frame_quat), global_pos - frame_pos)
    return local_pos, local_quat


def clean_xzplane_pose(pos, quat, gripper_correction=False):
    if gripper_correction:
        rotate_back = euler2quat(np.array([0, -np.pi / 2, 0]))  # rotate 180 degree
        quat = quat_mul(quat, rotate_back)
    euler = quat2euler(quat)
    return np.array([pos[0], pos[2], euler[1]])


def clean_4d_pose(pos, quat, gripper_correction=False):
    if gripper_correction:
        rotate_back = euler2quat(np.array([0, -np.pi / 2, 0]))  # rotate 180 degree
        quat = quat_mul(quat, rotate_back)
    euler = quat2euler(quat)
    return np.array([pos[0], pos[1], pos[2], euler[1]])


def clean_6d_pose(pos, quat, gripper_correction=False):
    return np.concatenate([pos, quat])


def sample_unit_vector():
    v = np.random.randn(3)
    v /= np.linalg.norm(v)
    return v


def add_pose_noise(pos, quat, pos_noise=0, ori_noise=0):
    noisy_pos = pos + sample_unit_vector() * pos_noise

    axis = sample_unit_vector()
    angle = np.random.uniform(0, ori_noise/180*np.pi)
    additional_quat = np.array([np.cos(angle/2), np.sin(angle)*axis[0], np.sin(angle)*axis[1], np.sin(angle)*axis[2]]) #wxyz
    noisy_quat = quat_mul(additional_quat, quat)
    return noisy_pos, noisy_quat

def get_error(R_list, t_list, data):
    "given the predicted translation and rotation, compute the error"
    rotation_loss_list = []
    translation_loss_list = []
    succ_list = []
    for k, rotation in enumerate(R_list):
        current_obj_mat = data['current_obj_pose'].reshape(-1,4,4)[k][:3,:3]
        current_obj_position = data['current_obj_pose'].reshape(-1,4,4)[k][:3,3]
        next_obj_mat = data['next_obj_pose'].reshape(-1,4,4)[k][:3,:3]
        next_obj_position = data['next_obj_pose'].reshape(-1,4,4)[k][:3,3]
        pred_next_obj_mat = R_list[k] @ current_obj_mat
        rotation_error = angle_diff(T.mat2quat(pred_next_obj_mat.cpu().numpy()), T.mat2quat(next_obj_mat.cpu().numpy()))/np.pi*180
        translation_error = np.linalg.norm((R_list[k]@current_obj_position + t_list[k] - next_obj_position).cpu().numpy())
        succ = np.float64(translation_error<0.03 and rotation_error<10)
        rotation_loss_list.append(rotation_error)
        translation_loss_list.append(translation_error)
        succ_list.append(succ)
    rotation_loss = np.array(rotation_loss_list).mean()
    translation_loss = np.array(translation_loss_list).mean()
    success_rate = np.array(succ_list).mean()
    return rotation_loss, translation_loss, success_rate

def get_gt_transformation_scale(R_list, t_list, data):
    current_obj_position = data['current_obj_pose'].reshape(-1,4,4)
    batch_size = current_obj_position.shape[0]
    gt_rotation_scale = []
    gt_translation_scale = []
    for k in range(batch_size):
        current_obj_mat = data['current_obj_pose'].reshape(-1,4,4)[k][:3,:3]
        current_obj_position = data['current_obj_pose'].reshape(-1,4,4)[k][:3,3]
        next_obj_mat = data['next_obj_pose'].reshape(-1,4,4)[k][:3,:3]
        next_obj_position = data['next_obj_pose'].reshape(-1,4,4)[k][:3,3]
        rotation_diff = angle_diff(T.mat2quat(next_obj_mat.cpu().numpy()), T.mat2quat(current_obj_mat.cpu().numpy()))/np.pi*180
        translation_diff = np.linalg.norm((next_obj_position - current_obj_position).cpu().numpy())
        gt_rotation_scale.append(rotation_diff)
        gt_translation_scale.append(translation_diff)
    rotation_loss_list = []
    translation_loss_list = []
    for k, rotation in enumerate(R_list):
        current_obj_mat = data['current_obj_pose'].reshape(-1,4,4)[k][:3,:3]
        current_obj_position = data['current_obj_pose'].reshape(-1,4,4)[k][:3,3]
        next_obj_mat = data['next_obj_pose'].reshape(-1,4,4)[k][:3,:3]
        next_obj_position = data['next_obj_pose'].reshape(-1,4,4)[k][:3,3]
        pred_next_obj_mat = R_list[k] @ current_obj_mat
        rotation_error = angle_diff(T.mat2quat(pred_next_obj_mat.cpu().numpy()), T.mat2quat(next_obj_mat.cpu().numpy()))/np.pi*180
        translation_error = np.linalg.norm((R_list[k]@current_obj_position + t_list[k] - next_obj_position).cpu().numpy())
        rotation_loss_list.append(rotation_error)
        translation_loss_list.append(translation_error)
    return gt_rotation_scale, rotation_loss_list, gt_translation_scale, translation_loss_list

bin_env/util.py
HACMan-2023/HACMan
import numpy as np
from bin_env.utils.rotations import *
import cv2


# -------------------- Generic ----------------------------
def get_intrinsics(fovy, img_width, img_height):
    # fovy = self.sim.model.cam_fovy[cam_no]
    aspect = float(img_width) / img_height
    fovx = 2 * np.arctan(np.tan(np.deg2rad(fovy) * 0.5) * aspect)
    fovx = np.rad2deg(fovx)
    cx = img_width / 2.
    cy = img_height / 2.
    fx = cx / np.tan(np.deg2rad(fovx / 2.))
    fy = cy / np.tan(np.deg2rad(fovy / 2.))
    K = np.zeros((3,3), dtype=np.float64)
    K[2][2] = 1
    K[0][0] = fx
    K[1][1] = fy
    K[0][2] = cx
    K[1][2] = cy
    return K


def depth2xyz(depth, cam_K):
    h, w = depth.shape
    ymap, xmap = np.meshgrid(np.arange(w), np.arange(h))

    x = ymap
    y = xmap
    z = depth

    x = (x - cam_K[0,2]) * z / cam_K[0,0]
    y = (y - cam_K[1,2]) * z / cam_K[1,1]

    xyz = np.stack([x, y, z], axis=2)
    return xyz


def visualize_point_cloud_from_nparray(d, c=None, vis_coordinate=False):
    if c is not None:
        if len(c.shape) == 3:
            c = c.reshape(-1, 3)
        if c.max() > 1:
            c = c.astype(np.float64)/256

    import open3d as o3d
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(d)
    if c is not None:
        pcd.colors = o3d.utility.Vector3dVector(c)

    if vis_coordinate:
        mesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5)
        o3d.visualization.draw_geometries([mesh, pcd])
    else:
        o3d.visualization.draw_geometries([pcd])


# -------------------- MuJoCo Specific ----------------------------
def get_transformation_matrix(pos, quat):
    arr = np.identity(4)
    arr[:3, :3] = quat2mat(quat)
    arr[:3, 3] = pos
    return arr


def get_transformation(env, camera_name=None):
    if camera_name is None:
        camera_name = env.camera_names[0]
    cam_id = env.sim.model.camera_name2id(camera_name)
    cam_pos = env.sim.model.cam_pos[cam_id]
    cam_quat = env.sim.model.cam_quat[cam_id]
    cam_quat = quat_mul(cam_quat, euler2quat([np.pi, 0, 0]))
    return get_transformation_matrix(cam_pos, cam_quat)


def convert_depth(env, depth):
    # Convert depth into meter
    extent = env.sim.model.stat.extent
    near = env.sim.model.vis.map.znear * extent
    far = env.sim.model.vis.map.zfar * extent
    depth_m = depth * 2 - 1
    depth_m = (2 * near * far) / (far + near - depth_m * (far - near))
    # Check this as well: https://github.com/deepmind/dm_control/blob/master/dm_control/mujoco/engine.py#L734
    return depth_m


def get_object_point_cloud(env, depth, img):
    depth = convert_depth(env, depth)
    full_pc = get_point_cloud(env, depth)
    obj_mask = get_obj_mask(img)
    pc = full_pc[obj_mask.reshape(-1),:]
    return pc


def get_point_cloud(env, depth, camera_name=None):
    # make sure to convert the raw depth image from MuJoCo using convert_depth
    # output is flattened
    if camera_name is None:
        camera_name = env.camera_names[0]
    fovy = env.sim.model.cam_fovy[env.sim.model.camera_name2id(camera_name)]
    K = get_intrinsics(fovy, depth.shape[0], depth.shape[1])
    pc = depth2xyz(depth, K)
    pc = pc.reshape(-1, 3)

    transform = get_transformation(env, camera_name=camera_name)
    new_pc = np.ones((pc.shape[0], 4))
    new_pc[:, :3] = pc
    new_pc = (transform @ new_pc.transpose()).transpose()
    return new_pc[:, :-1]


def get_obj_mask(img):
    # Given an [n,n,3] image, output a [n,n] mask of the red object in the scene
    # Assume the object is red
    img = img[:, :, -1::-1]  # RGB to BGR
    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # lower mask (0-10)
    lower_red = np.array([0, 50, 50])
    upper_red = np.array([10, 255, 255])
    mask0 = cv2.inRange(img_hsv, lower_red, upper_red)

    # upper mask (170-180)
    lower_red = np.array([170, 50, 50])
    upper_red = np.array([180, 255, 255])
    mask1 = cv2.inRange(img_hsv, lower_red, upper_red)

    # join my masks
    mask = mask0 + mask1

    # # set my output img to zero everywhere except my mask
    # output_img = img.copy()
    # output_img[np.where(mask == 0)] = 0
    # cv2.imwrite('color_img.jpg', output_img)
    return mask != 0


def add_additive_noise_to_xyz(
    xyz_img,
    gp_rescale_factor_range=[12, 20],
    gaussian_scale_range=[0.0, 0.003],
    valid_mask=None,
):
    """Add (approximate) Gaussian Process noise to ordered point cloud
    @param xyz_img: a [H x W x 3] ordered point cloud
    """
    xyz_img = xyz_img.copy()

    H, W, C = xyz_img.shape

    # Additive noise: Gaussian process, approximated by zero-mean anisotropic Gaussian random variable,
    #                 which is rescaled with bicubic interpolation.
    gp_rescale_factor = np.random.randint(
        gp_rescale_factor_range[0], gp_rescale_factor_range[1]
    )
    gp_scale = np.random.uniform(gaussian_scale_range[0], gaussian_scale_range[1])

    small_H, small_W = (np.array([H, W]) / gp_rescale_factor).astype(int)
    additive_noise = np.random.normal(
        loc=0.0, scale=gp_scale, size=(small_H, small_W, C)
    )
    additive_noise = cv2.resize(additive_noise, (W, H), interpolation=cv2.INTER_CUBIC)
    if valid_mask is not None:
        # use this to add to the image
        xyz_img[valid_mask, :] += additive_noise[valid_mask, :]
    else:
        xyz_img += additive_noise

    return xyz_img


def dropout_random_ellipses(
    depth_img, dropout_mean, gamma_shape=10000, gamma_scale=0.0001
):
    """Randomly drop a few ellipses in the image for robustness.
    This is adapted from the DexNet 2.0 code.
    Their code: https://github.com/BerkeleyAutomation/gqcnn/blob/75040b552f6f7fb264c27d427b404756729b5e88/gqcnn/sgd_optimizer.py
    @param depth_img: a [s] set of depth z values
    """
    depth_img = depth_img.copy()

    # Sample number of ellipses to dropout
    num_ellipses_to_dropout = np.random.poisson(dropout_mean)

    # Sample ellipse centers
    nonzero_pixel_indices = np.array(
        np.where(depth_img > 0)
    ).T  # Shape: [#nonzero_pixels x 2]
    dropout_centers_indices = np.random.choice(
        nonzero_pixel_indices.shape[0], size=num_ellipses_to_dropout
    )
    dropout_centers = nonzero_pixel_indices[
        dropout_centers_indices, :
    ]  # Shape: [num_ellipses_to_dropout x 2]

    # Sample ellipse radii and angles
    x_radii = np.random.gamma(gamma_shape, gamma_scale, size=num_ellipses_to_dropout)
    y_radii = np.random.gamma(gamma_shape, gamma_scale, size=num_ellipses_to_dropout)
    angles = np.random.randint(0, 360, size=num_ellipses_to_dropout)

    # Dropout ellipses
    for i in range(num_ellipses_to_dropout):
        center = dropout_centers[i, :]
        x_radius = np.round(x_radii[i]).astype(int)
        y_radius = np.round(y_radii[i]).astype(int)
        angle = angles[i]

        # dropout the ellipse
        # mask is always 2d even if input is not
        mask = np.zeros(depth_img.shape[:2])
        mask = cv2.ellipse(
            mask,
            tuple(center[::-1]),
            (x_radius, y_radius),
            angle=angle,
            startAngle=0,
            endAngle=360,
            color=1,
            thickness=-1,
        )
        depth_img[mask == 1] = 0

    return depth_img

bin_env/utils/point_cloud_utils.py
HACMan-2023/HACMan
import numpy as np
import torch
from robosuite.utils import transform_utils
import scipy.spatial.transform as spt


def calculate_non_z_angle_diff(a_, b_, input_wxyz):
    a = np.copy(a_)
    b = np.copy(b_)

    # Convert to xyzw
    if input_wxyz:
        a = a[[1, 2, 3, 0]]
        b = b[[1, 2, 3, 0]]
    
    # Calculate the difference
    a = spt.Rotation.from_quat(a)
    b = spt.Rotation.from_quat(b)

    z = np.array([0, 0, 1])
    a_z = a.apply(z)
    b_z = b.apply(z)

    rad_diff = np.arccos(a_z * b_z / (np.linalg.norm(a_z) * np.linalg.norm(b_z)))
    diff = rad_diff * 180 / np.pi
    
    return diff


def calculate_axis_diff(a_, b_, input_wxyz, axis='z', output_wxyz=True):
    a = np.copy(a_)
    b = np.copy(b_)

    # Convert to xyzw
    if input_wxyz:
        a = a[[1, 2, 3, 0]]
        b = b[[1, 2, 3, 0]]
    
    # Calculate the difference
    a = spt.Rotation.from_quat(a)
    b = spt.Rotation.from_quat(b)
    diff = a.inv() * b

    if axis == 'z':
        z = np.array([0, 0, 1])

    new_z = diff.apply(z)

    # Calculate the angle difference
    rad_diff = np.arccos(new_z * z / (np.linalg.norm(new_z) * np.linalg.norm(z)))
    diff = rad_diff * 180 / np.pi

    return diff

def decompose_pose_mat(pose, cat=False, wxyz=True):
    pos, quat = transform_utils.mat2pose(pose)
    if wxyz:
        quat = transform_utils.convert_quat(quat, to="wxyz")
    if cat:
        return np.concatenate([pos, quat])
    else:
        return pos, quat

def to_pose_mat(pos, quat):
    # Input quat follows wxyz convention (MuJoCo)
    mat = transform_utils.quat2mat(transform_utils.convert_quat(quat))
    pose = transform_utils.make_pose(pos, mat)
    return pose


def transform_point_cloud(current_object_pose, next_object_pose, pointcloud):
    if len(current_object_pose.shape) <= 2:
        current_object_pose = np.expand_dims(current_object_pose, axis=0)
        next_object_pose = np.expand_dims(next_object_pose, axis=0)
        pointcloud = np.expand_dims(pointcloud, axis=0)
    
    bs = current_object_pose.shape[0]
    transformed_pcds = []
    for i in range(bs):
        transform = next_object_pose[i].dot(transform_utils.pose_inv(current_object_pose[i]))
    
        transformed_pcd = np.ones((len(pointcloud[i]), 4))
        transformed_pcd[:, :3] = pointcloud[i]
        transformed_pcd = transform.dot(transformed_pcd.transpose()).transpose()[:, :3]
        transformed_pcds.append(transformed_pcd)
    
    if len(transformed_pcds) == 1:
        transformed_pcd = transformed_pcds[0]
        return transformed_pcd
    
    else:
        transformed_pcds = np.stack(transformed_pcds, axis=0)
        return transformed_pcds
bin_env/utils/transformations.py
HACMan-2023/HACMan
import os, sys
import gym
import torch
from torch import nn
import numpy as np
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
from stable_baselines3.common.type_aliases import TensorDict
from stable_baselines3.common.torch_layers import create_mlp
from hacman.networks.common import init_network
from hacman.utils.transformations import transform_point_cloud, decompose_pose_tensor
from torch_geometric.data import Data, Batch
from torch_geometric.nn import MLP


class PointCloudExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, normalize_pos_param=None, 
                 preprocessing_fn='flow'):
        # Normalize_range is in the form of [offset, scale]
        super().__init__(observation_space, features_dim=128, ) # feature_dim=32 for pt
        config = {"model": "pn",
                  "dropout": 0.,
                  "fps_deterministic": False,
                  "pos_in_feature": False,
                  "normalize_pos": True,
                  "normalize_pos_param": normalize_pos_param}
        if 'batch' in preprocessing_fn:
            # backward compatibility
            preprocessing_fn = preprocessing_fn[:-6]
        self.preprocessing_fn = preprocessing_fn
        self.preprocessing = globals()['preprocessing_'+preprocessing_fn] # Get processing function by name
        input_channels = 1 # Mask channel
        if preprocessing_fn == 'flow' or preprocessing_fn == 'flow_v0':
            input_channels += 3
        elif preprocessing_fn == 'goal_pose' or preprocessing_fn == 'goal_pose_v0':
            input_channels += 12
        elif preprocessing_fn == 'goal_pose_quat':
            input_channels += 7
        self.backbone = init_network(config, input_channels=input_channels, output_channels=[])
        
    def forward(self, observations: TensorDict, per_point_output=False) -> torch.Tensor:
        batch_size = observations['object_pcd_points'].shape[0]

        if 'v0' in self.preprocessing_fn:
            # x2 slower
            data_list = []
            for batch_id in range(batch_size):
                datum = self.preprocessing(observations, batch_id)
                data_list.append(datum)
            data = Batch.from_data_list(data_list)
            out = self.backbone(data.x, data.pos, data.batch)            
        else:
            data_x, data_pos, data_batch = self.preprocessing(observations)
            out = self.backbone(data_x, data_pos, data_batch)
        
        if per_point_output == False:
            out = out.reshape(batch_size, int(out.shape[0]/batch_size), -1)[:, :, :]
            # Only output the feature at poke_idx
            out_list = []
            for batch_id in range(batch_size):
                poke_idx = observations['poke_idx'][batch_id].detach().cpu().numpy()
                out_list.append(out[batch_id, poke_idx, :].reshape(1, -1))
            out = torch.cat(out_list, dim=0)
        
        return out
    
    
class PointCloudGlobalExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict,
                 preprocessing_fn='flow',
                 include_gripper=False, normalize_pos_param=None):
        # Normalize_range is in the form of [offset, scale]
        feature_dim = 1024
        if include_gripper:
            feature_dim += 7
        super().__init__(observation_space, features_dim=feature_dim)
        config = {"model": "pnc",
                  "dropout": 0.,
                  "fps_deterministic": False,
                  "pos_in_feature": False,
                  "normalize_pos": True,
                  "normalize_pos_param": normalize_pos_param}
        
        self.preprocessing_fn = preprocessing_fn
        self.preprocessing = globals()['preprocessing_'+preprocessing_fn] # Get processing function by name
        input_channels = 1 # Mask channel
        if preprocessing_fn == 'flow' or preprocessing_fn == 'flow_v0':
            input_channels += 3
        elif preprocessing_fn == 'goal_pose' or preprocessing_fn == 'goal_pose_v0':
            input_channels += 12
        elif preprocessing_fn == 'goal_pose_quat':
            input_channels += 7
        self.backbone = init_network(config, input_channels=input_channels, output_channels=[])
        self.include_gripper = include_gripper
        
    def forward(self, observations: TensorDict) -> torch.Tensor:
        batch_size = observations['object_pcd_points'].shape[0]

        if 'v0' in self.preprocessing_fn:
            # x2 slower
            data_list = []
            for batch_id in range(batch_size):
                datum = self.preprocessing(observations, batch_id)
                data_list.append(datum)
            data = Batch.from_data_list(data_list)
            out = self.backbone(data.x, data.pos, data.batch)            
        else:
            data_x, data_pos, data_batch = self.preprocessing(observations)
            out = self.backbone(data_x, data_pos, data_batch)
        
        if self.include_gripper:
            gripper_pose = observations['gripper_pose']
            gripper_pose = decompose_pose_tensor(gripper_pose, cat=True)    # Shape: [N, 3 + 4]
            out = torch.cat([out, gripper_pose], dim=1)
        
        return out


class StatesExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, include_gripper=False):
        features_dim = 17
        if include_gripper:
            features_dim += 7
        super().__init__(observation_space, features_dim=features_dim)
        # self.net = create_mlp(32, 128, [64])
        # self.net = MLP([32, 64, 128])
        self.include_gripper = include_gripper

    def forward(self, observations: TensorDict) -> torch.Tensor:
        batch_size = observations['object_pose'].shape[0]
        # for batch_id in range(batch_size):
        #     datum = preprocessing(observations, batch_id)
        #     data_list.append(datum)
        # data = Batch.from_data_list(data_list)
        # out = self.backbone(data.x, data.pos, data.batch)
        object_pose = decompose_pose_tensor(observations['object_pose'], cat=True)
        goal_pose = decompose_pose_tensor(observations['goal_pose'], cat=True)
        object_size = observations['object_size']
        obs_list = [object_pose, goal_pose, object_size]
        
        if self.include_gripper:
            gripper_pose = decompose_pose_tensor(observations['gripper_pose'], cat=True)
            obs_list.append(gripper_pose)
        
        obs = torch.cat(obs_list, dim=1)
        obs = obs.view([batch_size, -1])
        
        return obs


# Input Preprocessing Functions
def preprocessing_flow(obs):
    
    # Flow as point features.
    # Use object pose and goal pose to calculate the flow
    dtype = obs['object_pcd_points'].dtype
    device = obs['object_pcd_points'].device

    # Object points    
    object_pcd_points = obs['object_pcd_points']
    
    transform = obs['goal_pose'] @ torch.linalg.inv(obs['object_pose'])
    # transform = torch.linalg.solve(obs['object_pose'], obs['goal_pose'], left=False) # T*current_pose = goal_pose    
    # faster but requires pytorch 1.13
    
    ones = torch.ones(object_pcd_points.shape[0], object_pcd_points.shape[1], 1, dtype=dtype, device=device)
    padded_points = torch.cat([object_pcd_points, ones], dim=-1)
    flow = torch.matmul(transform, padded_points.transpose(1, 2)).transpose(1, 2)[:,:,:3] - object_pcd_points
    mask = torch.ones((object_pcd_points.shape[0], object_pcd_points.shape[1], 1), dtype=dtype, device=device)
    object_feature = torch.cat([mask, flow], axis=-1)
    
    # Backround points
    background_pcd_points = obs['background_pcd_points']
    background_feature = torch.zeros((background_pcd_points.shape[0], background_pcd_points.shape[1], 4), dtype=dtype, device=device)
                    
    data_pos = torch.cat([object_pcd_points, background_pcd_points], axis=1)
    data_pos = data_pos.reshape(-1, 3)
    data_x = torch.cat([object_feature, background_feature], axis=1)
    data_x = data_x.reshape(-1, data_x.shape[-1])
    data_batch = torch.arange(object_pcd_points.shape[0], device=device).repeat_interleave(object_pcd_points.shape[1] + background_pcd_points.shape[1])
    return data_x, data_pos, data_batch

# Concatenate goal points
def preprocessing_no_flow(obs):
    dtype = obs['object_pcd_points'].dtype
    device = obs['object_pcd_points'].device
    batch_size = obs['object_pcd_points'].shape[0]
    num_obj_points = obs['object_pcd_points'].shape[1]
    num_bg_points = obs['background_pcd_points'].shape[1]

    # Object points    
    object_pcd_points = obs['object_pcd_points']
    
    transform = obs['goal_pose'] @ torch.linalg.inv(obs['object_pose'])
    # transform = torch.linalg.solve(obs['object_pose'], obs['goal_pose'], left=False) # T*current_pose = goal_pose    
    # faster but requires pytorch 1.13
    
    ones = torch.ones(batch_size, num_obj_points, 1, dtype=dtype, device=device)
    padded_points = torch.cat([object_pcd_points, ones], dim=-1)
    transformed_pcd = torch.matmul(transform, padded_points.transpose(1, 2)).transpose(1, 2)[:,:,:3]
    
    # Backround points
    background_pcd_points = obs['background_pcd_points']
                    
    full_points = torch.cat([object_pcd_points, background_pcd_points, transformed_pcd], axis=1)
    
    # Object Mask: Object=1 Backgound=0 Goal=-1
    mask = torch.zeros((batch_size, num_obj_points*2 + num_bg_points, 1), dtype=dtype, device=device)
    mask[:, :num_obj_points] = 1.
    mask[:, num_obj_points+num_bg_points:] = -1.

    # Flatten points across batch
    data_pos = full_points.reshape(-1, 3)
    data_x = mask.reshape(-1, 1)
    data_batch = torch.arange(batch_size, device=device).repeat_interleave(num_obj_points*2 + num_bg_points)
    return data_x, data_pos, data_batch


def preprocessing_goal_pose_quat(obs):
    return preprocessing_goal_pose(obs, quat=True)


def preprocessing_goal_pose(obs, quat=False):
    # Concatenate target transformation in point features
    dtype = obs['object_pcd_points'].dtype
    device = obs['object_pcd_points'].device
    batch_size = obs['object_pcd_points'].shape[0]
    num_obj_points = obs['object_pcd_points'].shape[1]
    num_bg_points = obs['background_pcd_points'].shape[1]

    # Object points    
    object_pcd_points = obs['object_pcd_points']
    
    transform = obs['goal_pose'] @ torch.linalg.inv(obs['object_pose'])
    # transform = torch.linalg.solve(obs['object_pose'], obs['goal_pose'], left=False) # T*current_pose = goal_pose    
    # faster but requires pytorch 1.13    
    
    mask = torch.ones((batch_size, num_obj_points, 1), dtype=dtype, device=device)
    if quat:
        goal_feature_obj = decompose_pose_tensor(transform, cat=True)
    else:
        goal_feature_obj = transform[:, :3, :].reshape(batch_size, -1)
    goal_feature_obj = goal_feature_obj.unsqueeze(1).repeat(1, num_obj_points, 1)
    object_feature = torch.cat([mask, goal_feature_obj], axis=-1)
    
    # Backround points
    background_pcd_points = obs['background_pcd_points']
    if quat:
        goal_feature_bg = torch.zeros(7, device=device, dtype=dtype)
        goal_feature_bg[-1] = 1.
    else:
        goal_feature_bg = torch.eye(4, device=device, dtype=dtype)[:3, :].flatten()
    goal_feature_bg = goal_feature_bg.repeat(batch_size, num_bg_points, 1)
    # goal_feature_bg = decompose_pose_tensor(torch.eye(4, device=full_points.device).unsqueeze(0), cat=True)[0]
    bg_mask = torch.zeros((background_pcd_points.shape[0], background_pcd_points.shape[1], 1), dtype=dtype, device=device)
    background_feature = torch.cat([bg_mask, goal_feature_bg], axis=-1)
                        
    data_pos = torch.cat([object_pcd_points, background_pcd_points], axis=1)
    data_pos = data_pos.reshape(-1, 3)
    data_x = torch.cat([object_feature, background_feature], axis=1)
    data_x = data_x.reshape(-1, data_x.shape[-1])
    data_batch = torch.arange(object_pcd_points.shape[0], device=device).repeat_interleave(object_pcd_points.shape[1] + background_pcd_points.shape[1])
    return data_x, data_pos, data_batch

hacman/algos/feature_extractors/feature_extractors.py
HACMan-2023/HACMan
import numpy as np
import torch as th
from torch.nn import functional as F
from stable_baselines3.td3.td3 import TD3
from stable_baselines3.common.preprocessing import preprocess_obs as pobs
from typing import Any, Dict, List, Optional, Type, Union

import gym
import torch as th
from torch import nn

from stable_baselines3.common.policies import ContinuousCritic
from stable_baselines3.common.torch_layers import (
    BaseFeaturesExtractor,
    FlattenExtractor,
)
from stable_baselines3.common.type_aliases import Schedule
from stable_baselines3.td3.policies import TD3Policy, Actor


class HACManTD3(TD3):
    def __init__(self, *args, temperature=0.01, **kwargs):
        super().__init__(*args, **kwargs)
        self.temperature = temperature
        
    def get_next_q_values(self, replay_data):
        obs_tensor = replay_data.next_observations
        batch_size = obs_tensor['object_pcd_points'].shape[0]
        object_pcd_size = obs_tensor['object_pcd_points'].shape[1]
        preprocessed_obs = pobs(obs_tensor, self.actor.observation_space, normalize_images=self.actor.normalize_images)
        
        # Get per point action only for object points
        features = self.actor_target.features_extractor(preprocessed_obs, per_point_output=True)
        features = features.reshape(batch_size, -1, features.shape[-1])[:, :object_pcd_size, :]
        actions = self.actor_target.mu(features)

        # Add noise        
        noise = actions.clone().data.normal_(0, self.target_policy_noise)
        noise = noise.clamp(-self.target_noise_clip, self.target_noise_clip)
        actions = (actions + noise).clamp(-1, 1)
        
        # Get per point feature
        features = self.critic_target.features_extractor(preprocessed_obs, per_point_output=True)
        features = features.reshape(batch_size, -1, features.shape[-1])[:, :object_pcd_size, :]
        # Concatenate with per point critic feature and get per point value
        out1 = self.critic_target.q_networks[0](th.cat([features, actions], dim=2)).reshape(batch_size, -1, 1)
        out2 = self.critic_target.q_networks[1](th.cat([features, actions], dim=2)).reshape(batch_size, -1, 1)
        next_q_values = th.cat([out1, out2], dim=2)
        
        if self.mean_q:
            next_q_values = th.mean(next_q_values, dim=2, keepdim=True)
        else:
            next_q_values, _ = th.min(next_q_values, dim=2, keepdim=True)
        
        # object_pcd_points dim: batch x pcd_size x 3
        action_score = th.nn.Softmax(dim=1)(next_q_values/self.temperature)
        next_q_values = th.sum(next_q_values*action_score, dim=1)
            
        return next_q_values
    
    def get_actor_loss(self, replay_data):
        obs_tensor = replay_data.observations
        batch_size = obs_tensor['object_pcd_points'].shape[0]
        object_pcd_size = obs_tensor['object_pcd_points'].shape[1]
        preprocessed_obs = pobs(obs_tensor, self.actor.observation_space, normalize_images=self.actor.normalize_images)
        
        # Get per point action only for object points
        features = self.actor.features_extractor(preprocessed_obs, per_point_output=True)
        features = features.reshape(batch_size, -1, features.shape[-1])[:, :object_pcd_size, :]
        actions = self.actor.mu(features)
        
        # Get per point feature
        features = self.critic.features_extractor(preprocessed_obs, per_point_output=True)
        features = features.reshape(batch_size, -1, features.shape[-1])[:, :object_pcd_size, :]
        # Concatenate with per point critic feature and get per point value
        qvalue = self.critic.q_networks[0](th.cat([features, actions], dim=2)).reshape(batch_size, -1, 1)
        action_score = th.nn.Softmax(dim=1)(qvalue/self.temperature)
        
        actor_loss = -th.sum(qvalue*action_score, dim=1).mean()
        return actor_loss


hacman/algos/hacman_td3.py
HACMan-2023/HACMan
import numpy as np
import torch
from stable_baselines3.common.preprocessing import preprocess_obs as pobs


class LocationPolicy(object):
    def __init__(self) -> None:
        pass
        
    def load_model(self, model):
        pass
    
    def get_action(self, obs):
        pass


class LocationPolicyWithArgmaxQ(LocationPolicy):
    def __init__(self, model=None, temperature=1., deterministic=False, vis_only=False, egreedy=0.) -> None:
        self.model = model
        self.temperature = temperature
        self.egreedy = egreedy
        self.deterministic = deterministic
        self.vis_only = vis_only # only use the scores for visualization
        return
    
    def load_model(self, model):
        self.model = model
        return
    
    def get_action(self, obs):  
        assert self.model, "LocationPolicyWithArgmaxQ: needs to load a model."      
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        obs_tensor = {}
        for key in self.model.actor.observation_space.spaces.keys():
            try:
                if key != 'action_location_score':
                    obs_tensor[key] = torch.tensor(obs[key]).to(device).unsqueeze(dim=0)
            except:
                print(f'obs_tensor[key] = torch.tensor(obs[key]).to(device).unsqueeze(dim=0)')
                print(f'key:{key}, value:{obs[key]}')
        
        # Need to use eval. Otherwise this will affect batch norm.
        was_training = self.model.policy.training
        self.model.policy.set_training_mode(False)
        with torch.no_grad():
            preprocessed_obs = pobs(obs_tensor, self.model.actor.observation_space, normalize_images=self.model.actor.normalize_images)
            
            if self.model.actor.features_extractor.__class__.__name__ == 'PointCloudGlobalExtractor':
                # Get action
                actions = self.model.actor(preprocessed_obs)
                # Get per point feature
                features = self.model.critic.features_extractor(preprocessed_obs, per_point_output=True)
                # Concatenate with per point critic feature and get per point value
                actions = torch.cat([actions]*features.shape[0], dim=0)
                out = self.model.critic.q_networks[0](torch.cat([features, actions], dim=1))
            else:
                # Get per point action
                features = self.model.actor.features_extractor(preprocessed_obs, per_point_output=True)
                actions = self.model.actor.mu(features)
                # Get per point feature
                features = self.model.critic.features_extractor(preprocessed_obs, per_point_output=True)
                # Concatenate with per point critic feature and get per point value
                out = self.model.critic.q_networks[0](torch.cat([features, actions], dim=1))
        self.model.policy.set_training_mode(was_training)
        
        action_score = out[: obs['object_pcd_points'].shape[0]].reshape(-1)
        action_score = torch.nn.Softmax(dim=0)(action_score/self.temperature).cpu().detach().numpy()
        
        if self.vis_only:
            poke_idx = np.random.randint(len(obs['object_pcd_points']))
        elif self.deterministic:
            poke_idx = np.argmax(action_score)
        else:
            # Uniformly random before learning_starts 
            if self.model.num_timesteps < self.model.learning_starts or np.random.rand() < self.egreedy:
                poke_idx = np.random.choice(np.arange(len(action_score)))
            else:
                poke_idx = np.random.choice(np.arange(len(action_score)), p=action_score)
        
        poke_idx = np.expand_dims(poke_idx, axis=0)

        action = {'action_location_score':action_score,
                  'poke_idx': poke_idx}
        return action

class RandomLocation(LocationPolicy):
    def __init__(self) -> None:
        return
    
    def get_action(self, obs):
        # Choose an index from the observation point cloud
        points = obs['object_pcd_points']
        if points.ndim == 2:
            poke_idx = np.random.randint(len(points), size=(1,))
            return {'poke_idx': poke_idx,
                    'action_location_score': np.zeros(len(points))}
        elif points.ndim == 3: # batched
            bs, n_points, _ = points.shape
            poke_idx = np.random.randint(n_points, size=(bs, 1))
            return {'poke_idx': poke_idx,
                    'action_location_score': np.zeros((bs, n_points))}

hacman/algos/location_policy.py
HACMan-2023/HACMan
import numpy as np
import torch as th
from torch.nn import functional as F
from stable_baselines3.td3.td3 import TD3
from stable_baselines3.common.preprocessing import preprocess_obs as pobs
from typing import Any, Dict, List, Optional, Type, Union

import gym
import torch as th
from torch import nn

from stable_baselines3.common.policies import ContinuousCritic
from stable_baselines3.common.torch_layers import (
    BaseFeaturesExtractor,
    FlattenExtractor,
)
from stable_baselines3.common.type_aliases import Schedule
from stable_baselines3.td3.policies import TD3Policy, Actor


class MixTD3Policy(TD3Policy):

    def __init__(
        self,
        observation_space: gym.spaces.Space,
        action_space: gym.spaces.Space,
        lr_schedule: Schedule,
        net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
        activation_fn: Type[nn.Module] = nn.ReLU,
        actor_features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
        critic_features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
        actor_features_extractor_kwargs: Optional[Dict[str, Any]] = None,
        critic_features_extractor_kwargs: Optional[Dict[str, Any]] = None,
        normalize_images: bool = True,
        optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
        optimizer_kwargs: Optional[Dict[str, Any]] = None,
        n_critics: int = 2,
        share_features_extractor: bool = False,
    ):
        self.actor_features_extractor_class = actor_features_extractor_class
        self.critic_features_extractor_class = critic_features_extractor_class
        if actor_features_extractor_kwargs is None:
            actor_features_extractor_kwargs = {}
        if critic_features_extractor_kwargs is None:
            critic_features_extractor_kwargs = {}
        self.actor_features_extractor_kwargs = actor_features_extractor_kwargs
        self.critic_features_extractor_kwargs = critic_features_extractor_kwargs
        
        super().__init__(
            observation_space,
            action_space,
            lr_schedule=lr_schedule,
            net_arch=net_arch,
            activation_fn=activation_fn,
            features_extractor_class=None,
            features_extractor_kwargs=None,
            normalize_images=normalize_images,
            optimizer_class=optimizer_class,
            optimizer_kwargs=optimizer_kwargs,
            n_critics=n_critics,
            share_features_extractor=share_features_extractor,            
        )

    def make_actor(self, features_extractor: Optional[BaseFeaturesExtractor] = None) -> Actor:
        actor_kwargs = self.actor_kwargs.copy()
        features_extractor = self.actor_features_extractor_class(self.observation_space, **self.actor_features_extractor_kwargs)
        actor_kwargs.update(dict(features_extractor=features_extractor, features_dim=features_extractor.features_dim))    
        return Actor(**actor_kwargs).to(self.device)

    def make_critic(self, features_extractor: Optional[BaseFeaturesExtractor] = None) -> ContinuousCritic:        
        critic_kwargs = self.critic_kwargs.copy()
        features_extractor = self.critic_features_extractor_class(self.observation_space, **self.critic_features_extractor_kwargs)
        critic_kwargs.update(dict(features_extractor=features_extractor, features_dim=features_extractor.features_dim))    
        return ContinuousCritic(**critic_kwargs).to(self.device)

hacman/algos/mix_td3.py
HACMan-2023/HACMan
import os
import functools

from stable_baselines3.td3.policies import TD3Policy
from stable_baselines3 import TD3

from hacman.algos.hacman_td3 import HACManTD3
from hacman.algos.mix_td3 import MixTD3Policy
from hacman.algos.feature_extractors.feature_extractors import PointCloudExtractor, PointCloudGlobalExtractor, StatesExtractor


def add_model_config(parser):
    parser.add_argument("--load_ckpt", default=None, type=str, help="Ckpt path. Set to \"latest\" to use the latest checkpoint.")
    parser.add_argument("--algo", default='TD3', type=str, help="RL algorithm")
    parser.add_argument("--gradient_steps", default=40, type=int, help="Gradient step per env step")
    parser.add_argument("--batch_size", default=64, type=int, help="Batch Size")
    parser.add_argument("--ent_coef", default='auto', type=str, help="Entropy Coefficient mode")
    parser.add_argument("--clip_grad_norm", default=None, type=float, help="Clip gradient norm for critic")
    parser.add_argument("--clamp_critic_max", default=None, type=float, help="Clamp critic value")
    parser.add_argument("--clamp_critic_min", default=None, type=float, help="Clamp critic value")
    parser.add_argument("--gamma", default=0.99, type=float, help="Discount factor")
    parser.add_argument("--actor_update_interval", default=4, type=int, help="Actor update interval")
    parser.add_argument("--target_update_interval", default=4, type=int, help="Target update interval")
    parser.add_argument("--initial_timesteps", default=10000, type=int, help="Initial env steps before training starts")
    parser.add_argument("--eval_freq", default=100, type=int, help="Eval per n env steps")
    parser.add_argument("--mean_q", action="store_true", help="Use mean Q instead of min Q")
    parser.add_argument("--learning_rate", default=1e-4, type=float, help="Learning rate")
    parser.add_argument("--share_features_extractor", action="store_true", help="share features extractor")
    parser.add_argument("--preprocessing_fn", default="flow", type=str, help="Input processing function")
    parser.add_argument('--net_arch', default=[128, 128, 128], nargs="+", type=int)

    # Baseline options
    parser.add_argument("--action_mode", default='per_point_action', type=str,
                        choices={'per_point_action', "regress_location_and_action", 'regress_action_only'},
                        help="Set action modes for different baselines.")
    parser.add_argument("--feature_mode", default='points', type=str,
                        choices={'points', 'states'})
    return


def setup_model(config, env, result_folder, normalize_param=None):
    # ----------- Feature extractors ------------ 
    if config["feature_mode"] == "states":
        if config["action_mode"] == "regress_location_and_action":
            features_extractor_class = StatesExtractor
        elif config["action_mode"] == "regress_action_only":
            features_extractor_class = functools.partial(StatesExtractor, include_gripper=True)
        else:
            raise ValueError
        
    if config["feature_mode"] == "points":
        if config["action_mode"] == "per_point_action":
            features_extractor_class = functools.partial(PointCloudExtractor, preprocessing_fn=config['preprocessing_fn'])
        elif config["action_mode"] == "regress_location_and_action":
            features_extractor_class = PointCloudGlobalExtractor
        elif config["action_mode"] == "regress_action_only":
            features_extractor_class = functools.partial(PointCloudGlobalExtractor, include_gripper=True)
        else:
            raise ValueError
        
    if config['feature_mode'] == 'points' and normalize_param is not None:
        features_extractor_class =  functools.partial(features_extractor_class, normalize_pos_param=normalize_param)

    # ----------- Model ------------ 
    if config['algo'] == 'TD3':
        policy = functools.partial(TD3Policy,
                        features_extractor_class=features_extractor_class,
                        share_features_extractor=config["share_features_extractor"],
                        net_arch=config['net_arch'])
        td3_kwargs = dict(batch_size=config['batch_size'], gamma=config['gamma'],
                verbose=1, learning_starts=config['initial_timesteps'], learning_rate=config['learning_rate'],
                actor_update_interval=config['actor_update_interval'], 
                target_update_interval=config['target_update_interval'], train_freq=1,
                tensorboard_log=result_folder+'/tb', gradient_steps=config['gradient_steps'], 
                clip_critic_grad_norm=config['clip_grad_norm'], clamp_critic_min=config['clamp_critic_min'],
                clamp_critic_max=config['clamp_critic_max'], seed=config["seed"],
                mean_q=config['mean_q'])
        if config['load_ckpt'] is None:
            model = TD3(policy, env, **td3_kwargs)
        else:
            model = TD3.load(path=config['load_ckpt'], env=env, **td3_kwargs) # Overwrite hyperparameters of the loaded model
            print(f"Loaded policy: {config['load_ckpt']}")
    elif config['algo'] == 'HybridTD3':
        policy = functools.partial(TD3Policy,
                        features_extractor_class=features_extractor_class,
                        share_features_extractor=False,
                        net_arch=config['net_arch'])
        hybrid_td3_kwargs = dict(
            batch_size=config['batch_size'], gamma=config['gamma'],
            verbose=1, learning_starts=config['initial_timesteps'], learning_rate=config['learning_rate'],
            actor_update_interval=config['actor_update_interval'],
            target_update_interval=config['target_update_interval'], train_freq=1,
            tensorboard_log=result_folder+'/tb', gradient_steps=config['gradient_steps'],
            clip_critic_grad_norm=config['clip_grad_norm'], clamp_critic_min=config['clamp_critic_min'],
            clamp_critic_max=config['clamp_critic_max'], seed=config["seed"],
            mean_q=config['mean_q'], temperature=config['location_model_temperature']
        )
        if config['load_ckpt'] is None:
            model = HACManTD3(policy, env, **hybrid_td3_kwargs)
        else:
            model = HACManTD3.load(path=config['load_ckpt'], env=env, **hybrid_td3_kwargs) # Overwrite hyperparameters of the loaded model
            print(f"Loaded policy: {config['load_ckpt']}")
        
    elif config['algo'] == 'TD3MixArch':
        policy = functools.partial(MixTD3Policy,
                        actor_features_extractor_class=PointCloudGlobalExtractor,
                        critic_features_extractor_class=PointCloudExtractor,
                        net_arch=config['net_arch'])
        model = TD3(policy, env, batch_size=config['batch_size'], gamma=config['gamma'],
                verbose=1, learning_starts=config['initial_timesteps'], learning_rate=config['learning_rate'],
                actor_update_interval=config['actor_update_interval'],
                target_update_interval=config['target_update_interval'], train_freq=1,
                tensorboard_log=result_folder+'/tb', gradient_steps=config['gradient_steps'],
                clip_critic_grad_norm=config['clip_grad_norm'], clamp_critic_min=config['clamp_critic_min'],
                clamp_critic_max=config['clamp_critic_max'], seed=config["seed"],
                mean_q=config['mean_q'])

    return model
hacman/algos/setup_model.py
HACMan-2023/HACMan
import os, sys

import numpy as np
import gym

from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv, _flatten_obs
from hacman.utils.plotly_utils import plot_pcd, plot_action, plot_pcd_with_score
from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, Union, Dict
from hacman.algos.location_policy import RandomLocation


class SubprocVecEnvwithLocationPolicy(SubprocVecEnv):
    def __init__(self, env_fns: List[Callable[[], gym.Env]], 
                 start_method: Optional[str] = None, 
                 location_model: Callable[[Dict], Dict] = None):
        
        self.location_model = location_model
        super().__init__(env_fns, start_method)
    
    def process_obs(self, obs):
        count = 0
        for sub_obs in obs:
            location_info = self.location_model.get_action(sub_obs)
            sub_obs.update(location_info)
            assert self.env_method('set_prev_obs', sub_obs, indices=count)
            count += 1
        return
    
    def step_wait(self) -> VecEnvStepReturn:
        results = [remote.recv() for remote in self.remotes]
        self.waiting = False
        obs, rews, dones, infos = zip(*results)
        self.process_obs(obs) # New
        return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos
    
    def reset(self) -> VecEnvObs:
        for remote in self.remotes:
            remote.send(("reset", None))
        obs = [remote.recv() for remote in self.remotes]
        self.process_obs(obs) # New
        return _flatten_obs(obs, self.observation_space)


class DummyVecEnvwithLocationPolicy(DummyVecEnv):
    def __init__(self, env_fns: List[Callable[[], gym.Env]],
                 location_model: Callable[[Dict], Dict] = None):
        
        self.location_model = location_model
        super().__init__(env_fns)
    
    def _save_obs(self, env_idx: int, obs: VecEnvObs) -> None:
        location_info = self.location_model.get_action(obs)
        obs.update(location_info)
        assert self.env_method('set_prev_obs', obs, indices=env_idx)
        return super()._save_obs(env_idx, obs)

hacman/envs/location_policy_wrappers.py
HACMan-2023/HACMan
import numpy as np
import functools

from gym.wrappers.time_limit import TimeLimit

from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv

from hacman.envs.location_policy_wrappers import DummyVecEnvwithLocationPolicy, SubprocVecEnvwithLocationPolicy
from hacman.envs.setup_location_policy import setup_location_policy, add_location_policy_config


def add_env_config(parser):
    parser.add_argument("--env", default='simple_env', type=str, choices={'simple_env', 'hacman_bin_env'})
    parser.add_argument("--train_n_envs", default=20, type=int, help="Number of training envs in parallel")
    parser.add_argument("--eval_n_envs", default=10, type=int, help="Number of eval envs in parallel")
    parser.add_argument("--record_from_cam", default=None, type=str, help="Record video from camera.")
    parser.add_argument("--max_episode_steps", default=10, type=int, help="Max episode steps")

    # Location policy
    add_location_policy_config(parser)
    
    '''
    HACMan Bin Env Specific Configs
    '''
    # Object mesh
    parser.add_argument("--object_dataset", default='housekeep_all', type=str, choices={'cube', 'housekeep', 'housekeep_all'})
    parser.add_argument("--object_types", default=None, nargs='*', type=str, help="Load from specific categories")
    parser.add_argument("--object_name", default=None, type=str, help="Load from a specific object name")
    parser.add_argument("--object_split_train", default='train_thick', type=str, help="Load from a specific split")
    parser.add_argument("--object_split_eval", default='train_thick', type=str, help="Load from a specific split")
    parser.add_argument("--object_scale_range", default=[0.8, 1.2], nargs=2, type=float, help="Scale range for object size randomization")
    parser.add_argument("--object_size_limit", default=[0.05, 0.05], nargs=2, type=float, help="Max/min object size")

    parser.add_argument("--object_voxel_down_sample", default=0.005, type=float, help="Downsample object point cloud voxel size")
    parser.add_argument("--pcd_mode", default='reduced1', type=str, help="Point cloud sampling mode")
    parser.add_argument("--pcd_noise", default=None, nargs=4, type=float, help="Pcd noise")

    # Task
    parser.add_argument("--object_init_pose", default="any", type=str, help="Init pose sampling mode")
    parser.add_argument("--goal_mode", default='any_var_size', type=str, help="Goal sampling mode")
    parser.add_argument("--goal_mode_eval", default=None, type=str, help="Goal sampling mode")
    parser.add_argument("--success_threshold", default=0.03, type=float, help="Success threshold")
    
    # Bin
    parser.add_argument("--table_full_size", default=(0.45, 0.54, 0.107), nargs=3, type=float, help="Bin size")
    parser.add_argument("--table_size_rand", default=(0., 0., 0.), nargs=3, type=float, help="Bin size")
    parser.add_argument("--gripper_types", default='default', type=str, help="Gripper type")
    parser.add_argument("--friction_config", default='default', type=str, help="friction settings")

    # Actions
    parser.add_argument("--action_repeat", default=3, type=int, help="Number of action repeats")
    parser.add_argument("--location_noise", default=0., type=float, help="Location noise")
    parser.add_argument("--reward_scale", default=1., type=float, help="Reward scale")
    parser.add_argument("--fixed_ep_len", action="store_true", help="Fixed episode length")
    parser.add_argument("--reward_gripper_distance", default=None, type=float, help="Whether to use distance in the reward calculation")

    return

def setup_envs(config, eval_env_only=False): 
    train_env, eval_env = None, None
    # ------------- Train/eval wrapper setup ------------- 
    if config["action_mode"] == "per_point_action":
        location_model_train, location_model_eval = setup_location_policy(config)
        if config['debug']:
            train_wrapper = functools.partial(DummyVecEnvwithLocationPolicy, location_model=location_model_train)
            eval_wrapper = functools.partial(DummyVecEnvwithLocationPolicy, location_model=location_model_eval)
        else:
            train_wrapper = functools.partial(SubprocVecEnvwithLocationPolicy, location_model=location_model_train)
            eval_wrapper = functools.partial(SubprocVecEnvwithLocationPolicy, location_model=location_model_eval)
    else:
        if config['debug']:
            train_wrapper, eval_wrapper = DummyVecEnv, DummyVecEnv
        else:
            train_wrapper, eval_wrapper = SubprocVecEnv, SubprocVecEnv
    
    # ------------- Env Setup ------------- 
    env_configs = dict(env=config['env'],
                        reward_scale=config['reward_scale'],
                        reward_gripper_distance=config['reward_gripper_distance'],
                        goal_mode=config['goal_mode'],
                        pcd_mode=config['pcd_mode'],
                        action_mode=config['action_mode'],
                        action_repeat=config['action_repeat'],
                        object_dataset=config['object_dataset'],
                        object_types=config['object_types'],
                        object_name=config['object_name'],
                        object_scale_range=config['object_scale_range'],
                        object_init_pose=config['object_init_pose'],
                        object_split=config['object_split_train'],
                        object_size_limit=config['object_size_limit'],
                        fixed_ep_len=config['fixed_ep_len'],
                        table_full_size=config['table_full_size'],
                        table_size_rand=config['table_size_rand'],
                        location_noise=config['location_noise'],
                        pcd_noise=config['pcd_noise'],
                        object_voxel_down_sample=config['object_voxel_down_sample'],
                        friction_config=config['friction_config'],
                        success_threshold=config['success_threshold'],
                        gripper_types=config['gripper_types'],
                        max_episode_steps=config['max_episode_steps'],
                        )

    def make_env(config):
        if config['env'] == 'hacman_bin_env':
            from hacman.envs.sim_envs.hacman_bin_env import HACManBinEnv
            env = HACManBinEnv(**config)
        elif config['env'] == 'simple_env':
            from hacman.envs.sim_envs.simple_env import SimpleEnv
            env = SimpleEnv(**config)
        else:
            raise NotImplementedError
        env = TimeLimit(env, max_episode_steps=config['max_episode_steps'])
        return env

    # Train Env
    if not eval_env_only:
        make_train_env = functools.partial(make_env, env_configs)
        train_env = make_vec_env(make_train_env, n_envs=config['train_n_envs'], seed=config["seed"], 
                        vec_env_cls=train_wrapper)
    
    # Eval Env
    eval_env_config = env_configs.copy()
    if config['goal_mode_eval'] is not None:
        eval_env_config.update(goal_mode=config['goal_mode_eval'])
    eval_env_config.update(record_from_cam=config['record_from_cam'])
    eval_env_config.update(object_split=config['object_split_eval'])
    make_eval_env = functools.partial(make_env, eval_env_config)
    eval_env = make_vec_env(make_eval_env, n_envs=config['eval_n_envs'], seed=config["seed"], 
                            vec_env_cls=eval_wrapper)
    
    if eval_env_only:
        return eval_env
    
    return train_env, eval_env

hacman/envs/setup_envs.py
HACMan-2023/HACMan
from hacman.algos.location_policy import LocationPolicy, LocationPolicyWithArgmaxQ, RandomLocation
from hacman.envs.location_policy_wrappers import SubprocVecEnvwithLocationPolicy, DummyVecEnvwithLocationPolicy


def add_location_policy_config(parser):
    parser.add_argument("--location_model", default='argmaxQ', type=str, help="location model options")
    parser.add_argument("--location_model_temperature", default=0.1, type=float, help="location model options")
    parser.add_argument("--egreedy", default=0.1, type=float, help="Epsilon greedy for location")    
    return

def setup_location_policy(config):
    # Return the location policy for train_env and eval_env
    if config['location_model'] == 'random':
        location_model_train = RandomLocation()
        location_model_eval = location_model_train
    elif config['location_model'] == 'argmaxQ_vis':
        location_model_train = RandomLocation()
        location_model_eval = LocationPolicyWithArgmaxQ(model=None, vis_only=True) # Load model later
        return location_model_train, location_model_eval
    elif config['location_model'] == 'argmaxQ':
        location_model_train = LocationPolicyWithArgmaxQ(model=None, 
                                                        temperature=config['location_model_temperature'],
                                                        egreedy=config['egreedy'])
        location_model_eval = LocationPolicyWithArgmaxQ(model=None, 
                                                        temperature=config['location_model_temperature'],
                                                        deterministic=True)
    else:
        raise NotImplementedError
    
    return location_model_train, location_model_eval
hacman/envs/setup_location_policy.py
HACMan-2023/HACMan
import numpy as np
import gym

from hacman.utils.transformations import transform_point_cloud, sample_idx

class RandomLocation(object):
    def __init__(self) -> None:
        return
    
    def get_action(self, obs):
        # Choose an index from the observation point cloud
        poke_idx = np.random.randint(len(obs['object_pcd_points']))
        return {'poke_idx': poke_idx,
                'action_location_score': np.zeros(len(obs['object_pcd_points']))}

"""
Abstracted Gym wrapper as an interface for HACMAN
"""
class BaseEnv(gym.Env):
    def __init__(self,
            object_pcd_size=400,
            background_pcd_size=400, **kwargs):
        
        self.spec = None
        self.metadata = {}

        # Observation space
        self.object_pcd_size = object_pcd_size
        self.background_pcd_size = background_pcd_size
        self.observation_space = gym.spaces.Dict(
                spaces={
                    "poke_idx": gym.spaces.Box(-np.inf, np.inf, (1,)),
                    "object_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "goal_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    # "gripper_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "object_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.object_pcd_size, 3)),
                    "action_location_score": gym.spaces.Box(-np.inf, np.inf, (self.object_pcd_size,)),
                    "background_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.background_pcd_size, 3)),
                }
            )
        
        # Action space
        self.action_space = gym.spaces.Box(-1, 1, (3,))

        # Previous observation
        self.prev_obs = None

        # Action location policy
        self.random_location_policy = RandomLocation()

        super().__init__()
    

    def reset(self):
        obs = self.get_observations()
        self.prev_obs = obs
        return obs
    
    def set_prev_obs(self, obs):
        self.prev_obs = obs
        return True
        
    def step(self, action):
        # Take the environment step
        # Use the previous observation to get the location
        points = self.prev_obs['object_pcd_points']
        idx = self.prev_obs['poke_idx']
        location = points[idx]

        # TODO: Given the robot-object contact location "location" and the 
        # continuous paramter "action", Execute the action in simulation.
        pass
        
        # Compute the step outcome
        obs = self.get_observations()
        self.prev_obs = obs

        success, reward = self._evaluate_goal(obs, obs['goal_pose'])

        info = {"is_success": success,
                "action_param": action,
                "action_location": location,}

        done = False

        return obs, reward, done, info
        

    def get_observations(self):
        obs = {}
        # TODO: Change the follow placeholders
        obs['object_pose'] = np.eye(4)
        obs['goal_pose'] = np.eye(4)
        obs['object_pcd_points'] = np.zeros((self.object_pcd_size, 3))
        obs['background_pcd_points'] = np.zeros((self.object_pcd_size, 3))
    
        obj_idx = np.arange(obs['object_pcd_points'].shape[0])
        points = obs['object_pcd_points'][obj_idx]
        
        # Sample points to a fixed length
        obj_idx = sample_idx(len(points), self.object_pcd_size)
        obs['object_pcd_points'] = points[obj_idx, :]
        bg_idx = sample_idx(len(obs['background_pcd_points']), self.background_pcd_size)
        obs['background_pcd_points'] = obs['background_pcd_points'][bg_idx, :]
        
        # Choose location here:
        # Use a random location by default. Will be overwritten by the wrapper.
        location_info = self.random_location_policy.get_action(obs)
        obs.update(location_info)
        
        return obs
    
    def _evaluate_goal(self, obs, goal):
        Warning("This is a placeholder function. Please implement your own.")
        # Compute the flow
        old_pcd = obs['object_pcd_points']
        new_pcd = transform_point_cloud(obs['object_pose'], goal, old_pcd)
        flow = np.linalg.norm(new_pcd - old_pcd, axis=1)
        mean_flow = np.mean(flow)

        reward = -mean_flow
        success = mean_flow < 0.03

        return success, reward
    
if __name__ == "__main__":
    env = BaseEnv()
    env.reset()
    for _ in range(10):
        obs, reward, done, info = env.step(np.array([0.1, 0.2, 0.3]))
hacman/envs/sim_envs/base_env.py
HACMan-2023/HACMan
"""
Should different action modes be included here?
"""

import os, sys

import numpy as np
import imageio
import pickle
import gym

from .base_env import BaseEnv, sample_idx
from hacman.utils.transformations import to_pose_mat, transform_point_cloud, decompose_pose_mat

from hacman.utils.plotly_utils import plot_pcd, plot_action, plot_pcd_with_score
from bin_env.util import angle_diff
from bin_env.poke_env import PokeEnv
# import plotly.graph_objects as go
from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, Union, Dict
from hacman.algos.location_policy import RandomLocation


class HACManBinEnv(BaseEnv):
    def __init__(self, 
                 action_mode="per_point_action",
                 reward_scale=1.,
                 reward_mode="flow",
                 real_robot=False,
                 action_repeat=10,
                 fixed_ep_len=False,
                 reward_gripper_distance=None,
                 success_threshold=0.03,
                 **kwargs):
        super().__init__(
            object_pcd_size=400,
            background_pcd_size=400)
        
        self.real_robot = real_robot
        if self.real_robot:
            from franka_env_polymetis import FrankaPokeEnv
            self.env = FrankaPokeEnv(**kwargs)
        else:
            if action_mode == 'regress_action_only':
                # Only use position limits when the action continues from previous position
                kwargs['ignore_position_limits'] = False
            self.env = PokeEnv(**kwargs)
        
        self.action_repeat = action_repeat
        self.fixed_ep_len = fixed_ep_len
        self.reward_mode = reward_mode
        self.reward_scale = reward_scale
        self.reward_gripper_distance = reward_gripper_distance
        self.success_threshold = success_threshold

        # Override the base env's observation and action space
        self.action_mode = action_mode
        if self.action_mode == "per_point_action":
            self.observation_space = gym.spaces.Dict(
                spaces={
                    "poke_idx": gym.spaces.Box(-np.inf, np.inf, (1,)),
                    "object_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "goal_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "gripper_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "object_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.object_pcd_size, 3)),
                    "action_location_score": gym.spaces.Box(-np.inf, np.inf, (self.object_pcd_size,)),
                    "background_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.background_pcd_size, 3)),
                }
            )
            self.action_space = gym.spaces.Box(-1, 1, (3,))
            
        elif self.action_mode == "regress_location_and_action":
            self.observation_space = gym.spaces.Dict(
                spaces={
                    "object_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "goal_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "gripper_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "object_size": gym.spaces.Box(-np.inf, np.inf, (3,)),
                    "object_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.object_pcd_size, 3)),
                    "background_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.background_pcd_size, 3)),
                }
            )
            self.action_space = gym.spaces.Box(-1, 1, (6,))
        
        elif self.action_mode == "regress_action_only":
            self.observation_space = gym.spaces.Dict(
                spaces={
                    "object_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "goal_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "gripper_pose": gym.spaces.Box(-np.inf, np.inf, (4,4)),
                    "object_size": gym.spaces.Box(-np.inf, np.inf, (3,)),
                    "object_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.object_pcd_size, 3)),
                    "background_pcd_points": gym.spaces.Box(-np.inf, np.inf, (self.background_pcd_size, 3)),
                }
            )
            self.action_space = gym.spaces.Box(-1, 1, (3,))
        
        else:
            raise ValueError(f"Unknown action mode {self.action_mode}.")
        
                       
        self.goal = None
        
        self.step_count = 0
        self.missing_object_count = 0
        self.motion_failure_count = 0

    
    
    def reset(self, **kwargs):
        if self.real_robot:
            for i in range(10):
                raw_obs = self.env.reset(**kwargs)
                obs = self.process_observation(raw_obs)
                success, reward = self._evaluate_goal(obs, obs['goal_pose'])
                if success:
                    print(f'Resample goal because t=0 is at the goal: Attempt {i}')
                else:
                    break
            if success:
                print(f'Cannot find a goal within 10 attempts.')
                import ipdb; ipdb.set_trace()
        else:
            if self.action_mode == "regress_action_only":
                raw_obs = self.env.reset(
                    **kwargs, start_gripper_above_obj=True)
            else:
                raw_obs = self.env.reset(**kwargs)
                
            self.goal = self.env.sample_goal(raw_obs)
            self.env.set_goal(self.goal)
            obs = self.process_observation(raw_obs)
                
        self.prev_obs = obs
        return obs
    
    def step(self, action):
        start_location = None
        if self.action_mode == "per_point_action":
            # Per point regress action
            points = self.prev_obs['object_pcd_points']
            normals = self.prev_obs['object_pcd_normals']
            idx = self.prev_obs['poke_idx']
            location = points[idx]
            normal = normals[idx]
            
        elif self.action_mode == "regress_location_and_action":
            # Global regress location and action
            assert len(action) == 6
            
            # Raw action is in range [-1, 1]. We need to map the location to
            # the range of the bin.
            location = action[3:6]
            
            obj_pos, obj_ori = decompose_pose_mat(self.prev_obs['object_pose'], wxyz=True)
            
            cube_size = np.max(self.env.get_cube_size())
            bbox_half_size = cube_size / 2.0 + 0.02

            bbox_center = np.copy(obj_pos)
            bbox_center[2] = bbox_half_size + self.env.table_offset[2]

            location = location * bbox_half_size + bbox_center
            
            start_location = np.copy(location)
            start_location[2] = 2 * bbox_half_size + self.env.table_offset[2]
            # location += offset
            
            action = action[0:3]
            action[2] = - abs(action[2])
            normal = None
            
        elif self.action_mode == "regress_action_only":
            # Global regress action (continue from previous location)
            assert len(action) == 3
            action[2] = - abs(action[2])
            location = None
            normal = None
        
        else:
            raise ValueError(f"Unknown action mode {self.action_mode}.")
        
        raw_obs, _, _, poke_info = self.env.poke(
            location, normal, action, action_repeat=self.action_repeat,
            start_location=start_location)

        # Calculate step outcome
        self.step_count += 1
        if poke_info['poke_success'] == False:
            if poke_info['box_in_the_view'] == False:
                self.missing_object_count += 1
                print(f"{self.missing_object_count/self.step_count*100:.2f}% missing object within "
                    f"{self.step_count} total steps")
            else:
                self.motion_failure_count += 1
                if self.motion_failure_count % 100 == 0:
                    print(f"{self.motion_failure_count/self.step_count*100:.2f}% motion failure within "
                        f"{self.step_count} total steps")
        obs = self.process_observation(raw_obs)
        self.prev_obs = obs
            
        success, reward = self._evaluate_goal(obs, obs['goal_pose'])
        
        info = {"is_success": success,
                "action_param": action,
                "action_location": location,
                "poke_success":poke_info['poke_success'],
                "box_in_the_view": poke_info['box_in_the_view'],
                "object_name": poke_info['object_name'],
                }
                
        if "cam_frames" in poke_info.keys():
            info["cam_frames"] = poke_info["cam_frames"]
            
        if "fitness" in poke_info.keys():
            info["fitness"] = poke_info["fitness"]
                
        done = False
        if not self.fixed_ep_len:
            done = success
            
        return obs, reward, done, info
    
    def process_observation(self, obs):
        # Process the raw observation from robosuit
        new_obs = {}
        new_obs['object_pose'] = to_pose_mat(obs['cube_pos'], obs['cube_quat'])
        if 'goal_pose' in obs.keys():
            new_obs['goal_pose'] = obs['goal_pose']
        else:
            new_obs['goal_pose'] = self.goal.copy()

        if 'object_size' in obs.keys():
            new_obs['object_size'] = obs['object_size']

        if 'object_size' in obs.keys():
            new_obs['object_size'] = obs['object_size']

        # Add the gripper pose
        ee_pos, ee_quat = obs['robot0_eef_pos'], obs['robot0_eef_quat']
        new_obs['gripper_pose'] = to_pose_mat(ee_pos, ee_quat)

        if not 'object_pcd_points' in obs.keys() or obs['object_pcd_points'] is None:
            new_obs['object_pcd_points'] = None
            new_obs['object_pcd_normals'] = None
            new_obs['background_pcd_points'] = None     
            new_obs['poke_idx'] = None
            return new_obs

        obj_idx = np.arange(obs['object_pcd_points'].shape[0])
        points = obs['object_pcd_points'][obj_idx]
        normals = obs['object_pcd_normals'][obj_idx]
        
        # Sample points to a fixed length
        obj_idx = sample_idx(len(points), self.object_pcd_size)
        new_obs['object_pcd_points'] = points[obj_idx, :]
        new_obs['object_pcd_normals'] = normals[obj_idx, :]
        bg_idx = sample_idx(len(obs['background_pcd_points']), self.background_pcd_size)
        new_obs['background_pcd_points'] = obs['background_pcd_points'][bg_idx, :]
        
        if self.action_mode == "per_point_action":
            # Us a random policy by default. Will be overwritten by the wrapper.
            location_info = self.random_location_policy.get_action(new_obs)
            new_obs.update(location_info)

        if self.real_robot:
            # Save goal frames
            new_obs['raw_goal_frames'] = obs['goal_frames']
            
        return new_obs
    
    def _evaluate_goal(self, obs, goal):
        if obs['object_pcd_points'] is None:
            reward = 0
            success = False
            return success, reward
    
        if self.reward_mode == "flow":
            old_pcd = obs['object_pcd_points']
            new_pcd = transform_point_cloud(obs['object_pose'], goal, old_pcd)
            flow = np.linalg.norm(new_pcd - old_pcd, axis=-1)
            mean_flow = np.mean(flow, axis=-1)
            reward = -mean_flow

            # Add a distance term
            if self.action_mode in ["regress_action_only", "regress_location_and_action"] and (self.reward_gripper_distance is not None):
                gripper_pos, _ = decompose_pose_mat(obs['gripper_pose'])
                dists = np.linalg.norm(old_pcd - gripper_pos, axis=1)
                min_dist = np.min(dists)
                reward -= max(min_dist - 0.05, 0.0) * self.reward_gripper_distance
            
            success = mean_flow < self.success_threshold
        
        elif self.reward_mode == "pose_diff":
            DeprecationWarning("Pose diff reward is deprecated.")
            object_pos, object_ori = decompose_pose_mat(obs['object_pose'])
            goal_pos, goal_ori = decompose_pose_mat(goal)
            pos_diff = np.linalg.norm(object_pos - goal_pos)
            ori_diff = angle_diff(object_ori, goal_ori) / np.pi * 180.0
            reward = - (50 * pos_diff + 0.2 * ori_diff)
            success = - reward < self.success_threshold
        
        else:
            raise NotImplementedError

        return success, reward

hacman/envs/sim_envs/hacman_bin_env.py
HACMan-2023/HACMan
import numpy as np
from .base_env import BaseEnv, RandomLocation
from hacman.utils.transformations import transform_point_cloud, to_pose_mat, sample_idx
import gym

"""
Abstracted Gym wrapper as an interface for HACMAN
"""
class SimpleEnv(BaseEnv):
    def __init__(self,
            object_pcd_size=400,
            background_pcd_size=400, **kwargs):
        self.sim = SimpleCubeSim()
        super().__init__(object_pcd_size, background_pcd_size, **kwargs)
    
        self.reset()

    def reset(self):
        self.sim.reset()

        # Get the goal pose
        goal_pos = self.sim.get_goal_pos()
        self.goal = self._cube_pos_to_pose_mat(goal_pos)
        
        obs = self.get_observations()
        self.prev_obs = obs
        return obs
        

    def step(self, motion):
        # Take the environment step
        # Use the previous observation to get the location
        points = self.prev_obs['object_pcd_points']
        idx = self.prev_obs['poke_idx']
        location = points[idx.item()]

        # Given the robot-object contact location "location" and the 
        # continuous paramter "motion", Execute the action in simulation.
        self.sim.poke(location, motion)
        
        # Compute the step outcome
        obs = self.get_observations()
        self.prev_obs = obs

        success, reward = self._evaluate_goal(obs, obs['goal_pose'])

        info = {"is_success": success,
                "action_param": motion,
                "action_location": location,}

        done = False

        return obs, reward, done, info
        

    def get_observations(self):
        obs = {}

        obs['object_pose'] = self._cube_pos_to_pose_mat(self.sim.get_cube_pos())
        obs['goal_pose'] = np.copy(self.goal)

        # Point cloud segementation
        points, segmentation = self.sim.get_point_cloud()
        object_points = points[segmentation == 1]        # Cube points
        background_points = points[segmentation == 0]    # Table points
        
        # Sample points to a fixed length
        obj_idx = sample_idx(len(object_points), self.object_pcd_size)
        obs['object_pcd_points'] = object_points[obj_idx, :]
        bg_idx = sample_idx(len(background_points), self.background_pcd_size)
        obs['background_pcd_points'] = background_points[bg_idx, :]
        
        # Choose location here:
        # Use a random location by default. Will be overwritten by the wrapper.
        location_info = self.random_location_policy.get_action(obs)
        obs.update(location_info)
        
        return obs


    def get_normalize_param(self):
        offset = np.zeros(3)
        scale = np.ones(3) * self.sim.table_size
        return offset, scale


    def _evaluate_goal(self, obs, goal):
        # Compute the flow
        old_pcd = obs['object_pcd_points']
        new_pcd = transform_point_cloud(obs['object_pose'], goal, old_pcd)
        flow = np.linalg.norm(new_pcd - old_pcd, axis=1)
        mean_flow = np.mean(flow)

        reward = -mean_flow
        success = mean_flow < 0.05

        return success, reward
    
    def _cube_pos_to_pose_mat(self, pos):
        pos, quat = np.array([*pos, 0]), np.array([0, 0, 0, 1])
        pose_mat = to_pose_mat(pos, quat, input_wxyz=False)
        return pose_mat


class SimpleCubeSim(object):
    """A simple cube simulation for poking and interacting with a cube on a table.

    The simulation initializes a cube with a given size and table size. It provides
    methods for resetting the cube state, poking the cube at a specific location,
    obtaining the goal pose (which is always at the center of the table), retrieving
    the current pose of the cube, and generating a point cloud representing the table
    and the translated cube.

    **See scripts/test_simple_env.py for visualization.**

    Usage:
        sim = SimpleCubeSim()          # Create a SimpleCubeSim object
        sim.reset()                    # Reset the cube state
        sim.poke(location, motion)     # Poke the cube at a specific location with motion
        goal_pose = sim.get_goal_pose()      # Get the goal pose
        object_pose = sim.get_object_pose()  # Get the current pose of the cube
        point_cloud, segmentation = sim.get_point_cloud()
                                    # Get the point cloud representing
                                    the table and the cube, along with the segmentation mask

    Attributes:
        cube_size (float): Half size of the cube.
        table_size (float): Half size of the table.
        table_points (ndarray): Points representing the table in the point cloud.
        cube_points (ndarray): Points representing the cube in the point cloud.
    """
    def __init__(self):
        self.cube_size = 0.1    # cube half size
        self.table_size= 1.0    # table half size

        # For point cloud synthesis
        self.table_points = self._render_table_points(res=0.02)
        self.cube_points = self._render_cube_points(res=0.02)

        self.reset()   # Initialize the cube state
    
    def reset(self, state=None):
        if state is not None:
            self.cube_state = np.copy(state)
        else:
            # self.cube_state = np.random.uniform(
            #     -self.table_size + self.cube_size, self.table_size - self.cube_size, size=2)
            self.cube_state = np.array([0.5, 0.5])
    
    def poke(self, location, motion):
        """
        Simplified poke simulation.

        A force is applied to the cube only when it is pushing
        on a side face. Pulling and poking on the top face does not do anything.

        Args:
            location: 3D location of the poke on the cube
            motion: 3D motion of the poke (delta translation, x, y, z)
        """
        
        if np.isclose(location[2], 2 * self.cube_size): # Top face
            translation = np.zeros(2)
        else:
            normal = self._get_normal(location)
            if np.dot(normal, motion) < 0:
                translation = motion[:2]
            else:
                translation = np.zeros(2)

        # Update the cube state
        self.cube_state[:2] += translation
    
    def get_goal_pos(self):
        goal = np.zeros(2)  # It is always at the center of the table
        return goal

    def get_cube_pos(self):
        return self.cube_state

    def get_point_cloud(self):
        # Synthesize the point cloud by combining the table and the translated cube points
        cube_x, cube_y = self.cube_state
        cube_points = self.cube_points + np.array([cube_x, cube_y, 0.0])
        points = np.concatenate([self.table_points, cube_points], axis=0)

        # Segmentation mask
        segmentation = np.ones(len(points), dtype=np.int32)
        segmentation[:len(self.table_points)] = 0
        return points, segmentation

    def _get_normal(self, location):
        cube_com = np.array([self.cube_state[0], self.cube_state[1], self.cube_size/2.0])
        relative_location = (location - cube_com) / self.cube_size
        
        face_direction_mask = np.isclose(np.abs(relative_location), 1.0)  # Determine if the face faces x, y, or z
        normal = relative_location * face_direction_mask
        return normal
    
    def _render_table_points(self, res=0.02):
        points = []
        for x in np.arange(-self.table_size, self.table_size, res):
            for y in np.arange(-self.table_size, self.table_size, res):
                points.append([x, y, 0.0])
        return np.array(points)

    def _render_cube_points(self, res=0.02):
        points = []
        # Cube sides
        for z in np.arange(0, 2 * self.cube_size, res):
            for x in np.arange(-self.cube_size, self.cube_size + 1e-3, res):
                points.append([x, -self.cube_size, z])
                points.append([x, self.cube_size, z])
            for y in np.arange(-self.cube_size + 1e-3, self.cube_size, res):
                points.append([-self.cube_size, y, z])
                points.append([self.cube_size, y, z])
        # Cube top
        for x in np.arange(-self.cube_size, self.cube_size + 1e-3, res):
            for y in np.arange(-self.cube_size, self.cube_size + 1e-3, res):
                points.append([x, y, 2 * self.cube_size])

        # Add a mask digit to the points
        return np.array(points)

hacman/envs/sim_envs/simple_env.py
HACMan-2023/HACMan
"""
TODO:
- Make a unified plotly plotting function
- Show more than one env?
"""

import os, sys

import numpy as np
import imageio
import pickle
import gym

# from bin_env import PokeEnv
from hacman.utils.transformations import to_pose_mat, transform_point_cloud, decompose_pose_mat

import wandb
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv, _flatten_obs
from hacman.utils.plotly_utils import plot_pcd, plot_action, plot_pcd_with_score
# from bin_env.util import angle_diff
import plotly.graph_objects as go
from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, Union, Dict
from hacman.algos.location_policy import RandomLocation

    
class WandbPointCloudRecorder(VecEnvWrapper):
    """
    Modified based on VecVideoRecorder.
    """

    def __init__(self, venv, global_step_counter=None,
        save_plotly=False, foldername=None, 
        record_video=False, upload_video=False,
        real_robot=False, log_plotly_once=True):

        VecEnvWrapper.__init__(self, venv)
        self.real_robot = real_robot
        # Log plotly once per evaluate_policy call. This is to save space during training.
        self.log_plotly_once = log_plotly_once
        self.env = venv
        # Temp variable to retrieve metadata
        temp_env = venv

        # Unwrap to retrieve metadata dict
        # that will be used by gym recorder
        while isinstance(temp_env, VecEnvWrapper):
            temp_env = temp_env.venv

        if isinstance(temp_env, DummyVecEnv) or isinstance(temp_env, SubprocVecEnv):
            metadata = temp_env.get_attr("metadata")[0]
        else:
            metadata = temp_env.metadata

        self.episode_count = -1
        self.global_step_counter = global_step_counter

        self.env.metadata = metadata
        self.save_plotly = save_plotly
        assert not (foldername is None and self.save_plotly), "Need to specify a folder for WandbPointCloudRecorder."
        self.foldername = os.path.join(foldername, 'plotly')
        os.makedirs(self.foldername, exist_ok=True)

        self.vis = []
        self.recording = False
        self.title = None
        self.step_count = 0
        self.plot_count = 0

        # Video recording
        self.vid_count = 0
        self.record_video = record_video
        self.upload_video = upload_video
        self.vid_dir = os.path.join(foldername, 'video')
        self.cam_frames = []

    def reset(self) -> VecEnvObs:
        # This reset will only be called once per evaluate_policy
        # regardless of how many eval epsiodes. This is because
        # VecEnv will reset themselves.
        obs = self.venv.reset()
        self.reset_plotly_logging(obs)
        self.reset_video_recording()
        return obs
    
    def reset_video_recording(self):
        if len(self.cam_frames) == 0:
            return

        for i in range(len(self.cam_frames)):
            if len(self.cam_frames[i]) == 0:
                continue
            global_step = self.get_global_step()
            is_success = "fail"
            filename = f"video_{self.vid_count}_{global_step}_{i}_{is_success}.mp4"
            self.save_video(self.cam_frames[i], filename)
            self.cam_frames[i].clear()

    def reset_plotly_logging(self, obs):
        self.recording = True
        self.vis = []
        self.title = None        
        self.step_count = 0
        self.episode_count += 1
        self.vis.append(plot_pcd('background', obs['background_pcd_points'][0], 'lightgrey', size=2))

        if self.real_robot:
            # Ground truth goal pcd
            goal_pcd = self.venv.envs[0].unwrapped.env.goal_pcd.voxel_down_sample(0.01)
            goal_pcd = np.asarray(goal_pcd.points)
            self.vis.append(plot_pcd(f'goal', goal_pcd, 'blue'))
            # Estimated goal pcd
            goal_pcd = transform_point_cloud(obs['object_pose'][0], obs['goal_pose'][0], obs['object_pcd_points'][0])
            self.vis.append(plot_pcd(f'g_0', goal_pcd, 'lightblue'))
        else:
            goal_pcd = transform_point_cloud(obs['object_pose'][0], obs['goal_pose'][0], obs['object_pcd_points'][0])
            self.vis.append(plot_pcd(f'g_0', goal_pcd, 'blue'))
        
        if 'action_location_score' in obs.keys():
            self.vis.append(plot_pcd_with_score(f'o_0', obs['object_pcd_points'][0], obs['action_location_score'][0]))
        else:
            self.vis.append(plot_pcd(f'o_0', obs['object_pcd_points'][0], 'yellow'))
        return
    
    def enable_video_recording(self):
        self.record_video = True
    
    def disable_video_recording(self):
        self.record_video = False

    def get_global_step(self):
        if self.global_step_counter is not None:
            global_step = self.global_step_counter()
        else:
            global_step = self.episode_count
        return global_step

    def step_wait(self) -> VecEnvStepReturn:
        obs, rews, dones, infos = self.venv.step_wait()
        
        if self.recording:
            self.title = f"Final reward: {rews[0]:.2f} success: {infos[0]['is_success']}"
            self.step_count += 1
            action_location = infos[0]['action_location'][0]

            if action_location is None:
                # When not predicting action location (contiguous action),
                # plot from previous gripper position
                if (not hasattr(self, "prev_obs")): # Issue with gym0.24 running one step before initialized
                    action_location = np.zeros(3)
                elif self.prev_obs is None:
                    action_location = np.zeros(3)
                else:
                    prev_gripper_pose = self.prev_obs['gripper_pose']
                    action_location, _ = decompose_pose_mat(prev_gripper_pose[0])
            
            if 'poke_success' in infos[0].keys() and not infos[0]['poke_success']:
                self.vis.append(plot_action(f'a_{self.step_count-1}', action_location, infos[0]['action_param'], color='orange'))
            else:
                self.vis.append(plot_action(f'a_{self.step_count-1}', action_location, infos[0]['action_param']))                
                
            if dones[0]:
                # DummyVecEnv will reset venv automatically and replace the returned obs
                obj_pcd = infos[0]['terminal_observation']['object_pcd_points']            
                self.vis.append(plot_pcd(f'o_final', obj_pcd, 'mediumpurple'))
                self.close_video_recorder()
                if not self.log_plotly_once:
                    self.reset_plotly_logging(obs)                
            else:
                if self.real_robot:                    # Estimated goal pcd
                    goal_pcd = transform_point_cloud(obs['object_pose'][0], obs['goal_pose'][0], obs['object_pcd_points'][0])
                    self.vis.append(plot_pcd(f'g_{self.step_count}', goal_pcd, 'lightblue'))
                else:
                    goal_pcd = transform_point_cloud(obs['object_pose'][0], obs['goal_pose'][0], obs['object_pcd_points'][0])
                    self.vis.append(plot_pcd(f'g_{self.step_count}', goal_pcd, 'blue'))
                    
                if 'action_location_score' in obs.keys():
                    self.vis.append(plot_pcd_with_score(f'o_{self.step_count}', obs['object_pcd_points'][0], obs['action_location_score'][0]))
                else:
                    self.vis.append(plot_pcd(f'o_{self.step_count}', obs['object_pcd_points'][0], 'purple'))
                self.vis.append(plot_pcd(f'o_next_{self.step_count-1}', obs['object_pcd_points'][0], 'yellow'))
        
        # Camera recording
        if self.record_video and len(infos[0]['cam_frames']) > 0:
            if len(self.cam_frames) == 0:
                 # Initialize cam_frames
                self.cam_frames = [[] for _ in range(len(infos))]
            
            # Append the poke frames to the corresponding cam_frames
            for i, info in enumerate(infos):
                if len(info['cam_frames']) > 0:
                    self.cam_frames[i].extend(info['cam_frames'])

                if dones[i]:
                    global_step = self.get_global_step()
                    is_success = "success" if infos[i]['is_success'] else "fail"
                    filename = f"video_{self.vid_count}_{global_step}_{i}_{is_success}.mp4"
                    self.save_video(self.cam_frames[i], filename)
                    self.cam_frames[i].clear()

        return obs, rews, dones, infos
    

    def save_video(self, frames, filename, fps=10) -> None:
        assert len(frames) > 0, "No frames to save"
        global_step = self.get_global_step()
        vid_path = os.path.join(self.vid_dir, filename)
        os.makedirs(self.vid_dir, exist_ok=True)

        # Save video locally
        with imageio.get_writer(vid_path, mode='I', fps=fps) as writer:
            for im in frames:
                im = np.flipud(im)
                writer.append_data(im)
        
        self.vid_count += 1

        if self.upload_video:
            wandb.log({"video": wandb.Video(vid_path, fps=fps, format="mp4"), 'global_steps': global_step})
        
        return

    def close_video_recorder(self) -> None:
        if self.recording:
            global_step = self.get_global_step()
            # fig = self.get_plotly_with_slidebar()
            fig = self.get_plotly_with_slidebar()
            # fig.show()  # Visualize locally
            wandb.log({"visualizations": fig, 'global_steps': global_step})

            if self.save_plotly:
                filename = f"plotly_{global_step}_{self.plot_count}.html"
                fig.write_html(os.path.join(self.foldername, filename), auto_play=False)  # Save as html locally
                self.plot_count += 1

        self.recording = False
        self.vis = []
        self.title = None
        self.step_count = 0
    
    def get_plotly_with_slidebar(self):
        vis_name2id = {}
        for i in range(len(self.vis)):
            vis_name2id[self.vis[i].name] = i

        fig = go.Figure(self.vis)
        
        # Default
        for dt in fig.data:
            dt.visible = False
        fig.data[vis_name2id['background']].visible = True    
        if f'goal' in vis_name2id.keys():
            fig.data[vis_name2id['goal']].visible = True 
        fig.data[vis_name2id[f'g_0']].visible = True
        fig.data[vis_name2id[f'o_0']].visible = True
        fig.data[vis_name2id[f'a_0']].visible = True

        fig.update_scenes(aspectmode='data')
        fig.update_layout(title=self.title)

        steps = []
        for i in range(self.step_count):
            step = dict(
                method="update",
                args=[{"visible": [False] * len(fig.data)},
                    {"title": self.title}],
            )
            step["args"][0]["visible"][vis_name2id['background']] = True
            if f'goal' in vis_name2id.keys():
                step["args"][0]["visible"][vis_name2id['goal']] = True
            if f'g_{i}' in vis_name2id.keys():
                step["args"][0]["visible"][vis_name2id[f'g_{i}']] = True                
            step["args"][0]["visible"][vis_name2id[f'o_{i}']] = True
            step["args"][0]["visible"][vis_name2id[f'a_{i}']] = True
            if i == self.step_count - 1:
                step["args"][0]["visible"][vis_name2id[f'o_final']] = True
            else:
                step["args"][0]["visible"][vis_name2id[f'o_next_{i}']] = True
            steps.append(step)

        sliders = [dict(
            active=0,    
            pad={"t": 50},
            steps=steps
        )]
        fig.update_layout(sliders=sliders)
        return fig
    
    def get_plotly_with_slidebar_v2(self):
        vis_name2id = {}
        for i in range(len(self.vis)):
            vis_name2id[self.vis[i].name] = i
        
        traces = {}
        for i in range(self.step_count):
            traces[i] = []
            traces[i].append(self.vis[vis_name2id[f'background']])
            if f'goal' in vis_name2id.keys():
                traces[i].append(self.vis[vis_name2id['goal']])
            if f'g_{i}' in vis_name2id.keys():
                traces[i].append(self.vis[vis_name2id[f'g_{i}']])
            traces[i].append(self.vis[vis_name2id[f'o_{i}']])
            traces[i].append(self.vis[vis_name2id[f'a_{i}']])
            
            if i == self.step_count - 1:
                traces[i].append(self.vis[vis_name2id[f'o_final']])
            else:
                traces[i].append(self.vis[vis_name2id[f'o_next_{i}']])
            
        # Final step
        i = self.step_count        
        traces[i] = []
        traces[i].append(self.vis[vis_name2id[f'background']])
        if f'goal' in vis_name2id.keys():
            traces[i].append(self.vis[vis_name2id['goal']])
        if f'g_{i-1}' in vis_name2id.keys():
            traces[i].append(self.vis[vis_name2id[f'g_{i-1}']])
        traces[i].append(self.vis[vis_name2id[f'o_final']])
        traces[i].append(self.vis[vis_name2id[f'a_{i-1}']])
            
        fig = go.Figure(
            data=traces[0],
            frames=[
                go.Frame(data=traces[k], name=str(k))
                for k in range(self.step_count+1)
            ]
        )
        fig.frames[-1].data[-1].visible = False
        # Every frame needs to have the same number of traces
        
        # Layout
        fig.update_layout(
            title=self.title,
            width=600,
            height=600,
            scene1=dict(aspectmode='data'),
            updatemenus=[
                {
                    "buttons": [
                        {
                            "args": [None, self.frame_args(500)],
                            "label": "&#9654;",  # play symbol
                            "method": "animate",
                        },
                        {
                            "args": [[None], self.frame_args(0)],
                            "label": "&#9724;",  # pause symbol
                            "method": "animate",
                        },
                    ],
                    "direction": "left",
                    "pad": {"r": 10, "t": 70},
                    "type": "buttons",
                    "x": 0.1,
                    "y": 0,
                }
            ],
            sliders= [
                {
                    "pad": {"b": 10, "t": 60},
                    "len": 0.9,
                    "x": 0.1,
                    "y": 0,
                    # "steps": steps,
                    "steps": [
                        {
                            "args": [[f.name], self.frame_args(0)],
                            "label": f'step {str(k)}',
                            "method": "animate",
                        }
                        for k, f in enumerate(fig.frames)
                    ],
                }
            ],
        )
        
        return fig
    
    def frame_args(self, duration):
        return {
            "frame": {"duration": duration},
            "mode": "immediate",
            "fromcurrent": True,
            "transition": {"duration": duration, "easing": "linear"},
        }
    

    def close(self) -> None:
        self.close_video_recorder()

    # def __del__(self):
    #     self.close()

hacman/envs/wandb_wrappers.py
HACMan-2023/HACMan
import torch
from hacman.networks.point_transformer import PointTransformerSegmentation
from hacman.networks.pointnet2 import PointNet2Classification, PointNet2Segmentation
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def add_model_config(parser):
    parser.add_argument("--model", default="pn", type=str, help="Model Name")
    parser.add_argument("--fps_deterministic", action="store_true", help="Turn off fps random start")
    parser.add_argument("--pos_in_feature", action="store_true", help="Include absolute point position in feature")
    parser.add_argument("--dropout", default=0., type=float, help="dropout")
    parser.add_argument("--normalize_pos", default=False, action='store_true')
    return

def init_network(config, input_channels, output_channels, mode=None):
    if config["model"] == 'pt':
        network = PointTransformerSegmentation(input_channels, output_channels, 
                                            dim_model=[32, 64, 128, 256, 512], 
                                            k=16, 
                                            fps_random_start=not config['fps_deterministic'],
                                            pos_in_feature=config["pos_in_feature"],
                                            normalize_pos=config["normalize_pos"]).to(device)
    elif config["model"] == 'pn':
        network = PointNet2Segmentation(input_channels, output_channels, 
                                    dropout=config['dropout'],
                                    fps_random_start=not config['fps_deterministic'],
                                    pos_in_feature=config["pos_in_feature"],
                                    normalize_pos=config["normalize_pos"],
                                    normalize_pos_param=config["normalize_pos_param"]).to(device)
    elif config['model'] == 'pnc':
        network = PointNet2Classification(input_channels, out_channels=output_channels, 
                dropout=config['dropout'],
                fps_random_start=not config['fps_deterministic'],
                pos_in_feature=config["pos_in_feature"],
                normalize_pos=config["normalize_pos"],
                normalize_pos_param=config["normalize_pos_param"],
                mode=mode).to(device)

    else:
        print(f"Model '{config['model']}' does not exist.")
        assert NotImplementedError
    return network

def load_network(network, load_path):
    loaded = torch.load(load_path, map_location=device)
    if 'model_state_dict' in loaded.keys():
        network.load_state_dict(loaded['model_state_dict'])
    else:
        # Backward compatibility
        network.load_state_dict(loaded)
    return network

hacman/networks/common.py
HACMan-2023/HACMan
import torch
import torch.nn.functional as F
from torch.nn import BatchNorm1d as BN
from torch.nn import Identity
from torch.nn import Linear as Lin
from torch.nn import ReLU
from torch.nn import Sequential as Seq
from torch_cluster import fps, knn_graph
from torch_scatter import scatter_max
from torch_geometric.nn import global_max_pool, global_mean_pool
from torch_geometric.nn.conv import PointTransformerConv
from torch_geometric.nn.pool import knn
from torch_geometric.nn.unpool import knn_interpolate

"""
Point transformer: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/point_transformer_segmentation.py#L63
"""


class TransformerBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.lin_in = Lin(in_channels, in_channels)
        self.lin_out = Lin(out_channels, out_channels)

        self.pos_nn = MLP([3, 64, out_channels], batch_norm=False)

        self.attn_nn = MLP([out_channels, 64, out_channels], batch_norm=False)

        self.transformer = PointTransformerConv(in_channels, out_channels,
                                                pos_nn=self.pos_nn,
                                                attn_nn=self.attn_nn)

    def forward(self, x, pos, edge_index):
        x = self.lin_in(x).relu()
        x = self.transformer(x, pos, edge_index)
        x = self.lin_out(x).relu()
        return x


class TransitionDown(torch.nn.Module):
    '''
        Samples the input point cloud by a ratio percentage to reduce
        cardinality and uses an mlp to augment features dimensionnality
    '''
    def __init__(self, in_channels, out_channels, ratio=0.25, k=16, fps_random_start=True):
        super().__init__()
        self.k = k
        self.ratio = ratio
        self.mlp = MLP([in_channels, out_channels])
        self.fps_random_start = fps_random_start

    def forward(self, x, pos, batch):
        # FPS sampling
        id_clusters = fps(pos, ratio=self.ratio, batch=batch, random_start=self.fps_random_start)

        # compute for each cluster the k nearest points
        sub_batch = batch[id_clusters] if batch is not None else None

        # beware of self loop
        id_k_neighbor = knn(pos, pos[id_clusters], k=self.k, batch_x=batch,
                            batch_y=sub_batch)

        # transformation of features through a simple MLP
        x = self.mlp(x)

        # Max pool onto each cluster the features from knn in points
        x_out, _ = scatter_max(x[id_k_neighbor[1]], id_k_neighbor[0],
                               dim_size=id_clusters.size(0), dim=0)

        # keep only the clusters and their max-pooled features
        sub_pos, out = pos[id_clusters], x_out
        return out, sub_pos, sub_batch


def MLP(channels, batch_norm=True):
    return Seq(*[
        Seq(Lin(channels[i - 1], channels[i]),
            BN(channels[i]) if batch_norm else Identity(), ReLU())
        for i in range(1, len(channels))
    ])


class TransitionUp(torch.nn.Module):
    '''
        Reduce features dimensionnality and interpolate back to higher
        resolution and cardinality
    '''
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.mlp_sub = MLP([in_channels, out_channels])
        self.mlp = MLP([out_channels, out_channels])

    def forward(self, x, x_sub, pos, pos_sub, batch=None, batch_sub=None):
        # transform low-res features and reduce the number of features
        x_sub = self.mlp_sub(x_sub)

        # interpolate low-res feats to high-res points
        x_interpolated = knn_interpolate(x_sub, pos_sub, pos, k=3,
                                         batch_x=batch_sub, batch_y=batch)

        x = self.mlp(x) + x_interpolated

        return x


class PointTransformerSegmentation(torch.nn.Module):
    def __init__(self, in_channels, out_channels, dim_model, k=16, fps_random_start=True, 
                 pos_in_feature=False, normalize_pos=False):
        super().__init__()
        self.k = k

        self.normalize_pos = normalize_pos
        self.pos_in_feature = pos_in_feature
        if self.pos_in_feature:
            in_channels += 3
            
        # dummy feature is created if there is none given
        in_channels = max(in_channels, 1)

        # first block
        self.mlp_input = MLP([in_channels, dim_model[0]])

        self.transformer_input = TransformerBlock(
            in_channels=dim_model[0],
            out_channels=dim_model[0],
        )

        # backbone layers
        self.transformers_up = torch.nn.ModuleList()
        self.transformers_down = torch.nn.ModuleList()
        self.transition_up = torch.nn.ModuleList()
        self.transition_down = torch.nn.ModuleList()

        for i in range(0, len(dim_model) - 1):

            # Add Transition Down block followed by a Point Transformer block
            self.transition_down.append(
                TransitionDown(in_channels=dim_model[i],
                               out_channels=dim_model[i + 1], k=self.k, fps_random_start=fps_random_start))

            self.transformers_down.append(
                TransformerBlock(in_channels=dim_model[i + 1],
                                 out_channels=dim_model[i + 1]))

            # Add Transition Up block followed by Point Transformer block
            self.transition_up.append(
                TransitionUp(in_channels=dim_model[i + 1],
                             out_channels=dim_model[i]))

            self.transformers_up.append(
                TransformerBlock(in_channels=dim_model[i],
                                 out_channels=dim_model[i]))

        # summit layers
        self.mlp_summit = MLP([dim_model[-1], dim_model[-1]], batch_norm=False)

        self.transformer_summit = TransformerBlock(
            in_channels=dim_model[-1],
            out_channels=dim_model[-1],
        )

        out_channels = [out_channels] if type(out_channels) is not list else out_channels
        
        mlps = []        
        for out_dim in out_channels:
            mlps.append(Seq(Lin(dim_model[0], 128), ReLU(), Lin(128, 128), ReLU(), Lin(128, 128), ReLU(), Lin(128, out_dim)))
        self.mlps = torch.nn.ModuleList(mlps)

    def forward(self, x, pos, batch=None):
        if self.normalize_pos:
            pos = pos.clone()
            mean_pos = global_mean_pool(pos, batch)
            mean_pos_expanded = torch.index_select(mean_pos, 0, batch)
            pos -= mean_pos_expanded
            
            max_pos = global_max_pool(pos.abs(), batch).max(dim=1)[0]
            max_pos_expanded = torch.index_select(max_pos, 0, batch).reshape(-1, 1)
            pos *= (1 / max_pos_expanded) * 0.999999
            
        # add dummy features in case there is none
        if x is None:
            x = torch.ones((pos.shape[0], 1)).to(pos.get_device())

        if self.pos_in_feature:
            x = torch.cat([x, pos], dim=1)

        out_x = []
        out_pos = []
        out_batch = []

        # first block
        x = self.mlp_input(x)
        edge_index = knn_graph(pos, k=self.k, batch=batch)
        x = self.transformer_input(x, pos, edge_index)

        # save outputs for skipping connections
        out_x.append(x)
        out_pos.append(pos)
        out_batch.append(batch)

        # backbone down : #reduce cardinality and augment dimensionnality
        for i in range(len(self.transformers_down)):
            x, pos, batch = self.transition_down[i](x, pos, batch=batch)
            edge_index = knn_graph(pos, k=self.k, batch=batch)
            x = self.transformers_down[i](x, pos, edge_index)

            out_x.append(x)
            out_pos.append(pos)
            out_batch.append(batch)

        # summit
        x = self.mlp_summit(x)
        edge_index = knn_graph(pos, k=self.k, batch=batch)
        x = self.transformer_summit(x, pos, edge_index)

        # backbone up : augment cardinality and reduce dimensionnality
        n = len(self.transformers_down)
        for i in range(n):
            x = self.transition_up[-i - 1](x=out_x[-i - 2], x_sub=x,
                                           pos=out_pos[-i - 2],
                                           pos_sub=out_pos[-i - 1],
                                           batch_sub=out_batch[-i - 1],
                                           batch=out_batch[-i - 2])

            edge_index = knn_graph(out_pos[-i - 2], k=self.k,
                                   batch=out_batch[-i - 2])
            x = self.transformers_up[-i - 1](x, out_pos[-i - 2], edge_index)

        if len(self.mlps) == 0:
            return x
        
        out_list = []
        for mlp in self.mlps:
            out_list.append(mlp(x))
        return torch.cat(out_list, dim=1) # F.log_softmax(out, dim=-1)

hacman/networks/point_transformer.py
HACMan-2023/HACMan
import torch
from torch_geometric.nn import MLP, PointConv, fps, global_max_pool, radius, global_mean_pool
from torch_geometric.nn import MLP, knn_interpolate

import numpy as np


def normalize_pcd(pos, batch, normalize_pos_param):
    pos = pos.clone()
    if normalize_pos_param is None:
        mean_pos = global_mean_pool(pos, batch)
        mean_pos_expanded = torch.index_select(mean_pos, 0, batch)
        pos -= mean_pos_expanded
        
        max_pos = global_max_pool(pos.abs(), batch).max(dim=1)[0]
        max_pos_expanded = torch.index_select(max_pos, 0, batch).reshape(-1, 1)
        pos *= (1 / max_pos_expanded) * 0.999999
    
    else:
        offset, full_scale = normalize_pos_param
        offset = torch.from_numpy(offset).float().to(pos.device)
        full_scale = torch.from_numpy(full_scale).float().to(pos.device)
        
        pos -= offset
        pos /= torch.max(full_scale/2)
    
    return pos

class SAModule(torch.nn.Module):
    def __init__(self, ratio, r, nn, fps_random_start=True):
        super().__init__()
        self.ratio = ratio
        self.r = r
        self.conv = PointConv(nn, add_self_loops=False)
        self.fps_random_start = fps_random_start

    def forward(self, x, pos, batch):
        idx = fps(pos, batch, ratio=self.ratio, random_start=self.fps_random_start)
        row, col = radius(pos, pos[idx], self.r, batch, batch[idx],
                          max_num_neighbors=64)
        edge_index = torch.stack([col, row], dim=0)
        x_dst = None if x is None else x[idx]
        x = self.conv((x, x_dst), (pos, pos[idx]), edge_index)
        pos, batch = pos[idx], batch[idx]
        return x, pos, batch


class GlobalSAModule(torch.nn.Module):
    def __init__(self, nn):
        super().__init__()
        self.nn = nn

    def forward(self, x, pos, batch):
        x = self.nn(torch.cat([x, pos], dim=1))
        x = global_max_pool(x, batch)
        pos = pos.new_zeros((x.size(0), 3))
        batch = torch.arange(x.size(0), device=batch.device)
        return x, pos, batch


class FPModule(torch.nn.Module):
    def __init__(self, k, nn):
        super().__init__()
        self.k = k
        self.nn = nn

    def forward(self, x, pos, batch, x_skip, pos_skip, batch_skip):
        x = knn_interpolate(x, pos, pos_skip, batch, batch_skip, k=self.k)
        if x_skip is not None:
            x = torch.cat([x, x_skip], dim=1)
        x = self.nn(x)
        return x, pos_skip, batch_skip


class PointNet2Segmentation(torch.nn.Module):
    def __init__(self, in_channels, out_channels, dropout=0, fps_random_start=True, 
                 pos_in_feature=False, normalize_pos=False, normalize_pos_param=None):
        super().__init__()
        
        self.normalize_pos = normalize_pos
        self.normalize_pos_param = normalize_pos_param
        self.pos_in_feature = pos_in_feature
        if self.pos_in_feature:
            in_channels += 3
        # Input channels account for both `pos` and node features.
        self.sa1_module = SAModule(0.2, 0.2, MLP([in_channels + 3, 64, 64, 128]), fps_random_start=fps_random_start)
        self.sa2_module = SAModule(0.25, 0.4, MLP([128 + 3, 128, 128, 256]), fps_random_start=fps_random_start)
        self.sa3_module = GlobalSAModule(MLP([256 + 3, 256, 512, 1024]))

        self.fp3_module = FPModule(1, MLP([1024 + 256, 256, 256]))
        self.fp2_module = FPModule(3, MLP([256 + 128, 256, 128]))
        self.fp1_module = FPModule(3, MLP([128 + in_channels, 128, 128, 128]))

        out_channels = [out_channels] if type(out_channels) is not list else out_channels
        mlps = []        
        for out_dim in out_channels:
            mlps.append(MLP([128, 128, 128, out_dim], dropout=dropout, batch_norm=False))
        self.mlps = torch.nn.ModuleList(mlps)

    def forward(self, x, pos, batch=None):
        if self.normalize_pos:
            pos = normalize_pcd(pos, batch, self.normalize_pos_param)
            
        if self.pos_in_feature:
            x = torch.cat([x, pos], dim=1)

        sa0_out = (x, pos, batch)
        sa1_out = self.sa1_module(*sa0_out)
        sa2_out = self.sa2_module(*sa1_out)
        sa3_out = self.sa3_module(*sa2_out)

        fp3_out = self.fp3_module(*sa3_out, *sa2_out)
        fp2_out = self.fp2_module(*fp3_out, *sa1_out)
        x, _, _ = self.fp1_module(*fp2_out, *sa0_out)

        if len(self.mlps) == 0:
            # Just use a feature extractor
            out_list = [x]
        else:            
            out_list = []
            for mlp in self.mlps:
                out_list.append(mlp(x))
        return torch.cat(out_list, dim=1) # F.log_softmax(out, dim=-1)

class PointNet2Classification(torch.nn.Module):
    def __init__(self, in_channels, out_channels, dropout=0, fps_random_start=True, pos_in_feature=False, normalize_pos=False, 
                normalize_pos_param=None, mode=None):
        super().__init__()
        self.pos_in_feature = pos_in_feature
        if self.pos_in_feature:
            in_channels += 3
        # Input channels account for both `pos` and node features.
        self.sa1_module = SAModule(0.5, 0.2, MLP([in_channels + 3, 64, 64, 128]), fps_random_start=fps_random_start)
        self.sa2_module = SAModule(0.25, 0.4, MLP([128 + 3, 128, 128, 256]), fps_random_start=fps_random_start)
        self.sa3_module = GlobalSAModule(MLP([256 + 3, 256, 512, 1024]))
        self.normalize_pos = normalize_pos
        self.normalize_pos_param = normalize_pos_param

        
        
        out_channels = [out_channels] if type(out_channels) is not list else out_channels
        mlps = []        
        for out_dim in out_channels:
            mlps.append(MLP([1024, 512, 256, out_dim], dropout=dropout, batch_norm=False))
        self.mlps = torch.nn.ModuleList(mlps)
        self.mode = mode

    def forward(self, x, pos, batch=None):
        if self.normalize_pos:
            pos = normalize_pcd(pos, batch, self.normalize_pos_param)
            
        if self.pos_in_feature:
            x = torch.cat([x, pos], dim=1)
        sa0_out = (x, pos, batch)
        sa1_out = self.sa1_module(*sa0_out)
        sa2_out = self.sa2_module(*sa1_out)
        sa3_out = self.sa3_module(*sa2_out)
        x, pos, batch = sa3_out
        out_list = []
        
        if len(self.mlps) == 0:
            return x
        
        else:
            if self.mode == "binary_classification":
                for mlp in self.mlps:
                    out_list.append(torch.nn.functional.sigmoid(mlp(x)))
            else:
                for mlp in self.mlps:
                    out_list.append(mlp(x))
            return torch.cat(out_list, dim=1) # self.mlp(x).log_softmax(dim=-1)
        

if __name__ == '__main__':
    def count_parameters(model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    seg_model = PointNet2Segmentation(3, [])
    seg_count = count_parameters(seg_model)
    # 1381760 + 49152 = 1430912
    
    cls_model = PointNet2Classification(3, [])
    cls_count = count_parameters(cls_model)
    # 805120 + 524288 + 131072 = 1460480
    
    print(seg_count, cls_count)
hacman/networks/pointnet2.py
HACMan-2023/HACMan
import os

import gym
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

from stable_baselines3.common.callbacks import BaseCallback

class CustomizedCheckpointCallback(BaseCallback):
    """
    Modified based on CheckpointCallback
    
    Callback for saving a model every ``save_freq`` calls
    to ``env.step()``.
    By default, it only saves model checkpoints,
    you need to pass ``save_replay_buffer=True``,
    and ``save_vecnormalize=True`` to also save replay buffer checkpoints
    and normalization statistics checkpoints.

    .. warning::

      When using multiple environments, each call to  ``env.step()``
      will effectively correspond to ``n_envs`` steps.
      To account for that, you can use ``save_freq = max(save_freq // n_envs, 1)``

    :param save_freq_latest: Save latest model and replay buffer every ``save_freq_latest`` call of the callback.
    :param save_freq_checkpoint: Save model checkpoints every ``save_freq_checkpoint`` call of the callback.
    :param save_path: Path to the folder where the model will be saved.
    :param name_prefix: Common prefix to the saved models
    :param save_replay_buffer: Save the model replay buffer
    :param save_vecnormalize: Save the ``VecNormalize`` statistics
    :param verbose: Verbosity level: 0 for no output, 2 for indicating when saving model checkpoint
    """

    def __init__(
        self,
        save_freq_latest: int,
        save_freq_checkpoint: int,
        save_path: str,
        name_prefix: str = "rl_model",
        save_replay_buffer: bool = False,
        save_vecnormalize: bool = False,
        verbose: int = 0,
    ):
        super().__init__(verbose)
        self.save_freq_latest = save_freq_latest
        self.save_freq_checkpoint = save_freq_checkpoint
        self.save_path = save_path
        self.name_prefix = name_prefix
        self.save_replay_buffer = save_replay_buffer
        self.save_vecnormalize = save_vecnormalize

    def _init_callback(self) -> None:
        # Create folder if needed
        if self.save_path is not None:
            os.makedirs(self.save_path, exist_ok=True)

    def _checkpoint_path(self, checkpoint_type: str = "", extension: str = "", latest=False) -> str:
        """
        Helper to get checkpoint path for each type of checkpoint.

        :param checkpoint_type: empty for the model, "replay_buffer_"
            or "vecnormalize_" for the other checkpoints.
        :param extension: Checkpoint file extension (zip for model, pkl for others)
        :return: Path to the checkpoint
        """
        if latest:
            return os.path.join(self.save_path, f"{self.name_prefix}_{checkpoint_type}latest.{extension}")
        else:
            return os.path.join(self.save_path, f"{self.name_prefix}_{checkpoint_type}{self.num_timesteps}_steps.{extension}")

    def _on_step(self) -> bool:
        if self.n_calls % self.save_freq_latest == 0:
            self._save(latest=True, save_replay_buffer=self.save_replay_buffer)
        if self.n_calls % self.save_freq_checkpoint == 0:
            self._save(latest=False)
        return True
    
    def _save(self, latest=False, save_replay_buffer=False):
        model_path = self._checkpoint_path(extension="zip", latest=latest)
        self.model.save(model_path)
        if self.verbose >= 2:
            print(f"Saving model checkpoint to {model_path}")

        if save_replay_buffer and hasattr(self.model, "replay_buffer") and self.model.replay_buffer is not None:
            # If model has a replay buffer, save it too
            replay_buffer_path = self._checkpoint_path("replay_buffer_", extension="pkl", latest=latest)
            self.model.save_replay_buffer(replay_buffer_path)
            if self.verbose > 1:
                print(f"Saving model replay buffer checkpoint to {replay_buffer_path}")

        if self.save_vecnormalize and self.model.get_vec_normalize_env() is not None:
            # Save the VecNormalize statistics
            vec_normalize_path = self._checkpoint_path("vecnormalize_", extension="pkl", latest=latest)
            self.model.get_vec_normalize_env().save(vec_normalize_path)
            if self.verbose >= 2:
                print(f"Saving model VecNormalize to {vec_normalize_path}")
        return True


class ProgressBarCallback(BaseCallback):
    """
    Display a progress bar when training SB3 agent
    using tqdm and rich packages.
    """

    def __init__(self) -> None:
        super().__init__()
        if tqdm is None:
            raise ImportError(
                "You must install tqdm and rich in order to use the progress bar callback. "
                "It is included if you install stable-baselines with the extra packages: "
                "`pip install stable-baselines3[extra]`"
            )
        self.pbar = None

    def _on_training_start(self) -> None:
        # Initialize progress bar
        # Remove timesteps that were done in previous training sessions
        self.pbar = tqdm(total=self.locals["total_timesteps"] - self.model.num_timesteps)

    def _on_step(self) -> bool:
        # Update progress bar, we do num_envs steps per call to `env.step()`
        self.pbar.update(self.training_env.num_envs)
        return True

    def _on_training_end(self) -> None:
        # Close progress bar
        self.pbar.close()
hacman/sb3_utils/custom_callbacks.py
HACMan-2023/HACMan
import warnings
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import pickle
import gym
import numpy as np
from collections import defaultdict

from stable_baselines3.common import base_class
from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, VecMonitor, is_vecenv_wrapped


def evaluate_policy(
    model: "base_class.BaseAlgorithm",
    env: Union[gym.Env, VecEnv],
    n_eval_episodes: int = 10,
    deterministic: bool = True,
    render: bool = False,
    callback: Optional[Callable[[Dict[str, Any], Dict[str, Any]], None]] = None,
    reward_threshold: Optional[float] = None,
    return_episode_rewards: bool = False,
    warn: bool = True,
    save_path: str = None,
    pbar = None,
    per_object_succ = False,
    verbose = False,
) -> Union[Tuple[float, float], Tuple[List[float], List[int]]]:
    """
    Runs policy for ``n_eval_episodes`` episodes and returns average reward.
    If a vector env is passed in, this divides the episodes to evaluate onto the
    different elements of the vector env. This static division of work is done to
    remove bias. See https://github.com/DLR-RM/stable-baselines3/issues/402 for more
    details and discussion.

    .. note::
        If environment has not been wrapped with ``Monitor`` wrapper, reward and
        episode lengths are counted as it appears with ``env.step`` calls. If
        the environment contains wrappers that modify rewards or episode lengths
        (e.g. reward scaling, early episode reset), these will affect the evaluation
        results as well. You can avoid this by wrapping environment with ``Monitor``
        wrapper before anything else.

    :param model: The RL agent you want to evaluate.
    :param env: The gym environment or ``VecEnv`` environment.
    :param n_eval_episodes: Number of episode to evaluate the agent
    :param deterministic: Whether to use deterministic or stochastic actions
    :param render: Whether to render the environment or not
    :param callback: callback function to do additional checks,
        called after each step. Gets locals() and globals() passed as parameters.
    :param reward_threshold: Minimum expected reward per episode,
        this will raise an error if the performance is not met
    :param return_episode_rewards: If True, a list of rewards and episode lengths
        per episode will be returned instead of the mean.
    :param warn: If True (default), warns user about lack of a Monitor wrapper in the
        evaluation environment.
    :return: Mean reward per episode, std of reward per episode.
        Returns ([float], [int]) when ``return_episode_rewards`` is True, first
        list containing per-episode rewards and second containing per-episode lengths
        (in number of steps).
    """
    is_monitor_wrapped = False
    # Avoid circular import
    from stable_baselines3.common.monitor import Monitor

    if not isinstance(env, VecEnv):
        env = DummyVecEnv([lambda: env])

    is_monitor_wrapped = is_vecenv_wrapped(env, VecMonitor) or env.env_is_wrapped(Monitor)[0]

    if not is_monitor_wrapped and warn:
        warnings.warn(
            "Evaluation environment is not wrapped with a ``Monitor`` wrapper. "
            "This may result in reporting modified episode lengths and rewards, if other wrappers happen to modify these. "
            "Consider wrapping environment first with ``Monitor`` wrapper.",
            UserWarning,
        )

    n_envs = env.num_envs
    episode_rewards = []
    episode_lengths = []
    traj_list = []

    episode_counts = np.zeros(n_envs, dtype="int")
    # Divides episodes among different sub environments in the vector as evenly as possible
    episode_count_targets = np.array([(n_eval_episodes + i) // n_envs for i in range(n_envs)], dtype="int")

    current_rewards = np.zeros(n_envs)
    current_lengths = np.zeros(n_envs, dtype="int")
    success_buffer = []
    verbose_buffer = []
    if per_object_succ:
        per_object_results = defaultdict(list)
    observations = env.reset()
    
    traj_list.append(observations)
    states = None
    episode_starts = np.ones((env.num_envs,), dtype=bool)
    
    if callback is not None:
        callback(locals(), globals())
    
    while (episode_counts < episode_count_targets).any():
        actions, states = model.predict(observations, state=states, episode_start=episode_starts, deterministic=deterministic)
        observations, rewards, dones, infos = env.step(actions)
        traj_list.append((actions, observations, rewards, dones, infos))
        
        current_rewards += rewards
        current_lengths += 1
        for i in range(n_envs):
            if episode_counts[i] < episode_count_targets[i]:

                # unpack values so that the callback can access the local variables
                reward = rewards[i]
                done = dones[i]
                info = infos[i]
                episode_starts[i] = done

                if callback is not None:
                    callback(locals(), globals())

                if dones[i]:
                    if save_path is not None:
                        pickle.dump(traj_list, open(save_path, 'wb'))
        
                    if is_monitor_wrapped:
                        # Atari wrapper can send a "done" signal when
                        # the agent loses a life, but it does not correspond
                        # to the true end of episode
                        if "episode" in info.keys():
                            # Do not trust "done" with episode endings.
                            # Monitor wrapper includes "episode" key in info if environment
                            # has been wrapped with it. Use those rewards instead.
                            episode_rewards.append(info["episode"]["r"])
                            episode_lengths.append(info["episode"]["l"])
                            # Only increment at the real end of an episode
                            episode_counts[i] += 1
                    else:
                        episode_rewards.append(current_rewards[i])
                        episode_lengths.append(current_lengths[i])
                        episode_counts[i] += 1
                    current_rewards[i] = 0
                    current_lengths[i] = 0

                    maybe_is_success = info.get("is_success")
                    if maybe_is_success is not None:
                        if per_object_succ:
                            per_object_results[info['object_name']].append(maybe_is_success)
                        success_buffer.append(maybe_is_success)
                        verbose_buffer.append([
                            maybe_is_success,
                            i,
                            current_lengths[i],
                            reward
                        ])
                    
                    if pbar is not None:
                        pbar.update(1)

        if render:
            env.render()

    if pbar is not None:
        pbar.update(np.sum(dones))
        
    mean_reward = np.mean(episode_rewards)
    std_reward = np.std(episode_rewards)
    if reward_threshold is not None:
        assert mean_reward > reward_threshold, "Mean reward below threshold: " f"{mean_reward:.2f} < {reward_threshold:.2f}"
    if return_episode_rewards:
        return episode_rewards, episode_lengths, success_buffer
    
    if len(success_buffer) > 0:
        success_rate = np.average(success_buffer)

    if per_object_succ:
        return mean_reward, std_reward, success_rate, per_object_results
    else:
        if verbose:
            return mean_reward, std_reward, success_rate, verbose_buffer
        else:
            return mean_reward, std_reward, success_rate


hacman/sb3_utils/evaluation.py
HACMan-2023/HACMan
from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm

import os

def load_checkpoint(path, model: OffPolicyAlgorithm):
    assert path is not None
    
    if path == "latest":
        #TODO: find the path to the latest checkpoint
        pass
    
    dir = os.path.dirname(path)
    ckpt_name = path.split("/")[-1].rstrip(".zip")
    step_num = ckpt_name.split("_")[-2]
    
    replay_name = f"rl_model_replay_buffer_{step_num}_steps.pkl"
    replay_path = os.path.join(dir, replay_name)
    
    loaded_model = model.__class__.load(path, env=model.get_env())
    loaded_model.load_replay_buffer(replay_path)
    
    return loaded_model
        
hacman/sb3_utils/load_checkpoint.py
HACMan-2023/HACMan
import os
import numpy as np
import logging
import torch

def use_freer_gpu():
    """
    Set the program to use the freer GPU. Run this at the beginning of the program.
    """
    try:
        os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
        memory_available = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
        if len(memory_available) == 1:
            return os.environ['CUDA_VISIBLE_DEVICES']

        gpu_idx = np.argmax(memory_available).item()
        os.environ['CUDA_VISIBLE_DEVICES'] = f'{gpu_idx}'
        # torch.cuda.set_device(gpu_idx)
        print(f"Free GPU: {gpu_idx}")
        
    except:
        gpu_idx = 0
    
    return gpu_idx
hacman/utils/launch_utils.py
HACMan-2023/HACMan
import numpy as np
import plotly.graph_objects as go
import matplotlib.pyplot as plt

def format_color(color):
    if type(color) == str:
        color_str = color
    elif type(color) == np.ndarray and len(color) == 3:
        color = color.astype(np.int32)
        color_str = f'rgb({color[0]},{color[1]},{color[2]})'
    else:
        color = color.astype(np.int32)
        color_str = [f'rgb({color[k, 0]},{color[k, 1]},{color[k, 2]})' for k in range(color.shape[0])]
    return color_str

def plot_pcd(name, points, color, size=3):
    return go.Scatter3d(x=points[:, 0], y=points[:, 1], z=points[:, 2], mode="markers", 
                        marker = dict(color=format_color(color), size=size), name=name)

def plot_pcd_with_score(name, points, action_score, size=3):
    assert type(action_score) == np.ndarray
    action_score = action_score.reshape(-1)
    action_score = (action_score - action_score.min())/(action_score.max()-action_score.min()+1e-7)
    object_color_id = (action_score*255).astype(np.int32)
    object_color = []
    for cid in object_color_id:
        try:
            object_color.append(plt.get_cmap('plasma').colors[cid])
        except:
            print(f'cid={cid} gives an error. Will use 0 instead.')
            object_color.append(plt.get_cmap('plasma').colors[0])
    object_color = np.array(object_color)*255
    return plot_pcd(name, points, object_color, size)

def plot_action(name, start, direction, color='red', size=3):
    direction = direction*0.02*3  # action scale=0.02. steps=10
    if start is None: x, y, z = 0, 0, 0
    else: x, y, z = start[0], start[1], start[2]
    u, v, w = direction[0], direction[1], direction[2]
    return go.Scatter3d(x=[x, x + u], y=[y, y + v], z=[z, z + w], mode='lines',
                        line=dict(color=format_color(color), width=10), name=name)
hacman/utils/plotly_utils.py
HACMan-2023/HACMan
"""
Utility functions of matrix and vector transformations.

NOTE: convention for quaternions is (x, y, z, w)
"""

import math
import numpy as np

# JIT_decorator copied from robosuite util code
# from robosuite.utils.numba import jit_decorator
import numba
ENABLE_NUMBA = True
CACHE_NUMBA = True
def jit_decorator(func):
    if ENABLE_NUMBA:
        return numba.jit(nopython=True, cache=CACHE_NUMBA)(func)
    return func


PI = np.pi
EPS = np.finfo(float).eps * 4.

# axis sequences for Euler angles
_NEXT_AXIS = [1, 2, 0, 1]

# map axes strings to/from tuples of inner axis, parity, repetition, frame
_AXES2TUPLE = {
    "sxyz": (0, 0, 0, 0),
    "sxyx": (0, 0, 1, 0),
    "sxzy": (0, 1, 0, 0),
    "sxzx": (0, 1, 1, 0),
    "syzx": (1, 0, 0, 0),
    "syzy": (1, 0, 1, 0),
    "syxz": (1, 1, 0, 0),
    "syxy": (1, 1, 1, 0),
    "szxy": (2, 0, 0, 0),
    "szxz": (2, 0, 1, 0),
    "szyx": (2, 1, 0, 0),
    "szyz": (2, 1, 1, 0),
    "rzyx": (0, 0, 0, 1),
    "rxyx": (0, 0, 1, 1),
    "ryzx": (0, 1, 0, 1),
    "rxzx": (0, 1, 1, 1),
    "rxzy": (1, 0, 0, 1),
    "ryzy": (1, 0, 1, 1),
    "rzxy": (1, 1, 0, 1),
    "ryxy": (1, 1, 1, 1),
    "ryxz": (2, 0, 0, 1),
    "rzxz": (2, 0, 1, 1),
    "rxyz": (2, 1, 0, 1),
    "rzyz": (2, 1, 1, 1),
}

_TUPLE2AXES = dict((v, k) for k, v in _AXES2TUPLE.items())


def convert_quat(q, to="xyzw"):
    """
    Converts quaternion from one convention to another.
    The convention to convert TO is specified as an optional argument.
    If to == 'xyzw', then the input is in 'wxyz' format, and vice-versa.

    Args:
        q (np.array): a 4-dim array corresponding to a quaternion
        to (str): either 'xyzw' or 'wxyz', determining which convention to convert to.
    """
    if q.ndim == 1:
        if to == "xyzw":
            return q[[1, 2, 3, 0]]
        if to == "wxyz":
            return q[[3, 0, 1, 2]]
    elif q.ndim == 2:
        if to == "xyzw":
            return q[:, [1, 2, 3, 0]]
        if to == "wxyz":
            return q[:, [3, 0, 1, 2]]
    else:
        raise ValueError("convert_quat: input quaternion must be 1 or 2 dimensional")
    
    raise Exception("convert_quat: choose a valid `to` argument (xyzw or wxyz)")


def quat_multiply(quaternion1, quaternion0):
    """
    Return multiplication of two quaternions (q1 * q0).

    E.g.:
    >>> q = quat_multiply([1, -2, 3, 4], [-5, 6, 7, 8])
    >>> np.allclose(q, [-44, -14, 48, 28])
    True

    Args:
        quaternion1 (np.array): (x,y,z,w) quaternion
        quaternion0 (np.array): (x,y,z,w) quaternion

    Returns:
        np.array: (x,y,z,w) multiplied quaternion
    """
    x0, y0, z0, w0 = quaternion0
    x1, y1, z1, w1 = quaternion1
    return np.array(
        (
            x1 * w0 + y1 * z0 - z1 * y0 + w1 * x0,
            -x1 * z0 + y1 * w0 + z1 * x0 + w1 * y0,
            x1 * y0 - y1 * x0 + z1 * w0 + w1 * z0,
            -x1 * x0 - y1 * y0 - z1 * z0 + w1 * w0,
        ),
        dtype=np.float32,
    )


def quat_conjugate(quaternion):
    """
    Return conjugate of quaternion.

    E.g.:
    >>> q0 = random_quaternion()
    >>> q1 = quat_conjugate(q0)
    >>> q1[3] == q0[3] and all(q1[:3] == -q0[:3])
    True

    Args:
        quaternion (np.array): (x,y,z,w) quaternion

    Returns:
        np.array: (x,y,z,w) quaternion conjugate
    """
    return np.array(
        (-quaternion[0], -quaternion[1], -quaternion[2], quaternion[3]),
        dtype=np.float32,
    )


def quat_inverse(quaternion):
    """
    Return inverse of quaternion.

    E.g.:
    >>> q0 = random_quaternion()
    >>> q1 = quat_inverse(q0)
    >>> np.allclose(quat_multiply(q0, q1), [0, 0, 0, 1])
    True

    Args:
        quaternion (np.array): (x,y,z,w) quaternion

    Returns:
        np.array: (x,y,z,w) quaternion inverse
    """
    return quat_conjugate(quaternion) / np.dot(quaternion, quaternion)


def quat_distance(quaternion1, quaternion0):
    """
    Returns distance between two quaternions, such that distance * quaternion0 = quaternion1

    Args:
        quaternion1 (np.array): (x,y,z,w) quaternion
        quaternion0 (np.array): (x,y,z,w) quaternion

    Returns:
        np.array: (x,y,z,w) quaternion distance
    """
    return quat_multiply(quaternion1, quat_inverse(quaternion0))


def quat_slerp(quat0, quat1, fraction, shortestpath=True):
    """
    Return spherical linear interpolation between two quaternions.

    E.g.:
    >>> q0 = random_quat()
    >>> q1 = random_quat()
    >>> q = quat_slerp(q0, q1, 0.0)
    >>> np.allclose(q, q0)
    True

    >>> q = quat_slerp(q0, q1, 1.0)
    >>> np.allclose(q, q1)
    True

    >>> q = quat_slerp(q0, q1, 0.5)
    >>> angle = math.acos(np.dot(q0, q))
    >>> np.allclose(2.0, math.acos(np.dot(q0, q1)) / angle) or \
        np.allclose(2.0, math.acos(-np.dot(q0, q1)) / angle)
    True

    Args:
        quat0 (np.array): (x,y,z,w) quaternion startpoint
        quat1 (np.array): (x,y,z,w) quaternion endpoint
        fraction (float): fraction of interpolation to calculate
        shortestpath (bool): If True, will calculate the shortest path

    Returns:
        np.array: (x,y,z,w) quaternion distance
    """
    q0 = unit_vector(quat0[:4])
    q1 = unit_vector(quat1[:4])
    if fraction == 0.0:
        return q0
    elif fraction == 1.0:
        return q1
    d = np.dot(q0, q1)
    if abs(abs(d) - 1.0) < EPS:
        return q0
    if shortestpath and d < 0.0:
        # invert rotation
        d = -d
        q1 *= -1.0
    angle = math.acos(np.clip(d, -1, 1))
    if abs(angle) < EPS:
        return q0
    isin = 1.0 / math.sin(angle)
    q0 *= math.sin((1.0 - fraction) * angle) * isin
    q1 *= math.sin(fraction * angle) * isin
    q0 += q1
    return q0


def random_quat(rand=None):
    """
    Return uniform random unit quaternion.

    E.g.:
    >>> q = random_quat()
    >>> np.allclose(1.0, vector_norm(q))
    True
    >>> q = random_quat(np.random.random(3))
    >>> q.shape
    (4,)

    Args:
        rand (3-array or None): If specified, must be three independent random variables that are uniformly distributed
            between 0 and 1.

    Returns:
        np.array: (x,y,z,w) random quaternion
    """
    if rand is None:
        rand = np.random.rand(3)
    else:
        assert len(rand) == 3
    r1 = np.sqrt(1.0 - rand[0])
    r2 = np.sqrt(rand[0])
    pi2 = math.pi * 2.0
    t1 = pi2 * rand[1]
    t2 = pi2 * rand[2]
    return np.array(
        (np.sin(t1) * r1, np.cos(t1) * r1, np.sin(t2) * r2, np.cos(t2) * r2),
        dtype=np.float32,
    )


def random_axis_angle(angle_limit=None, random_state=None):
    """
    Samples an axis-angle rotation by first sampling a random axis
    and then sampling an angle. If @angle_limit is provided, the size
    of the rotation angle is constrained. 

    If @random_state is provided (instance of np.random.RandomState), it
    will be used to generate random numbers.

    Args:
        angle_limit (None or float): If set, determines magnitude limit of angles to generate
        random_state (None or RandomState): RNG to use if specified

    Raises:
        AssertionError: [Invalid RNG]
    """
    if angle_limit is None:
        angle_limit = 2. * np.pi

    if random_state is not None:
        assert isinstance(random_state, np.random.RandomState)
        npr = random_state
    else:
        npr = np.random

    # sample random axis using a normalized sample from spherical Gaussian.
    # see (http://extremelearning.com.au/how-to-generate-uniformly-random-points-on-n-spheres-and-n-balls/)
    # for why it works.
    random_axis = npr.randn(3)
    random_axis /= np.linalg.norm(random_axis)
    random_angle = npr.uniform(low=0., high=angle_limit)
    return random_axis, random_angle


def vec(values):
    """
    Converts value tuple into a numpy vector.

    Args:
        values (n-array): a tuple of numbers

    Returns:
        np.array: vector of given values
    """
    return np.array(values, dtype=np.float32)


def mat4(array):
    """
    Converts an array to 4x4 matrix.

    Args:
        array (n-array): the array in form of vec, list, or tuple

    Returns:
        np.array: a 4x4 numpy matrix
    """
    return np.array(array, dtype=np.float32).reshape((4, 4))


def mat2pose(hmat):
    """
    Converts a homogeneous 4x4 matrix into pose.

    Args:
        hmat (np.array): a 4x4 homogeneous matrix

    Returns:
        2-tuple:

            - (np.array) (x,y,z) position array in cartesian coordinates
            - (np.array) (x,y,z,w) orientation array in quaternion form
    """
    pos = hmat[:3, 3]
    orn = mat2quat(hmat[:3, :3])
    return pos, orn


@jit_decorator
def mat2quat(rmat):
    """
    Converts given rotation matrix to quaternion.

    Args:
        rmat (np.array): 3x3 rotation matrix

    Returns:
        np.array: (x,y,z,w) float quaternion angles
    """
    M = np.asarray(rmat).astype(np.float32)[:3, :3]

    m00 = M[0, 0]
    m01 = M[0, 1]
    m02 = M[0, 2]
    m10 = M[1, 0]
    m11 = M[1, 1]
    m12 = M[1, 2]
    m20 = M[2, 0]
    m21 = M[2, 1]
    m22 = M[2, 2]
    # symmetric matrix K
    K = np.array(
        [
            [m00 - m11 - m22, np.float32(0.0), np.float32(0.0), np.float32(0.0)],
            [m01 + m10, m11 - m00 - m22, np.float32(0.0), np.float32(0.0)],
            [m02 + m20, m12 + m21, m22 - m00 - m11, np.float32(0.0)],
            [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],
        ]
    )
    K /= 3.0
    # quaternion is Eigen vector of K that corresponds to largest eigenvalue
    w, V = np.linalg.eigh(K)
    inds = np.array([3, 0, 1, 2])
    q1 = V[inds, np.argmax(w)]
    if q1[0] < 0.0:
        np.negative(q1, q1)
    inds = np.array([1, 2, 3, 0])
    return q1[inds]


def euler2mat(euler):
    """
    Converts euler angles into rotation matrix form

    Args:
        euler (np.array): (r,p,y) angles

    Returns:
        np.array: 3x3 rotation matrix

    Raises:
        AssertionError: [Invalid input shape]
    """

    euler = np.asarray(euler, dtype=np.float64)
    assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler)

    ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0]
    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
    cc, cs = ci * ck, ci * sk
    sc, ss = si * ck, si * sk

    mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64)
    mat[..., 2, 2] = cj * ck
    mat[..., 2, 1] = sj * sc - cs
    mat[..., 2, 0] = sj * cc + ss
    mat[..., 1, 2] = cj * sk
    mat[..., 1, 1] = sj * ss + cc
    mat[..., 1, 0] = sj * cs - sc
    mat[..., 0, 2] = -sj
    mat[..., 0, 1] = cj * si
    mat[..., 0, 0] = cj * ci
    return mat


def mat2euler(rmat, axes="sxyz"):
    """
    Converts given rotation matrix to euler angles in radian.

    Args:
        rmat (np.array): 3x3 rotation matrix
        axes (str): One of 24 axis sequences as string or encoded tuple (see top of this module)

    Returns:
        np.array: (r,p,y) converted euler angles in radian vec3 float
    """
    try:
        firstaxis, parity, repetition, frame = _AXES2TUPLE[axes.lower()]
    except (AttributeError, KeyError):
        firstaxis, parity, repetition, frame = axes

    i = firstaxis
    j = _NEXT_AXIS[i + parity]
    k = _NEXT_AXIS[i - parity + 1]

    M = np.array(rmat, dtype=np.float32, copy=False)[:3, :3]
    if repetition:
        sy = math.sqrt(M[i, j] * M[i, j] + M[i, k] * M[i, k])
        if sy > EPS:
            ax = math.atan2(M[i, j], M[i, k])
            ay = math.atan2(sy, M[i, i])
            az = math.atan2(M[j, i], -M[k, i])
        else:
            ax = math.atan2(-M[j, k], M[j, j])
            ay = math.atan2(sy, M[i, i])
            az = 0.0
    else:
        cy = math.sqrt(M[i, i] * M[i, i] + M[j, i] * M[j, i])
        if cy > EPS:
            ax = math.atan2(M[k, j], M[k, k])
            ay = math.atan2(-M[k, i], cy)
            az = math.atan2(M[j, i], M[i, i])
        else:
            ax = math.atan2(-M[j, k], M[j, j])
            ay = math.atan2(-M[k, i], cy)
            az = 0.0

    if parity:
        ax, ay, az = -ax, -ay, -az
    if frame:
        ax, az = az, ax
    return vec((ax, ay, az))


def pose2mat(pose):
    """
    Converts pose to homogeneous matrix.

    Args:
        pose (2-tuple): a (pos, orn) tuple where pos is vec3 float cartesian, and
            orn is vec4 float quaternion.

    Returns:
        np.array: 4x4 homogeneous matrix
    """
    homo_pose_mat = np.zeros((4, 4), dtype=np.float32)
    homo_pose_mat[:3, :3] = quat2mat(pose[1])
    homo_pose_mat[:3, 3] = np.array(pose[0], dtype=np.float32)
    homo_pose_mat[3, 3] = 1.
    return homo_pose_mat


@jit_decorator
def quat2mat(quaternion):
    """
    Converts given quaternion to matrix.

    Args:
        quaternion (np.array): (x,y,z,w) vec4 float angles

    Returns:
        np.array: 3x3 rotation matrix
    """
    # awkward semantics for use with numba
    inds = np.array([3, 0, 1, 2])
    q = np.asarray(quaternion).copy().astype(np.float32)[inds]

    n = np.dot(q, q)
    if n < EPS:
        return np.identity(3)
    q *= math.sqrt(2.0 / n)
    q2 = np.outer(q, q)
    return np.array(
        [
            [1.0 - q2[2, 2] - q2[3, 3], q2[1, 2] - q2[3, 0], q2[1, 3] + q2[2, 0]],
            [q2[1, 2] + q2[3, 0], 1.0 - q2[1, 1] - q2[3, 3], q2[2, 3] - q2[1, 0]],
            [q2[1, 3] - q2[2, 0], q2[2, 3] + q2[1, 0], 1.0 - q2[1, 1] - q2[2, 2]],
        ]
    )


def quat2axisangle(quat):
    """
    Converts quaternion to axis-angle format.
    Returns a unit vector direction scaled by its angle in radians.

    Args:
        quat (np.array): (x,y,z,w) vec4 float angles

    Returns:
        np.array: (ax,ay,az) axis-angle exponential coordinates
    """
    # clip quaternion
    if quat[3] > 1.:
        quat[3] = 1.
    elif quat[3] < -1.:
        quat[3] = -1.

    den = np.sqrt(1. - quat[3] * quat[3])
    if math.isclose(den, 0.):
        # This is (close to) a zero degree rotation, immediately return
        return np.zeros(3)

    return (quat[:3] * 2. * math.acos(quat[3])) / den


def axisangle2quat(vec):
    """
    Converts scaled axis-angle to quat.

    Args:
        vec (np.array): (ax,ay,az) axis-angle exponential coordinates

    Returns:
        np.array: (x,y,z,w) vec4 float angles
    """
    # Grab angle
    angle = np.linalg.norm(vec)

    # handle zero-rotation case
    if math.isclose(angle, 0.):
        return np.array([0., 0., 0., 1.])

    # make sure that axis is a unit vector
    axis = vec / angle

    q = np.zeros(4)
    q[3] = np.cos(angle / 2.)
    q[:3] = axis * np.sin(angle / 2.)
    return q


def pose_in_A_to_pose_in_B(pose_A, pose_A_in_B):
    """
    Converts a homogenous matrix corresponding to a point C in frame A
    to a homogenous matrix corresponding to the same point C in frame B.

    Args:
        pose_A (np.array): 4x4 matrix corresponding to the pose of C in frame A
        pose_A_in_B (np.array): 4x4 matrix corresponding to the pose of A in frame B

    Returns:
        np.array: 4x4 matrix corresponding to the pose of C in frame B
    """

    # pose of A in B takes a point in A and transforms it to a point in C.

    # pose of C in B = pose of A in B * pose of C in A
    # take a point in C, transform it to A, then to B
    # T_B^C = T_A^C * T_B^A
    return pose_A_in_B.dot(pose_A)


def pose_inv(pose):
    """
    Computes the inverse of a homogeneous matrix corresponding to the pose of some
    frame B in frame A. The inverse is the pose of frame A in frame B.

    Args:
        pose (np.array): 4x4 matrix for the pose to inverse

    Returns:
        np.array: 4x4 matrix for the inverse pose
    """

    # Note, the inverse of a pose matrix is the following
    # [R t; 0 1]^-1 = [R.T -R.T*t; 0 1]

    # Intuitively, this makes sense.
    # The original pose matrix translates by t, then rotates by R.
    # We just invert the rotation by applying R-1 = R.T, and also translate back.
    # Since we apply translation first before rotation, we need to translate by
    # -t in the original frame, which is -R-1*t in the new frame, and then rotate back by
    # R-1 to align the axis again.

    pose_inv = np.zeros((4, 4))
    pose_inv[:3, :3] = pose[:3, :3].T
    pose_inv[:3, 3] = -pose_inv[:3, :3].dot(pose[:3, 3])
    pose_inv[3, 3] = 1.0
    return pose_inv


def _skew_symmetric_translation(pos_A_in_B):
    """
    Helper function to get a skew symmetric translation matrix for converting quantities
    between frames.

    Args:
        pos_A_in_B (np.array): (x,y,z) position of A in frame B

    Returns:
        np.array: 3x3 skew symmetric translation matrix
    """
    return np.array(
        [
            0.,
            -pos_A_in_B[2],
            pos_A_in_B[1],
            pos_A_in_B[2],
            0.,
            -pos_A_in_B[0],
            -pos_A_in_B[1],
            pos_A_in_B[0],
            0.,
        ]
    ).reshape((3, 3))


def vel_in_A_to_vel_in_B(vel_A, ang_vel_A, pose_A_in_B):
    """
    Converts linear and angular velocity of a point in frame A to the equivalent in frame B.

    Args:
        vel_A (np.array): (vx,vy,vz) linear velocity in A
        ang_vel_A (np.array): (wx,wy,wz) angular velocity in A
        pose_A_in_B (np.array): 4x4 matrix corresponding to the pose of A in frame B

    Returns:
        2-tuple:

            - (np.array) (vx,vy,vz) linear velocities in frame B
            - (np.array) (wx,wy,wz) angular velocities in frame B
    """
    pos_A_in_B = pose_A_in_B[:3, 3]
    rot_A_in_B = pose_A_in_B[:3, :3]
    skew_symm = _skew_symmetric_translation(pos_A_in_B)
    vel_B = rot_A_in_B.dot(vel_A) + skew_symm.dot(rot_A_in_B.dot(ang_vel_A))
    ang_vel_B = rot_A_in_B.dot(ang_vel_A)
    return vel_B, ang_vel_B


def force_in_A_to_force_in_B(force_A, torque_A, pose_A_in_B):
    """
    Converts linear and rotational force at a point in frame A to the equivalent in frame B.

    Args:
        force_A (np.array): (fx,fy,fz) linear force in A
        torque_A (np.array): (tx,ty,tz) rotational force (moment) in A
        pose_A_in_B (np.array): 4x4 matrix corresponding to the pose of A in frame B

    Returns:
        2-tuple:

            - (np.array) (fx,fy,fz) linear forces in frame B
            - (np.array) (tx,ty,tz) moments in frame B
    """
    pos_A_in_B = pose_A_in_B[:3, 3]
    rot_A_in_B = pose_A_in_B[:3, :3]
    skew_symm = _skew_symmetric_translation(pos_A_in_B)
    force_B = rot_A_in_B.T.dot(force_A)
    torque_B = -rot_A_in_B.T.dot(skew_symm.dot(force_A)) + rot_A_in_B.T.dot(torque_A)
    return force_B, torque_B


def rotation_matrix(angle, direction, point=None):
    """
    Returns matrix to rotate about axis defined by point and direction.

    E.g.:
        >>> angle = (random.random() - 0.5) * (2*math.pi)
        >>> direc = numpy.random.random(3) - 0.5
        >>> point = numpy.random.random(3) - 0.5
        >>> R0 = rotation_matrix(angle, direc, point)
        >>> R1 = rotation_matrix(angle-2*math.pi, direc, point)
        >>> is_same_transform(R0, R1)
        True

        >>> R0 = rotation_matrix(angle, direc, point)
        >>> R1 = rotation_matrix(-angle, -direc, point)
        >>> is_same_transform(R0, R1)
        True

        >>> I = numpy.identity(4, numpy.float32)
        >>> numpy.allclose(I, rotation_matrix(math.pi*2, direc))
        True

        >>> numpy.allclose(2., numpy.trace(rotation_matrix(math.pi/2,
        ...                                                direc, point)))
        True

    Args:
        angle (float): Magnitude of rotation
        direction (np.array): (ax,ay,az) axis about which to rotate
        point (None or np.array): If specified, is the (x,y,z) point about which the rotation will occur

    Returns:
        np.array: 4x4 homogeneous matrix that includes the desired rotation
    """
    sina = math.sin(angle)
    cosa = math.cos(angle)
    direction = unit_vector(direction[:3])
    # rotation matrix around unit vector
    R = np.array(
        ((cosa, 0.0, 0.0), (0.0, cosa, 0.0), (0.0, 0.0, cosa)), dtype=np.float32
    )
    R += np.outer(direction, direction) * (1.0 - cosa)
    direction *= sina
    R += np.array(
        (
            (0.0, -direction[2], direction[1]),
            (direction[2], 0.0, -direction[0]),
            (-direction[1], direction[0], 0.0),
        ),
        dtype=np.float32,
    )
    M = np.identity(4)
    M[:3, :3] = R
    if point is not None:
        # rotation not around origin
        point = np.array(point[:3], dtype=np.float32, copy=False)
        M[:3, 3] = point - np.dot(R, point)
    return M


def clip_translation(dpos, limit):
    """
    Limits a translation (delta position) to a specified limit

    Scales down the norm of the dpos to 'limit' if norm(dpos) > limit, else returns immediately

    Args:
        dpos (n-array): n-dim Translation being clipped (e,g.: (x, y, z)) -- numpy array
        limit (float): Value to limit translation by -- magnitude (scalar, in same units as input)

    Returns:
        2-tuple:

            - (np.array) Clipped translation (same dimension as inputs)
            - (bool) whether the value was clipped or not
    """
    input_norm = np.linalg.norm(dpos)
    return (dpos * limit / input_norm, True) if input_norm > limit else (dpos, False)


def clip_rotation(quat, limit):
    """
    Limits a (delta) rotation to a specified limit

    Converts rotation to axis-angle, clips, then re-converts back into quaternion

    Args:
        quat (np.array): (x,y,z,w) rotation being clipped
        limit (float): Value to limit rotation by -- magnitude (scalar, in radians)

    Returns:
        2-tuple:

            - (np.array) Clipped rotation quaternion (x, y, z, w)
            - (bool) whether the value was clipped or not
    """
    clipped = False

    # First, normalize the quaternion
    quat = quat / np.linalg.norm(quat)

    den = np.sqrt(max(1 - quat[3] * quat[3], 0))
    if den == 0:
        # This is a zero degree rotation, immediately return
        return quat, clipped
    else:
        # This is all other cases
        x = quat[0] / den
        y = quat[1] / den
        z = quat[2] / den
        a = 2 * math.acos(quat[3])

    # Clip rotation if necessary and return clipped quat
    if abs(a) > limit:
        a = limit * np.sign(a) / 2
        sa = math.sin(a)
        ca = math.cos(a)
        quat = np.array([
            x * sa,
            y * sa,
            z * sa,
            ca
        ])
        clipped = True

    return quat, clipped


def make_pose(translation, rotation):
    """
    Makes a homogeneous pose matrix from a translation vector and a rotation matrix.

    Args:
        translation (np.array): (x,y,z) translation value
        rotation (np.array): a 3x3 matrix representing rotation

    Returns:
        pose (np.array): a 4x4 homogeneous matrix
    """
    pose = np.zeros((4, 4))
    pose[:3, :3] = rotation
    pose[:3, 3] = translation
    pose[3, 3] = 1.0
    return pose


def unit_vector(data, axis=None, out=None):
    """
    Returns ndarray normalized by length, i.e. eucledian norm, along axis.

    E.g.:
        >>> v0 = numpy.random.random(3)
        >>> v1 = unit_vector(v0)
        >>> numpy.allclose(v1, v0 / numpy.linalg.norm(v0))
        True

        >>> v0 = numpy.random.rand(5, 4, 3)
        >>> v1 = unit_vector(v0, axis=-1)
        >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=2)), 2)
        >>> numpy.allclose(v1, v2)
        True

        >>> v1 = unit_vector(v0, axis=1)
        >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=1)), 1)
        >>> numpy.allclose(v1, v2)
        True

        >>> v1 = numpy.empty((5, 4, 3), dtype=numpy.float32)
        >>> unit_vector(v0, axis=1, out=v1)
        >>> numpy.allclose(v1, v2)
        True

        >>> list(unit_vector([]))
        []

        >>> list(unit_vector([1.0]))
        [1.0]

    Args:
        data (np.array): data to normalize
        axis (None or int): If specified, determines specific axis along data to normalize
        out (None or np.array): If specified, will store computation in this variable

    Returns:
        None or np.array: If @out is not specified, will return normalized vector. Otherwise, stores the output in @out
    """
    if out is None:
        data = np.array(data, dtype=np.float32, copy=True)
        if data.ndim == 1:
            data /= math.sqrt(np.dot(data, data))
            return data
    else:
        if out is not data:
            out[:] = np.array(data, copy=False)
        data = out
    length = np.atleast_1d(np.sum(data * data, axis))
    np.sqrt(length, length)
    if axis is not None:
        length = np.expand_dims(length, axis)
    data /= length
    if out is None:
        return data


def get_orientation_error(target_orn, current_orn):
    """
    Returns the difference between two quaternion orientations as a 3 DOF numpy array.
    For use in an impedance controller / task-space PD controller.

    Args:
        target_orn (np.array): (x, y, z, w) desired quaternion orientation
        current_orn (np.array): (x, y, z, w) current quaternion orientation

    Returns:
        orn_error (np.array): (ax,ay,az) current orientation error, corresponds to
            (target_orn - current_orn)
    """
    current_orn = np.array(
        [current_orn[3], current_orn[0], current_orn[1], current_orn[2]]
    )
    target_orn = np.array([target_orn[3], target_orn[0], target_orn[1], target_orn[2]])

    pinv = np.zeros((3, 4))
    pinv[0, :] = [-current_orn[1], current_orn[0], -current_orn[3], current_orn[2]]
    pinv[1, :] = [-current_orn[2], current_orn[3], current_orn[0], -current_orn[1]]
    pinv[2, :] = [-current_orn[3], -current_orn[2], current_orn[1], current_orn[0]]
    orn_error = 2.0 * pinv.dot(np.array(target_orn))
    return orn_error


def get_pose_error(target_pose, current_pose):
    """
    Computes the error corresponding to target pose - current pose as a 6-dim vector.
    The first 3 components correspond to translational error while the last 3 components
    correspond to the rotational error.

    Args:
        target_pose (np.array): a 4x4 homogenous matrix for the target pose
        current_pose (np.array): a 4x4 homogenous matrix for the current pose

    Returns:
        np.array: 6-dim pose error.
    """
    error = np.zeros(6)

    # compute translational error
    target_pos = target_pose[:3, 3]
    current_pos = current_pose[:3, 3]
    pos_err = target_pos - current_pos

    # compute rotational error
    r1 = current_pose[:3, 0]
    r2 = current_pose[:3, 1]
    r3 = current_pose[:3, 2]
    r1d = target_pose[:3, 0]
    r2d = target_pose[:3, 1]
    r3d = target_pose[:3, 2]
    rot_err = 0.5 * (np.cross(r1, r1d) + np.cross(r2, r2d) + np.cross(r3, r3d))

    error[:3] = pos_err
    error[3:] = rot_err
    return error


@jit_decorator
def matrix_inverse(matrix):
    """
    Helper function to have an efficient matrix inversion function.

    Args:
        matrix (np.array): 2d-array representing a matrix

    Returns:
        np.array: 2d-array representing the matrix inverse
    """
    return np.linalg.inv(matrix)

hacman/utils/robosuite_transform_utils.py
HACMan-2023/HACMan
import numpy as np
import torch
import hacman.utils.robosuite_transform_utils as transform_utils
import scipy.spatial.transform as spt
from transforms3d import quaternions, euler
# from torchgeometry.core.conversions import rotation_matrix_to_quaternion
from pytorch3d.transforms import matrix_to_quaternion

def decompose_pose_tensor(pose: torch.Tensor, cat=False):
    # Note: It handles torch tensors rather than ndarray
    pos = pose[:, :3, 3]
    quat = matrix_to_quaternion(pose[:, :3, :3])
    quat = quat[:, (1, 2, 3, 0)] # Move w to the last to match with robosuite

    # Verification only (uncomment this to debug)
    # test_p = pose[0].detach().cpu().numpy()
    # test_pos, test_quat = transform_utils.mat2pose(test_p) # This does not work on torch Tensors

    if cat:
        return torch.cat([pos, quat], dim=1)
    else:
        return pos, quat

def decompose_pose_mat(pose, cat=False, wxyz=True):
    pos, quat = transform_utils.mat2pose(pose)
    if wxyz:
        quat = transform_utils.convert_quat(quat, to="wxyz")
    if cat:
        return np.concatenate([pos, quat])
    else:
        return pos, quat


def to_pose_mat(pos, quat, input_wxyz=True):
    """
    Convert position and quaternion to a pose matrix.

    Args:
        pos (ndarray): Position vector of shape (3,) representing the translation.
        quat (ndarray): Quaternion vector of shape (4,) representing the orientation.
        input_wxyz (bool): Indicates whether the input quaternion follows the wxyz convention (True) or xyzw convention (False).

    Returns:
        ndarray: Pose matrix of shape (4, 4) representing the pose.
    """
    if not input_wxyz:
        quat = quat[[3, 0, 1, 2]]  # Convert wxyz to xyzw convention

    mat = quaternions.quat2mat(quat)
    pose = np.eye(4)
    pose[:3, :3] = mat
    pose[:3, 3] = pos
    return pose



def inv_pose_mat(pose):
    """
    Compute the inverse of a pose matrix.

    Args:
        pose (ndarray): Pose matrix of shape (4, 4).

    Returns:
        ndarray: Inverse pose matrix of shape (4, 4).
    """
    pose_inv = np.zeros((4, 4))
    pose_inv[:3, :3] = pose[:3, :3].T
    pose_inv[:3, 3] = -pose_inv[:3, :3].dot(pose[:3, 3])
    pose_inv[3, 3] = 1.0
    return pose_inv


def transform_point_cloud(current_pose_mat, next_pose_mat, pointcloud):
    """
    Transform a point cloud from the current pose to the next pose.

    This function takes the current and next pose matrices along with the point cloud and performs the transformation.
    The point cloud can be either batched or non-batched.

    Args:
        current_pose_mat (ndarray): Current pose matrix of shape (batch_size, 4, 4) or (4, 4).
        next_pose_mat (ndarray): Next pose matrix of shape (batch_size, 4, 4) or (4, 4).
        pointcloud (ndarray): Point cloud of shape (batch_size, num_points, 3) or (num_points, 3).

    Returns:
        ndarray: Transformed point cloud of shape (batch_size, num_points, 3) or (num_points, 3).
    """
    # Expand dimensions if input is non-batched
    if current_pose_mat.ndim == 2:
        current_pose_mat = np.expand_dims(current_pose_mat, axis=0)
        next_pose_mat = np.expand_dims(next_pose_mat, axis=0)
        pointcloud = np.expand_dims(pointcloud, axis=0)

    bs = current_pose_mat.shape[0]
    transformed_pcds = []
    for i in range(bs):
        transform = np.dot(next_pose_mat[i], inv_pose_mat(current_pose_mat[i]))

        transformed_pcd = np.concatenate([pointcloud[i], np.ones((len(pointcloud[i]), 1))], axis=1)
        transformed_pcd = np.dot(transform, transformed_pcd.T).T[:, :3]
        transformed_pcds.append(transformed_pcd)

    return np.array(transformed_pcds).squeeze()


def sample_idx(old_length, new_length):
    """
    Sample indices from the range [0, old_length) to obtain a new list of indices of length new_length.

    Args:
        old_length (int): The length of the original index range.
        new_length (int): The desired length of the sampled index list.

    Returns:
        ndarray: The list of sampled indices of length new_length.
    """
    if old_length > new_length:
        idx = np.random.choice(old_length, new_length, replace=False)
    elif old_length < new_length:
        idx = np.arange(new_length)
        idx[old_length:] = np.random.choice(old_length, new_length - old_length, replace=True)
    else:
        idx = np.arange(old_length)
        
    return idx
    
hacman/utils/transformations.py
HACMan-2023/HACMan
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import os
from hacman.utils.launch_utils import use_freer_gpu
use_freer_gpu()
# print(f"Using CUDA Device {os.environ['CUDA_VISIBLE_DEVICES']}")

from stable_baselines3.common.callbacks import EvalCallback, CallbackList

import wandb
from hacman.envs.wandb_wrappers import WandbPointCloudRecorder
from hacman.sb3_utils.custom_callbacks import ProgressBarCallback, CustomizedCheckpointCallback
from hacman.sb3_utils.evaluation import evaluate_policy
from hacman.algos.setup_model import setup_model, add_model_config
from hacman.envs.setup_envs import setup_envs, add_env_config

import argparse
import torch_geometric

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Exp and logging config
    parser.add_argument("--ExpID", default=9999, type=int, help="Exp ID")
    parser.add_argument("--name", default="tmp", type=str, help="Exp Name")
    parser.add_argument("--seed", default=0, type=int, help="Random seed")
    parser.add_argument("--eval", default=None, type=int, help="Eval only. Eval epsiodes to run")
    parser.add_argument("--dirname", default=None, type=str, help="Path to save models")
    parser.add_argument("--debug", action="store_true", help="Use debug config")

    parser.add_argument("--n_eval_episodes", default=20, type=int, help="Number of eval episodes")
    parser.add_argument("--save_freq_latest", default=100, type=int, help="Save per n env steps")
    parser.add_argument("--save_freq_checkpoint", default=1000, type=int, help="Save per n env steps")
    parser.add_argument("--save_replay_buffer", action="store_true", help="Save buffer")
    parser.add_argument("--record_video", action="store_true", help="Record video.")
    parser.add_argument("--upload_video", action="store_true", help="Upload video to wandb.")

    add_model_config(parser)
    add_env_config(parser)
    
    args = parser.parse_args()
    config = vars(args)

    if args.dirname is None or args.debug:
        dirname = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results')
    else:
        dirname = args.dirname
    config['dirname'] = dirname
    config["fullname"] = f"Exp{args.ExpID:04d}-{args.name}"
    
    torch_geometric.seed_everything(config["seed"])
    print(config)

    # ------------- Logging ------------- 
    result_folder = os.path.join(config["dirname"], f'{config["fullname"]}-{config["seed"]}')
    if not os.path.exists(result_folder):
        os.makedirs(result_folder)

    run_id = wandb.util.generate_id()
    run = wandb.init(# project="gc-reorientation", entity="bwww", 
                     name=config["fullname"], config=config, id=run_id, 
                     dir=result_folder, sync_tensorboard=True)
    print(f'wandb run-id:{run.id}')
    print(f'Result folder: {result_folder}')
    
    # ------------- Task Setup -------------
    if config['eval'] is not None:
        eval_env = setup_envs(config, eval_env_only=True)
        model = setup_model(config, eval_env, result_folder, normalize_param=None)
        
        if hasattr(eval_env, "location_model"):
            eval_env.location_model.load_model(model)

        model.policy.eval()

    else:
        env, eval_env = setup_envs(config)
        model = setup_model(config, env, result_folder, normalize_param=None)
        
        if hasattr(env, "location_model") and hasattr(eval_env, "location_model"):
            env.location_model.load_model(model)
            eval_env.location_model.load_model(model)
            
    # ------------- Additional Setup for Logging Purposes ------------- 
    def global_step_counter():
        return model.num_timesteps
    eval_env = WandbPointCloudRecorder(eval_env, 
        global_step_counter=global_step_counter, 
        save_plotly=True, foldername=result_folder,
        record_video=config['record_video'],
        upload_video=config['upload_video'])
    if config['record_video']:
            eval_env.enable_video_recording()
    
    eval_callback = EvalCallback(
        eval_env, 
        n_eval_episodes=config['n_eval_episodes'], 
        eval_freq=config['eval_freq'],
        best_model_save_path=os.path.join(result_folder, "model"))
    ckpt_callback = CustomizedCheckpointCallback(
        save_freq_latest=config['save_freq_latest'],
        save_freq_checkpoint=config['save_freq_checkpoint'],
        save_path=os.path.join(result_folder, f"model-{run.id}"),
        save_replay_buffer=config['save_replay_buffer'],
        save_vecnormalize=True)  # TODO: confirm?
    progress_bar_callback = ProgressBarCallback()
    callback_list = CallbackList([ckpt_callback, eval_callback, progress_bar_callback])
    
    # ------------- Start Training -------------
    if config['eval'] is None:
        model.learn(
            total_timesteps=1000000,
            log_interval=10,
            callback=callback_list)

    # ------------- Eval -------------
    else:
        model.policy.eval()

        from tqdm import tqdm
        import numpy as np
        import pandas as pd

        pbar = tqdm(total=config['eval'])
        mean_reward, std_reward, succ, verbose_buffer = evaluate_policy(
            model, eval_env, n_eval_episodes=config['eval'], deterministic=True,
            save_path=os.path.join(result_folder, f'obs_list_{0}.pkl'), verbose=True, pbar=pbar)
        
        uncertainty = 1.96 * np.sqrt(succ * (1 - succ) / config['eval'])
        print(f"succ={succ:.3f} +/- {uncertainty:.3f}. mean_reward={mean_reward:.2f} +/- {std_reward}")

        # Export the verbose buffer
        eval_dir = os.path.join(result_folder, 'evals')
        if not os.path.exists(eval_dir):
            os.makedirs(eval_dir)
        import pickle
        with open(os.path.join(eval_dir, f'verbose_eval.pkl'), 'wb') as f:
            pickle.dump(verbose_buffer, f)
scripts/run.py
HACMan-2023/HACMan
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

import open3d as o3d

def test_simple_cube_sim():
    # Import the SimpleCubeSim class and create an instance
    from hacman.envs.simple_env import SimpleCubeSim
    sim = SimpleCubeSim()
    sim.reset(np.zeros(2))

    # Get the point cloud and segmentation before the poke
    point_cloud_before, segmentation_before = sim.get_point_cloud()

    # Perform a poke on the cube
    location = [0.1, 0.1, 0.1]  # Example poke location
    motion = [-0.2, -0.4, 0.0]   # Example poke motion
    sim.poke(location, motion)

    # Get the point cloud and segmentation after the poke
    point_cloud_after, segmentation_after = sim.get_point_cloud()

    # Visualize the point clouds with segmentation color
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    cmap = ListedColormap(['#1f77b4', '#f2f2f2'])

    # Before poke
    ax1.scatter(point_cloud_before[:, 0], point_cloud_before[:, 1], c=segmentation_before, cmap=cmap)
    ax1.scatter(location[0], location[1], c='C2', marker='o', label='Contact Location')
    ax1.scatter(location[0]+motion[0], location[1]+motion[1], c='C1', marker='.', label='__nolegend__')
    ax1.quiver(location[0], location[1], motion[0], motion[1], color='C1', label='Motion')
    ax1.set_title('Point Cloud Before Poke')
    ax1.set_xlabel('X')
    ax1.set_ylabel('Y')
    ax1.legend()

    # After poke
    ax2.scatter(point_cloud_after[:, 0], point_cloud_after[:, 1], c=segmentation_after, cmap=cmap)
    ax2.set_title('Point Cloud After Poke')
    ax2.set_xlabel('X')
    ax2.set_ylabel('Y')

    plt.show()


def test_simple_env_2d(trials=10):
    # Import the SimpleEnv class
    from hacman.envs.sim_envs.simple_env import SimpleEnv

    # Create a SimpleEnv object
    env = SimpleEnv()

    for i in range(trials):
        # Reset the environment
        obs = env.reset()

        # Store initial observation
        initial_object_pcd_points = obs['object_pcd_points']
        initial_background_pcd_points = obs['background_pcd_points']

        # Perform a step in the environment
        action = np.random.random(3) * 2 - 1
        next_obs, reward, done, info = env.step(action)

        # Store updated observation
        updated_object_pcd_points = next_obs['object_pcd_points']
        updated_background_pcd_points = next_obs['background_pcd_points']

        # Create matplotlib figure and axes
        plt.subplots(figsize=(7, 7))
        plt.xlabel("X")
        plt.ylabel("Y")

        # Plot points
        plt.scatter(initial_background_pcd_points[:, 0], initial_background_pcd_points[:, 1], color='lightgrey', label="Previous Background Points")
        plt.scatter(updated_background_pcd_points[:, 0], updated_background_pcd_points[:, 1], color='grey', label="Current Background Points")
        plt.scatter(initial_object_pcd_points[:, 0], initial_object_pcd_points[:, 1], color='orange', label="Previous Object Points")
        plt.scatter(updated_object_pcd_points[:, 0], updated_object_pcd_points[:, 1], color='blue', label="Current Object Points")
        

        # Plot the action line
        action_location = np.round(info["action_location"], 2)
        action_param = np.round(info["action_param"], 2)
        action_endpoint = action_location + action_param
        plt.plot([action_location[0], action_endpoint[0]], [action_location[1], action_endpoint[1]], color='red', label="Action")

        # Show the plot
        plt.title(f"Point Cloud Visualization \n Location: {action_location} \n Param: {action_param} \n Reward: {reward}")
        plt.legend()
        plt.show()

def test_simple_env_3d(trials=10):
    # Create a SimpleEnv object
    from hacman.envs.sim_envs.simple_env import SimpleEnv
    env = SimpleEnv()

    for _ in range(trials):
        # Reset the environment
        obs = env.reset()

        # Store initial observation
        initial_object_pcd_points = obs['object_pcd_points']
        initial_background_pcd_points = obs['background_pcd_points']

        # Create Open3D point cloud objects for initial observation
        initial_object_pcd = o3d.geometry.PointCloud()
        initial_object_pcd.points = o3d.utility.Vector3dVector(initial_object_pcd_points)
        initial_object_pcd.paint_uniform_color([1.0, 0.65, 0.0])  # Orange color for previous object points

        initial_background_pcd = o3d.geometry.PointCloud()
        initial_background_pcd.points = o3d.utility.Vector3dVector(initial_background_pcd_points)
        initial_background_pcd.paint_uniform_color([0.9, 0.9, 0.9])  # Light grey color for previous background points

        # Visualize the initial point clouds using Open3D
        visualizer = o3d.visualization.Visualizer()
        visualizer.create_window()
        visualizer.add_geometry(initial_background_pcd)
        visualizer.add_geometry(initial_object_pcd)

        # Perform a step in the environment
        action = np.random.random(3) * 2 - 1
        next_obs, reward, done, info = env.step(action)
        print("Reward: ", reward)

        # Store updated observation
        updated_object_pcd_points = next_obs['object_pcd_points']
        updated_background_pcd_points = next_obs['background_pcd_points']

        # Create Open3D point cloud objects for updated observation
        updated_object_pcd = o3d.geometry.PointCloud()
        updated_object_pcd.points = o3d.utility.Vector3dVector(updated_object_pcd_points)
        updated_object_pcd.paint_uniform_color([0.0, 0.0, 1.0])  # Blue color for current object points

        updated_background_pcd = o3d.geometry.PointCloud()
        updated_background_pcd.points = o3d.utility.Vector3dVector(updated_background_pcd_points)
        updated_background_pcd.paint_uniform_color([0.7, 0.7, 0.7])  # Grey color for current background points

        # Create an Open3D line segment to represent the action
        action_location = info["action_location"]
        action_param = info["action_param"]
        action_endpoint = action_location + action_param

        line = o3d.geometry.LineSet()
        line.points = o3d.utility.Vector3dVector([action_location, action_endpoint])
        line.lines = o3d.utility.Vector2iVector([[0, 1]])
        line.colors = o3d.utility.Vector3dVector([[1.0, 0.0, 0.0]])  # Red color for the action line

        # Visualize the updated point clouds and action line using Open3D
        visualizer.add_geometry(updated_background_pcd)
        visualizer.add_geometry(updated_object_pcd)
        visualizer.add_geometry(line)
        visualizer.run()
        visualizer.destroy_window()

if __name__ == '__main__':
    # test_simple_cube_sim()
    test_simple_env_2d()
    # test_simple_env_3d()
scripts/test_simple_env.py
HACMan-2023/HACMan
from distutils.core import setup

setup(
    name='HACMan',
    version='0.1dev',
    packages=['hacman','bin_env'],
    license='MIT License',
    long_description='Learning Hybrid Actor-Critic Maps for 6D Non-Prehensile Manipulation',
)
setup.py
HACMan-2023/HACMan