<jupyter_start><jupyter_text>Generate Solvable Goals<jupyter_code>def generateTool(tool_template):
    objects=list()
    for part in tool_template['main']:
        obj_type = PrimitiveShape.CUBOID
        obj= Shape.create(type=obj_type,
                            size=[part[0], part[1], 0.05],
                            color=colors[3],
                            static=True, respondable=True)
        pos = [center_robot[0]+part[2],center_robot[1]+part[3], 0.90]
        obj.set_position(pos)
        obj.set_orientation([0,0,part[4]])
        objects.append(obj)
    return objects
import collision
def check_collision(obj_info, obj_info2):
    # Object 1    
    if obj_info['type']==0:
        center_object = collision.Vector(obj_info['position'][0],obj_info['position'][1])
        obj1 =collision.Circle(center_object, (obj_info['shape'][0])/2+0.05)
    elif obj_info['type']==1:
        center_object = collision.Vector(obj_info['position'][0],obj_info['position'][1])
        obj1 =collision.Poly.from_box(center_object, obj_info['shape'][0]+0.05,obj_info['shape'][1]+0.05)
        obj1.angle = obj_info['angle']

    # Object 2

    center_object = collision.Vector(obj_info2['position'][0],obj_info2['position'][1])
    obj2 =collision.Circle(center_object, obj_info2['shape'][0]/2)

    return collision.collide(obj1,obj2)
from os.path import dirname, join, abspath
from pyrep import PyRep
from pyrep.robots.arms.ur10 import UR10
from pyrep.objects.shape import Shape
from pyrep.objects.dummy import Dummy
from pyrep.objects.joint import Joint
from pyrep.backend import sim
from pyrep.const import PrimitiveShape,JointType
from pyrep.const import ConfigurationPathAlgorithms as Algos
from pyrep.errors import ConfigurationPathError
from pyrep.backend.sim import simGetObjectPosition,simGetObjectHandle
import numpy as np
import math
import copy
import time
import signal
from multiprocessing import Process, Value, Lock

colors=dict()
colors[0]=[0.2,0.2,0.2]
colors[1]=[0,1,0]
colors[2]=[0,0,1]
colors[3]=[1,1,0]
colors[4]=[0,1,1]
colors[5]=[1,0,0]
colors[6]=[0,0.5,0]
colors[7]=[0,0,0.5]
colors[8]=[0.5,0.5,0]
colors[9]=[0,0.5,0.5]
colors[10]=[0.5,0,0.5]
colors[11]=[0.5,1,0]
colors[12]=[0,0.5,1]
colors[13]=[1,1,0.5]
colors[14]=[0.5,1,1]
colors[15]=[1,0.5,1]
import os
import pickle
import json
import random

SCENE_FILE = 'sim_files/empty_ode.ttt'
pr = PyRep()
pr.launch(SCENE_FILE, headless=True)
pr.set_simulation_timestep(0.05)
pr.start()    
### For video recording in headless = False
ur10_model = pr.import_model('sim_files/ur10.ttm')
table_model = pr.import_model('sim_files/table.ttm')
# Reposition Simulation View
for _ in range(1000):
    pr.step()
    
ur10_model.remove()
table_model.remove()
### Draw Initialization Example ###
colors=dict()
colors[0]=[0.2,0.2,0.2]
colors[1]=[0,0,1]
colors[2]=[0,1,0]
colors[3]=[1,1,0]
colors[4]=[0,1,1]
colors[5]=[1,0,1]
colors[6]=[0,0.5,0]
colors[7]=[0,0,0.5]
colors[8]=[0.5,0.5,0]
colors[9]=[0,0.5,0.5]
colors[10]=[0.5,0,0.5]
colors[11]=[0.5,1,0]
colors[12]=[0,0.5,1]
colors[13]=[1,1,0.5]
colors[14]=[0.5,1,1]
colors[15]=[1,0.5,1]
tool_names = list(tools.keys())

generation_area_center =np.array([-0.425,0])
generation_area=np.array([0.6,1])

start = generation_area_center-generation_area/2
end = generation_area_center+generation_area/2
xx = np.linspace(start[0],end[0],4)
yy = np.linspace(start[1],end[1],6)
obj_posses = np.transpose([np.tile(xx, len(yy)), np.repeat(yy, len(xx))])

for tool_index in range(1):
    fig, ax = plt.subplots(figsize=(14,10))

    tool_name = tool_names[tool_index]
    tool = tools[tool_name]
    print(tool_name)
    patch = patches.Rectangle((generation_area_center[0]-generation_area[0]/2,
                           generation_area_center[1]-generation_area[1]/2),
                          generation_area[0], generation_area[1],ec='blue', fc='blue',alpha=0.25)
    ax.add_patch(patch)    
    for part in tool['main']:
        obj_info2 = dict()

        obj_info2['type'] = 1
        obj_info2['shape'] = [part[0],part[1]]
        obj_info2['position'] = [np.cos(tool['ori'])*part[2]+np.sin(tool['ori'])*part[3],
                                 np.cos(tool['ori'])*part[3]+np.sin(tool['ori'])*part[2]]
        obj_info2['angle'] = part[4] + tool['ori']
        c=colors[2]
        for k in [-1,1]:
            for l in [-1,1]:
                patch = patches.Rectangle((obj_info2['position'][0],obj_info2['position'][1]), k*obj_info2['shape'][0]/2, l*obj_info2['shape'][1]/2,ec='black', fc=c)
                patch.angle=np.rad2deg(obj_info2['angle'])     
                patch.set_hatch('///')
                ax.add_patch(patch)

    for init_position_ind in range(obj_posses.shape[0]): ## For each init position
        init_position = obj_posses[init_position_ind]
        obj_info = dict()

        obj_info['type'] = 0
        obj_info['shape'] = [0.10 , 0.10]
        obj_info['position'] = [init_position[0], init_position[1]]
        obj_info['angle'] = 0    
        circle = patches.Circle(obj_info['position'], 0.05, color=colors[3],ec='black')
        ax.add_patch(circle)
    ax.set_xlim(0.5,-1.1)
    ax.set_ylim(-0.75,0.75)        
    plt.show()        
tool_names = list(tools.keys())

push_dirs = [(1,0),(-1,0),(0,1),(0,-1)]
action_list=list()
trial = 0
for i in range(4):
    for j in range(4):
        for k in range(4):
            action_list.append((i,j,k))
            
generation_area_center =np.array([-0.425,0])
generation_area=np.array([0.6,1])

start = generation_area_center-generation_area/2
end = generation_area_center+generation_area/2
xx = np.linspace(start[0],end[0],4)
yy = np.linspace(start[1],end[1],6)
obj_posses = np.transpose([np.tile(xx, len(yy)), np.repeat(yy, len(xx))])

position_pairs = dict()
pbar = tqdm(total=24*7*64)
for init_position_ind in range(24): ## For each init position

    init_position = obj_posses[init_position_ind]

    position_pairs[tuple(init_position.tolist())]=list()
    obj_info = dict()

    obj_info['type'] = 0
    obj_info['shape'] = [0.10 , 0.10]
    obj_info['position'] = [init_position[0], init_position[1]]
    obj_info['angle'] = 0    

    for tool_index in range(7):
        # Check if tool is viable for this position
        isCollision=False

        tool_name = tool_names[tool_index]
        tool = tools[tool_name]

        for part in tool['main']:
            obj_info2 = dict()

            obj_info2['type'] = 1
            obj_info2['shape'] = [part[0],part[1]]
            obj_info2['position'] = [np.cos(tool['ori'])*part[2]+np.sin(tool['ori'])*part[3],
                                     np.cos(tool['ori'])*part[3]+np.sin(tool['ori'])*part[2]]
            obj_info2['angle'] = part[4] + tool['ori']

            isCollision = isCollision or check_collision(obj_info2, obj_info)
        if isCollision:
            pbar.update(64)
            continue

        for action_ind, action in enumerate(action_list):
            ur10_model = pr.import_model('sim_files/ur10.ttm')
            table_model = pr.import_model('sim_files/table.ttm')
            for _ in range(10):
                pr.step()
            # Start Robot
            agent = UR10()
            agent.max_velocity=0.10
            agent.max_acceleration=1
            agent.max_jerk=10
            agent_hand = Shape(simGetObjectHandle('UR10_hand'))
            try:
                path = agent.get_path( #[-math.pi,0,math.pi]
                    position=[center_robot[0],center_robot[1],0.90],euler=[-math.pi,0,0],ignore_collisions=True)
            except ConfigurationPathError as _:
                pass
            path.set_to_end()
            pr.step()
            for _ in range(10):
                pr.step()
            # Sample Object Start Position

            start_position_x = center_robot[0] +init_position[0]
            start_position_y = center_robot[1] +init_position[1]
            obj_start_position = (start_position_x,start_position_y)

            objects= generateTool(tool)
            objects.insert(0,agent_hand)
            grouped_shape= pr.group_objects(objects)
            sim.simSetObjectSpecialProperty(grouped_shape.get_handle(),1)
            sim.simSetObjectParent(grouped_shape.get_handle(),simGetObjectHandle('UR10_connection'),True)
            ## PUSH Action
            start_pos = [center_robot[0],center_robot[1],0.75]
            start_ori = tool['ori']
            ## Repositioning STEP 1
            step_positions=[
                start_pos.copy(),
                start_pos
            ]
            step_positions[0][2]=0.95
            for step in range(2):
                try:
                    path = agent.get_path(
                        position=step_positions[step],euler=[-math.pi,0,start_ori],ignore_collisions=True)
                except ConfigurationPathError as _:
                    print('Could not find path')
                    continue
                if step==0:
                    path.set_to_end()
                    pr.step()
                else:
                    done=False
                    while not done:
                        done = path.step()
                        pr.step()
            ###PUT OBJECT###

            obj_type = PrimitiveShape.CYLINDER
            obj1= Shape.create(type=obj_type,
                                size=[0.10, 0.10, 0.10],
                                color=colors[4],
                                static=False, respondable=True)
            pos = [obj_start_position[0],obj_start_position[1], 0.76]

            obj1.set_position(pos) 
            pr.step()
            next_wps=list()
            curr_wp = start_pos.copy()

            for i in range(3):
                wp=[curr_wp[0]+push_dirs[action[i]][0]*0.20,curr_wp[1]+push_dirs[action[i]][1]*0.20,curr_wp[2]]
                next_wps.append(wp)
                curr_wp = wp.copy()
            for step in range(len(next_wps)):
                try:
                    path = agent.get_path( #[-math.pi,0,math.pi]
                        position=next_wps[step],euler=[-math.pi,0,start_ori],ignore_collisions=True)
                except ConfigurationPathError as _:
#                     print('Could not find path')
                    continue
                done=False
                while not done:
                    done = path.step()
                    pr.step()
            obj1_pos = obj1.get_position()
            target_position = np.array(obj1_pos[:2])  - np.array(center_robot)
            pr.step()
            position_pairs[tuple(init_position.tolist())].append(target_position)
            ur10_model.remove()
            table_model.remove()
            obj1.remove()
            pr.step()
            pbar.update(1)
pbar.close()
with open('position_pairs.pickle', 'wb') as handle:
    pickle.dump(position_pairs, handle)
## If data is already generated
with open('position_pairs.pickle', 'rb') as handle:
    position_pairs = pickle.load(handle)
### DRAW DEMOS ###
all_demos=list()
for i, init in enumerate(position_pairs.keys()):
    fig, ax = plt.subplots(figsize=(14,10))
    demos=list()
    for goal in position_pairs[init]:
        if np.linalg.norm(np.array(goal)-np.array(init))>5e-2:
            demos.append(np.array(goal))
    while len(demos)>0:
        goal = demos.pop()
        all_demos.append((init,goal))
        plt.arrow(init[0],init[1],goal[0]-init[0],goal[1]-init[1],width=0.01)
        demos = list(filter(lambda point: np.linalg.norm(np.array(point)-np.array(goal))>0.05,demos))
        

    circle = patches.Circle(init, 0.05, color=colors[3],ec='black')
    ax.add_patch(circle)
    ax.set_xlim(0.5,-1.1)
    ax.set_ylim(-0.75,0.75)   
    fig.savefig('Init_pos_'+str(i)+'.png')
    plt.show()
objects_start_positions=list()
objects_target_positions=list()
for demo in all_demos:
    objects_start_positions.append((demo[0][0]+center[0],demo[0][1]+center[1]))
    objects_target_positions.append((demo[1][0]+center[0],demo[1][1]+center[1]))
## Save start and target positions for future task usage
with open('obj_st_pos.pickle', 'wb') as handle:
    pickle.dump(objects_start_positions, handle)
with open('obj_tar_pos.pickle', 'wb') as handle:
    pickle.dump(objects_target_positions, handle)<jupyter_output><empty_output><jupyter_text>Find Solutions via Network<jupyter_code>center = (0.37,-0.19)

stick = {
    'main': [
        (0.14, 0.08, -0.12, 0, 0, 0.50),
        (0.14, 0.08, -0.26, 0, 0, 0.50),
        (0.14, 0.08, -0.40, 0, 0, 0.50),
        (0.14, 0.08, -0.54, 0, 0, 0.50)        
    ],
    (1,0): (-0.66, 0.00),
    (-1,0):( 0.00, 0.00),
    (0,1): (-0.33,-0.09),
    (0,-1):(-0.33, 0.09),
    'edges':[
        (1,2),
        (2,3),
        (3,4),
    ]    
}

T_shaped = {
    'main': [
        (0.14, 0.08, -0.12, 0.00, 0, 0.50),
        (0.14, 0.08, -0.26, 0.00, 0, 0.50),
        (0.14, 0.08, -0.40, 0.00, 0, 0.50),
        (0.16, 0.08, -0.51, 0.00, np.pi/2, 0.50)        
    ],
    (1,0): (-0.60, 0.00),
    (-1,0):( 0.00, 0.00),
    (0,1): (-0.30,-0.09),
    (0,-1):(-0.30, 0.09),
    'edges':[
        (1,2),
        (2,3),
        (3,4),
    ]    
}

L_shaped = {
    'main': [
        (0.14, 0.08, -0.12, 0.00, 0, 0.50),
        (0.14, 0.08, -0.26, 0.00, 0, 0.50),
        (0.14, 0.08, -0.40, 0.00, 0, 0.50),
        (0.16, 0.08, -0.43, 0.12, np.pi/2, 0.50)        
    ],
    (1,0): (-0.52, 0.00),
    (-1,0):( 0.00, 0.00),
    (0,1): (-0.26,-0.09),
    (0,-1):(-0.26, 0.09),
    'edges':[
        (1,2),
        (2,3),
        (3,4),
    ]    
}

inv_L_shaped = {
    'main': [
        (0.14, 0.08, -0.12, 0.00, 0, 0.50),
        (0.14, 0.08, -0.26, 0.00, 0, 0.50),
        (0.14, 0.08, -0.40, 0.00, 0, 0.50),
        (0.16, 0.08, -0.43, -0.12, np.pi/2, 0.50)        
    ],
    (1,0): (-0.52, 0.00),
    (-1,0):( 0.00, 0.00),
    (0,1): (-0.26,-0.09),
    (0,-1):(-0.26, 0.09),
    'edges':[
        (1,2),
        (2,3),
        (3,4),
    ]    
}

Hook = {
    'main': [
        (0.14, 0.08, -0.12, 0.00, 0, 0.50),
        (0.14, 0.08, -0.26, 0.00, 0, 0.50),
        (0.14, 0.08, -0.40, 0.00, 0, 0.50),
        (0.16, 0.08, -0.09, 0.12, np.pi/2, 0.50)        
    ],
    (1,0): (-0.52, 0.00),
    (-1,0):( 0.00, 0.00),
    (0,1): (-0.26,-0.09),
    (0,-1):(-0.26, 0.09),
    'edges':[
        (1,2),
        (2,3),
        (1,4),
    ]    
}

inv_Hook = {
    'main': [
        (0.14, 0.08, -0.12, 0.00, 0, 0.50),
        (0.14, 0.08, -0.26, 0.00, 0, 0.50),
        (0.14, 0.08, -0.40, 0.00, 0, 0.50),
        (0.16, 0.08, -0.09, -0.12, np.pi/2, 0.50)        
    ],
    (1,0): (-0.52, 0.00),
    (-1,0):( 0.00, 0.00),
    (0,1): (-0.26,-0.09),
    (0,-1):(-0.26, 0.09),
    'edges':[
        (1,2),
        (2,3),
        (1,4),
    ]    
}
import numpy as np
order_list=[
    (-1, 0),
    ( 0, 1),
    (1, 0),
    ( 0, -1),
]
def generateTool(tool,orientation):
    tool_parts=list()
    for part in tool['main']:
        tool_part_info = dict()
        tool_part_info['pos'] = [center[0]+ np.cos(orientation) * part[2] - np.sin(orientation) * part[3],
                                 center[1]+ np.sin(orientation) * part[2] + np.cos(orientation) * part[3]]
        tool_part_info['shape'] = part[0:2]
        tool_part_info['angle'] = part[4]
        if np.abs(orientation-np.pi/2)<1e-4 or np.abs(orientation-3*np.pi/2)<1e-4:
            tool_part_info['angle'] = np.pi/2 - tool_part_info['angle']
        tool_parts.append(tool_part_info)
    rotation_count = int(orientation/(np.pi/2))
    robot_parts_after_rot=dict()
    for ind,rot_ind in enumerate(list(range(rotation_count,4)) + list(range(rotation_count))):
        new_pos = tool[order_list[rot_ind]]
        new_pos = [np.cos(orientation) * new_pos[0] - np.sin(orientation) * new_pos[1],
                   np.sin(orientation) * new_pos[0] + np.cos(orientation) * new_pos[1]]
        robot_parts_after_rot[order_list[ind]] = new_pos
    return tool_parts, robot_parts_after_rot

import collision
def check_collision(obj_info, obj_info2):
    # Object 1    
    if obj_info['type']==0:
        center_object = collision.Vector(obj_info['position'][0],obj_info['position'][1])
        obj1 =collision.Circle(center_object, (obj_info['shape'][0])/2+0.05)
    elif obj_info['type']==1:
        center_object = collision.Vector(obj_info['position'][0],obj_info['position'][1])
        obj1 =collision.Poly.from_box(center_object, obj_info['shape'][0]+0.05,obj_info['shape'][1]+0.05)
        obj1.angle = obj_info['angle']

    # Object 2

    center_object = collision.Vector(obj_info2['position'][0],obj_info2['position'][1])
    obj2 =collision.Circle(center_object, obj_info2['shape'][0]/2)

    return collision.collide(obj1,obj2)
import torch
from torch import FloatTensor as FT
Rewards=dict()
Selected_Tool_and_Action=dict()

animate=True
tools_master= [
        stick,
        L_shaped,
        inv_L_shaped,
        T_shaped,
        Hook,
        inv_Hook
    ]

possible_indices={
    'stick': [(0,0,'stick'),],
    'stick-left': [(0,np.pi/2,'stick-left'),],    
    'stick-right': [(0,3*np.pi/2,'stick-right'),], 
    'L_shaped': [(1,0,'L_shaped'),],
    'inv_L_shaped': [(2,0,'inv_L_shaped'),],
    'L_hook': [(1,np.pi/2,'L_hook'),],
    'R_hook': [(2,3*np.pi/2,'R_hook'),],
    'all': [(0,0,'stick'),(0,np.pi/2,'stick-left'),(0,3*np.pi/2,'stick-right'),
             (1,0,'L_shaped'),(2,0,'inv_L_shaped'),(1,np.pi/2,'L_hook'),(2,3*np.pi/2,'R_hook')]
}
for key in possible_indices.keys():
    if key == 'all':
        animate=True
    else:
        animate=False
    Rewards[key]=list()
    Selected_Tool_and_Action[key]=list()
    tools = [[tools_master[i[0]],i[1],i[2]] for i in possible_indices[key]]
    for trial in range(len(objects_start_positions)):

        goal_object_x = objects_start_positions[trial][0]
        goal_object_y = objects_start_positions[trial][1]

        target_position=objects_target_positions[trial]


        num_of_objects = 6 
        num_of_rel = num_of_objects*(num_of_objects-1)

        num_of_push= 3
        num_of_traj_tool_ori = (4** num_of_push) 
        num_of_traj_tool= num_of_traj_tool_ori
        num_of_traj = num_of_traj_tool*len(tools) 
        object_states = torch.zeros(num_of_traj,40*num_of_push,num_of_objects,8)
        object_shapes = torch.zeros(num_of_traj,num_of_objects,5)
        object_rels   = torch.zeros(num_of_traj,num_of_rel,4)
        selected_actions = torch.zeros(num_of_traj,num_of_push,2)
        robot_relative_posses_for_each_tool=dict()
        
        corresponding_tool = torch.zeros(num_of_traj,1)

        for tool_ind, tool_template_set in enumerate(tools):
            
            tool_template=tool_template_set[0]
            tool_ori = tool_template_set[1]
            
            tool_parts_positions,robot_relative_posses = generateTool(tool_template,tool_ori)
            robot_relative_posses_for_each_tool[tool_ind] = robot_relative_posses 
            st,end = tool_ind*num_of_traj_tool , (tool_ind+1)*num_of_traj_tool
            corresponding_tool[tool_ind*num_of_traj_tool:(tool_ind+1)*num_of_traj_tool] = tool_ind
            for i in range(4):
                tool_part = tool_parts_positions[i]
                object_shapes[st:end,i+1,:] = FT([0.0000, 1, 0.50, tool_part['shape'][0], tool_part['shape'][1]])
                object_states[st:end,0,i+1,:3] = FT([tool_part['pos'][0],tool_part['pos'][1],tool_part['angle']])
            cnt=0
            for i in range(num_of_objects):
                for j in range(num_of_objects):
                    if i!=j:
                        if (i,j) in tool_template['edges'] or (j,i) in tool_template['edges']: 
                            object_rels[st:end,cnt,:]=FT([0,1,0,0])
                        cnt=cnt+1

        object_shapes[:,0,:]=FT([1.0000, 0.0000, 1.0000, 0.1000, 0.1000])    
        object_shapes[:,5,:]=FT([0.0000, 0.0000, 0.5000, 0.1000, 0.1000])
        object_states[:,:,:,2]=object_states[:,0:1,:,2]
        object_states[:,0,5,:2] = FT([goal_object_x,goal_object_y])
        object_rels[:,:,0]=1

        def test_trajs(object_states_sub):
            #
            to_be_pred = object_states_sub.to(device)
            to_be_pred[:, 1:, num_of_objects-1:, :] = 0
            x = dict()
            x['objects_shape'] = object_shapes.to(device)
            x['relation_info'] = object_rels.to(device)
            for timestep in range(1, 40):
                x['objects_state'] = to_be_pred[:, timestep-1, :, :]
                pred = dataset.gp_pp(x)
                to_be_pred[:, timestep, num_of_objects-1:, :6] = pred[:,num_of_objects-2:,:]
            return to_be_pred

        push_dirs = [(1,0),(-1,0),(0,1),(0,-1)]
        for step in range(num_of_push): # 
            curr_center_posses= torch.zeros(num_of_traj,2)
            curr_center_posses[:,:] = FT([center[0],center[1]])
            for i in range(step):
                curr_center_posses[:,:]+= selected_actions[:,i,:]
            for tool_ind, tool_template_set in enumerate(tools):

                tool_template=tool_template_set[0]
                tool_ori = tool_template_set[1]
                st,end = tool_ind*num_of_traj_tool ,(tool_ind+1)*num_of_traj_tool
                relative_posses = robot_relative_posses_for_each_tool[tool_ind]
                foo = 4**(step)
                bar = 4**(num_of_push-step-1)
                for order in range(foo):
                    for action in range(4):
                        st2,end2 = st+order*bar*4+action*bar,st+order*bar*4+(action+1)*bar
                        object_states[st2:end2,step*40,0,:2] = curr_center_posses[st2:end2,:] + FT(relative_posses[push_dirs[action]])
                        push_dists = np.array(push_dirs[action])*0.20
                        selected_actions[st2:end2,step,:] = torch.from_numpy(push_dists).float()
                        for obj_ind in range(5):
                            for j in range(st2,end2):
                                obj_init_pos = object_states[j,step*40,obj_ind,:2]
                                object_states[j,step*40:(step+1)*40,obj_ind,0] = torch.linspace(obj_init_pos[0],obj_init_pos[0]+push_dists[0],40)
                                object_states[j,step*40:(step+1)*40,obj_ind,1] = torch.linspace(obj_init_pos[1],obj_init_pos[1]+push_dists[1],40)    


            object_states[:,step*40+1:(step+1)*40,:,3:5]=object_states[:,step*40+1:(step+1)*40,:,0:2]-object_states[:,step*40:(step+1)*40-1,:,0:2]
            object_states[:,step*40:(step+1)*40-1,:5,6:8]=object_states[:,step*40+1:(step+1)*40,:5,3:5]
            Net.eval()
            with torch.no_grad():
                predicted=test_trajs(object_states[:,step*40:(step+1)*40])
                object_states[:,step*40:(step+1)*40] = predicted
            if step!=num_of_push-1:
                object_states[:,(step+1)*40]=object_states[:,(step+1)*40-1]

        valid_indexes = []
        for i in range(num_of_traj):
            goal_object_info={
                'type':0,
                'position':[object_states[i,0,-1,0],object_states[i,0,-1,1]],
                'shape': [object_shapes[i,-1,3],object_shapes[i,-1,4]]
            }
            any_collision=False
            for part in range(5):
                part_info={
                    'type':1,
                    'position':[object_states[i,0,part,0],object_states[i,0,part,1]],
                    'angle': object_states[i,0,part,2],
                    'shape': [object_shapes[i,part,3],object_shapes[i,part,4]]
                }
                if part == 0:
                    part_info['type']=0
                any_collision = any_collision or check_collision(part_info,goal_object_info)
            if not any_collision:
                valid_indexes.append(i)
        if len(valid_indexes)>0:
            values,indices = torch.topk(-torch.norm(object_states[valid_indexes,:,5,:2]-FT(target_position),dim=2).reshape(-1),1)
            Rewards[key].append(values[0].item())
            
            index_best =(valid_indexes[int(np.floor(indices.numpy()[0]/120))],
                         int(np.floor(indices.numpy()[0]%120)))
            selected_tool= tools[int(corresponding_tool[index_best[0]].item())][2]
            selected_action = selected_actions[index_best[0]]
            Selected_Tool_and_Action[key].append((selected_tool,selected_action.cpu().numpy(),index_best[1]))
            top_indices=list()
            for ind in (np.floor(indices.numpy()/120)).astype(int):
                top_indices.append(valid_indexes[ind])
        else:
            val = -torch.norm(FT([goal_object_x,goal_object_y])- FT(target_position))
            Selected_Tool_and_Action[key].append((tools[0][2],0))

            Rewards[key].append(val)


        if animate:
            for top_ind,nx in enumerate(top_indices):
                fig, ax = plt.subplots(1, 1,figsize=(8,8))
                plt.close();
                titles= ['Predicted']
                for ax_ind in range(1):
                    ax.set_xlim((-0.8, 0.8))
                    ax.set_ylim((-0.8, 0.8))
                    ax.set_title(titles[ax_ind])

                colors=dict()
                colors[0]=[0.2,0.2,0.2]
                colors[1]=[0,0,1]
                colors[2]=[0,1,0]
                colors[3]=[1,1,0]
                colors[4]=[0,1,1]
                colors[5]=[1,0,1]
                colors[6]=[0,0.5,0]
                colors[7]=[0,0,0.5]
                colors[8]=[0.5,0.5,0]
                colors[9]=[0,0.5,0.5]
                colors[10]=[0.5,0,0.5]
                colors[11]=[0.5,1,0]
                colors[12]=[0,0.5,1]
                colors[13]=[1,1,0.5]
                colors[14]=[0.5,1,1]
                colors[15]=[1,0.5,1]
                obj_list=dict()
                obj_list[0]=[]
                obj_list[1]=[]

                edge_list=dict()
                edge_list[0]=[]
                edge_list[1]=[]
                for ax_ind in range(1):
                    for obj_index,obj in enumerate(object_shapes[nx]):
                        c=colors[1+obj_index]
                        if obj[1]==0:
                            circle = patches.Circle((0,0),obj[-1]/2,color=c,ec='black')
                            obj_list[ax_ind].append((0,circle))
                        else:
                            patch_mini_list=list()
                            for k in [-1,1]:
                                for l in [-1,1]:
                                    patch = patches.Rectangle((0,0), k*obj[-2]/2, l*obj[-1]/2,ec='black', fc=c)
                                    patch_mini_list.append(patch)
                            obj_list[ax_ind].append((1,patch_mini_list))
                    edg= object_rels[nx]
                    cnt=0
                    for i in range(num_of_objects):
                        for j in range(num_of_objects):
                            if i!=j: 
                                if edg[cnt,1]==1:
                                    c='black'
                                elif edg[cnt,2]==1:
                                    c='purple'
                                elif edg[cnt,3]==1:
                                    c='brown'
                                else:
                                    cnt=cnt+1
                                    continue
                                cnt=cnt+1
                                edge, = ax.plot([], [], color=c, lw=3) 
                                edge_list[ax_ind].append((edge,(i,j)))                
                def init():
                    for ax_ind in range(1):
                        for obj in obj_list[ax_ind]:
                            if obj[0]==0:
                                ax.add_patch(obj[1])
                                ax.add_patch(obj[1])

                            else:
                                for patch in obj[1]:
                                    ax.add_patch(patch)
                    goal_pos = patches.Rectangle((target_position[0]-0.025,target_position[1]-0.025), 0.05, 0.05,ec='red',alpha=0.25, fc='red')
                    ax.add_patch(goal_pos)
                    return []  

                def animate(i):
                    for ax_ind in range(1):
                        positions=object_states[nx,i,:,:3].cpu().numpy()
                        for obj_index, obj in enumerate(obj_list[ax_ind]):
                            if obj[0]==0:
                                obj[1].center=(positions[obj_index,0],positions[obj_index,1])           
                            else:
                                for patch in obj[1]:
                                    patch.xy=(positions[obj_index,0],positions[obj_index,1])       
                                    patch.angle=np.rad2deg(positions[obj_index,2])
                        for edge in edge_list[ax_ind]:
                            edge_line=edge[0]
                            edge_info=edge[1]
                            obj1=edge_info[0]
                            obj2=edge_info[1]
                            xx=[positions[obj1,0],positions[obj2,0]]
                            yy=[positions[obj1,1],positions[obj2,1]]
                            edge_line.set_data(xx,yy)        
                    return []
                anim = animation.FuncAnimation(fig, animate, init_func=init,
                                           frames=120, interval=50, blit=True)
                anim.save('best_tools_videos/'+str(key)+'_'+str(trial)+'best'+str(top_ind)+'.gif', writer='imagemagick', fps=30)
                display(HTML(anim.to_html5_video()))
error_cm = dict() 
succes_count_5 = dict()
succes_count_10 = dict()
for key in Rewards.keys():
    errors = -np.array(Rewards[key])
    error_cm[key] = np.mean(errors)
    succes_count_5[key] = np.sum(errors<0.05)
    succes_count_10[key] = np.sum(errors<0.10)    

cm_errors_for_all = -np.array(Rewards['all'])
np.save('simulation_all_trajectories_error.npy',cm_errors_for_all)
with open('error_cm_sim.pickle', 'wb') as handle:
    pickle.dump(error_cm, handle)
with open('succes_count_5_sim.pickle', 'wb') as handle:
    pickle.dump(succes_count_5, handle)
with open('succes_count_10_sim.pickle', 'wb') as handle:
    pickle.dump(succes_count_10, handle)    
Selected_Tool_and_Action2=dict()
for key in Selected_Tool_and_Action.keys():
    Selected_Tool_and_Action2[key]=list()
    for tup in Selected_Tool_and_Action[key]:
        action=tup[1]
        if type(action) is np.ndarray:
            action=action.tolist()
        if action==0:
#             print('lolw',tup)
            Selected_Tool_and_Action2[key].append((tup[0],0))
        else:
            Selected_Tool_and_Action2[key].append((tup[0],action,tup[2]))
objects_target_positions[0]
tasks_and_goals=dict()
tasks_and_goals['Init_obj_pos']=list()
tasks_and_goals['Target_obj_pos']=list()
for trial in range(166):
    
    start_position  = tuple(objects_start_positions[trial])
    target_position = tuple(objects_target_positions[trial])

    tasks_and_goals['Init_obj_pos'].append(start_position)
    tasks_and_goals['Target_obj_pos'].append(target_position)  
    
tasks_and_goals['tool_action'] = Selected_Tool_and_Action2
import json
with open('tool_to_transfer.json', 'w') as f:
    json.dump(tasks_and_goals, f)<jupyter_output><empty_output><jupyter_text>Rerun it in Simulation<jupyter_code>from os.path import dirname, join, abspath
from pyrep import PyRep
from pyrep.robots.arms.ur10 import UR10
from pyrep.objects.shape import Shape
from pyrep.objects.dummy import Dummy
from pyrep.objects.joint import Joint
from pyrep.backend import sim
from pyrep.const import PrimitiveShape,JointType
from pyrep.const import ConfigurationPathAlgorithms as Algos
from pyrep.errors import ConfigurationPathError
from pyrep.backend.sim import simGetObjectPosition,simGetObjectHandle
import numpy as np
import math
import copy
import time
import signal
from multiprocessing import Process, Value, Lock
colors=dict()
colors[0]=[0.2,0.2,0.2]
colors[1]=[0,1,0]
colors[2]=[0,0,1]
colors[3]=[1,1,0]
colors[4]=[0,1,1]
colors[5]=[1,0,0]
colors[6]=[0,0.5,0]
colors[7]=[0,0,0.5]
colors[8]=[0.5,0.5,0]
colors[9]=[0,0.5,0.5]
colors[10]=[0.5,0,0.5]
colors[11]=[0.5,1,0]
colors[12]=[0,0.5,1]
colors[13]=[1,1,0.5]
colors[14]=[0.5,1,1]
colors[15]=[1,0.5,1]
import os
import pickle
import json
import random

center = (-0.45,0)
center_old = (0.37,-0.19)

tools={
    'stick': { 
        'main': [
            (0.14, 0.08, -0.12, 0, 0, 0.50),
            (0.14, 0.08, -0.26, 0, 0, 0.50),
            (0.14, 0.08, -0.40, 0, 0, 0.50),
            (0.14, 0.08, -0.54, 0, 0, 0.50)        
        ],
        'ori':0
    },
    'stick-left': { 
        'main': [
            (0.14, 0.08, -0.12, 0, 0, 0.50),
            (0.14, 0.08, -0.26, 0, 0, 0.50),
            (0.14, 0.08, -0.40, 0, 0, 0.50),
            (0.14, 0.08, -0.54, 0, 0, 0.50)        
        ],
        'ori':-np.pi/2
    },
    'stick-right': { 
        'main': [
            (0.14, 0.08, -0.12, 0, 0, 0.50),
            (0.14, 0.08, -0.26, 0, 0, 0.50),
            (0.14, 0.08, -0.40, 0, 0, 0.50),
            (0.14, 0.08, -0.54, 0, 0, 0.50)        
        ],
        'ori':np.pi/2
    },
    'L_shaped': { 
        'main': [
            (0.14, 0.08, -0.12, 0.00, 0, 0.50),
            (0.14, 0.08, -0.26, 0.00, 0, 0.50),
            (0.14, 0.08, -0.40, 0.00, 0, 0.50),
            (0.16, 0.08, -0.43, 0.12, np.pi/2, 0.50)    
        ],
        'ori':0
    },
    'inv_L_shaped': { 
        'main': [
            (0.14, 0.08, -0.12, 0.00, 0, 0.50),
            (0.14, 0.08, -0.26, 0.00, 0, 0.50),
            (0.14, 0.08, -0.40, 0.00, 0, 0.50),
            (0.16, 0.08, -0.43, -0.12, np.pi/2, 0.50)         
        ],
        'ori':0
    },    
    'L_hook': { 
        'main': [
            (0.14, 0.08, -0.12, 0.00, 0, 0.50),
            (0.14, 0.08, -0.26, 0.00, 0, 0.50),
            (0.14, 0.08, -0.40, 0.00, 0, 0.50),
            (0.16, 0.08, -0.43, 0.12, np.pi/2, 0.50)    
        ],
        'ori':-np.pi/2
    },
    'R_hook': { 
        'main': [
            (0.14, 0.08, -0.12, 0.00, 0, 0.50),
            (0.14, 0.08, -0.26, 0.00, 0, 0.50),
            (0.14, 0.08, -0.40, 0.00, 0, 0.50),
            (0.16, 0.08, -0.43, -0.12, np.pi/2, 0.50)         
        ],
        'ori':np.pi/2
    },        
}

def generateTool(tool_template):
    objects=list()
    for part in tool_template['main']:
        obj_type = PrimitiveShape.CUBOID
        obj= Shape.create(type=obj_type,
                            size=[part[0], part[1], 0.05],
                            color=colors[3],
                            static=True, respondable=True)
        pos = [center[0]+part[2],center[1]+part[3], 0.90]
        obj.set_position(pos)
        obj.set_orientation([0,0,part[4]])
        objects.append(obj)
    return objects

center_old = (0.37,-0.19)

import json
with open('tool_to_transfer.json', 'r') as f:
    tasks_and_goals = json.load( f)

SCENE_FILE = 'sim_files/empty_ode.ttt'
pr = PyRep()
pr.launch(SCENE_FILE, headless=True)
pr.set_simulation_timestep(0.05)
pr.start()    

len(tasks_and_goals['Init_obj_pos'])
ur10_model = pr.import_model('sim_files/ur10.ttm')
table_model = pr.import_model('sim_files/table.ttm')
# Reposition Sim
for _ in range(1000):
    pr.step()
    
ur10_model.remove()
table_model.remove()
succes_count_all = dict()
rewards_all = dict()
for task_key in tasks_and_goals['tool_action'].keys():
    rewards = list()
    succes_count = 0
    print(task_key)
    for ind in range(166):
        ur10_model = pr.import_model('sim_files/ur10.ttm')
        table_model = pr.import_model('sim_files/table.ttm')

        agent = UR10()
        agent.max_velocity=0.10
        agent.max_acceleration=1
        agent.max_jerk=10
        agent_hand = Shape(simGetObjectHandle('UR10_hand'))
        try:
            path = agent.get_path( #[-math.pi,0,math.pi]
                position=[center[0],center[1],0.90],euler=[-math.pi,0,0],ignore_collisions=True)
        except ConfigurationPathError as _:
            print('Could not find path')
        path.set_to_end()
        pr.step()
        for _ in range(10):
            pr.step()

        new_goal_pos = tasks_and_goals['Init_obj_pos'][ind]-np.array(center_old) + np.array(center)
        new_target_pos = tasks_and_goals['Target_obj_pos'][ind]-np.array(center_old) + np.array(center)

        tool=tasks_and_goals['tool_action'][task_key][ind][0]
        selected_action = tasks_and_goals['tool_action'][task_key][ind][1]
        objects= generateTool(tools[tool])
        objects.insert(0,agent_hand)
        grouped_shape= pr.group_objects(objects)
        sim.simSetObjectSpecialProperty(grouped_shape.get_handle(),1)
        sim.simSetObjectParent(grouped_shape.get_handle(),simGetObjectHandle('UR10_connection'),True)
                ## PUSH Action

        start_pos = [center[0],center[1],0.75]
        start_ori = tools[tool]['ori']
        ## Repositioning STEP 1
        step_positions=[
            start_pos.copy(),
            start_pos
        ]
        step_positions[0][2]=0.95
        for step in range(2):
            try:
                path = agent.get_path( #[-math.pi,0,math.pi]
                    position=step_positions[step],euler=[-math.pi,0,start_ori],ignore_collisions=True)
            except ConfigurationPathError as _:
                print('Could not find path')
                continue

            if step==0:
                path.set_to_end()
                pr.step()
            else:
                done=False
                while not done:
                    done = path.step()
                    pr.step()

        ###PUT OBJECT###

        obj_type = PrimitiveShape.CYLINDER
        obj1= Shape.create(type=obj_type,
                            size=[0.10, 0.10, 0.10],
                            color=colors[4],
                            static=False, respondable=True)
        pos = [new_goal_pos[0],new_goal_pos[1], 0.76]
        obj1.set_position(pos)

        obj_type = PrimitiveShape.CYLINDER
        obj2= Shape.create(type=obj_type,
                            size=[0.2, 0.2, 0.01],
                            color=colors[5],
                            static=True, respondable=False)
        pos = [new_target_pos[0],new_target_pos[1], 0.705]
        obj2.set_position(pos)


        ###########
        if selected_action==0:
            print('') # No action
        else:
            push_len= tasks_and_goals['tool_action'][task_key][ind][2]+1
            next_wps=list()
            curr_wp = start_pos.copy()
            for i in range(3):
                curr_push_len = min(40,push_len-i*40)
                if curr_push_len<1:
                    break
                mult=curr_push_len/40.0
                wp=[curr_wp[0]+mult*float(selected_action[i][0]),curr_wp[1]+mult*float(selected_action[i][1]),curr_wp[2]]
                next_wps.append(wp)
                curr_wp = wp.copy()
            for step in range(len(next_wps)):
                try:
                    path = agent.get_path( #[-math.pi,0,math.pi]
                        position=next_wps[step],euler=[-math.pi,0,start_ori],ignore_collisions=True)
                except ConfigurationPathError as _:
                    print('Could not find path')
                    continue
                done=False
                while not done:
                    done = path.step()
                    pr.step()
        obj1_pos = obj1.get_position()
        obj2_pos = obj2.get_position()
        reward =np.linalg.norm([obj1_pos[0]-obj2_pos[0],obj1_pos[1]-obj2_pos[1]])
        if reward<0.10:
            succes_count=succes_count+1
        rewards.append(reward) 
        ur10_model.remove()
        table_model.remove()
        obj1.remove()
        obj2.remove()
    succes_count_all[task_key] =succes_count
    rewards_all[task_key] =rewards
    print(succes_count,np.mean(rewards))

print(succes_count_all)
print('--------')
for key in rewards_all.keys():
    print(key)
    print(np.mean(rewards_all[key])) 
print(rewards_all)
print(succes_count_all)
print('--------')
for key in rewards_all.keys():
    print(key)
    print(np.mean(rewards_all[key])) 
with open('succes_counts_transfer.pickle', 'wb') as f:
    pickle.dump(succes_count_all, f)
with open('rewards_all_transfer.pickle', 'wb') as f:
    pickle.dump(rewards_all, f)   
import json
with open('succes_count_all.json', 'w') as f:
    json.dump(succes_count_all, f)

with open('rewards_all.json', 'w') as f:
    json.dump(rewards_all, f)    
pr.start()
rewards_all['stick']
pr.stop()
pr.shutdown()<jupyter_output><empty_output><jupyter_text>Plot Time<jupyter_code>with open('error_cm_sim.pickle', 'rb') as handle:
    error_cm_network = pickle.load(handle)
with open('succes_count_10_sim.pickle', 'rb') as handle:
    succes_count_network = pickle.load(handle)    
with open('succes_counts_transfer.pickle', 'rb') as f:
    succes_count_sim=pickle.load(f)
with open('rewards_all_transfer.pickle', 'rb') as f:
    rewards_all_sim = pickle.load(f)   
error_cm_sim=dict()
for key in rewards_all_sim.keys():
    error_cm_sim[key] = np.mean(rewards_all_sim[key])
rewards_all_sim.keys()
names = ['stick', 'stick-left', 'stick-right', 'L_shaped', 'inv_L_shaped', 'L_hook', 'R_hook', 'all']
new_names = ['Stick', 'Left Rotated Stick','Right Rotated Stick', 'L-shaped', 'Inv-L-shaped', 
             'Left Rotated L-shaped ', 'Right Rotated Inv-L-shaped', 'All Tools']

prediction_errors=np.zeros((len(names),))
simulation_errors=np.zeros((len(names),))

succes_counts_prediction=np.zeros((len(names),))
succes_counts_simulation=np.zeros((len(names),))

for index, key in enumerate(names):
    prediction_errors[index] = error_cm_network[key]
    succes_counts_prediction[index] = succes_count_network[key]
    simulation_errors[index] = error_cm_sim[key]
    succes_counts_simulation[index] = succes_count_sim[key]    
fig,ax = plt.subplots(1,2,figsize=(16,6))
ax[0].bar(new_names,prediction_errors,color="gray")
ax[1].bar(new_names,simulation_errors,color="gray")
ax[0].set_ylabel("Mean Errors (cm)",fontsize=16)
ax[0].set_xlabel("Predictions",fontsize=16)
ax[1].set_xlabel("Simulation",fontsize=16)
ax[1].set_xticklabels(new_names, rotation = 45, ha="right",fontsize=12)
ax[0].set_xticklabels(new_names, rotation = 45, ha="right",fontsize=12)
ax[0].set_ylim(0,0.25)
ax[1].set_ylim(0,0.25)
fig.savefig("mean_errors_tools.png",bbox_inches='tight')
plt.show()
fig,ax = plt.subplots(1,2,figsize=(16,6))
ax[0].bar(new_names,succes_counts_prediction,color="gray")
ax[1].bar(new_names,succes_counts_simulation,color="gray")
ax[0].set_ylabel("Succes Counts",fontsize=16)
ax[0].set_xlabel("Predictions",fontsize=16)
ax[1].set_xlabel("Simulation",fontsize=16)
ax[1].set_xticklabels(new_names, rotation = 45, ha="right",fontsize=12)
ax[0].set_xticklabels(new_names, rotation = 45, ha="right",fontsize=12)
ax[0].set_ylim(0,140)
ax[1].set_ylim(0,140)
fig.savefig("succes_counts_tools.png",bbox_inches='tight')
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1605523782954387, 'fraction_numerical': 0.05816710838331706, 'mean_word_length': 2.1647292525014716, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '72796', 'license': 'MIT', 'n_tokens_mistral': 16740, 'n_tokens_neox': 15503, 'n_words': 2300, 'path': 'Tool_Experiment.ipynb', 'repo_name': 'Fzaero/Object-and-Relation-Centric-Representations-for-Push-Effect-Prediction', 'size': '43014'}