create_test_arrangements.py

import random

# creates few object arrangements for testing

try:
    import sim
except:
    print('--------------------------------------------------------------')
    print('"sim.py" could not be imported. This means very probably that')
    print('either "sim.py" or the remoteApi library could not be found.')
    print('Make sure both are in the same folder as this file,')
    print('or appropriately adjust the file "sim.py"')
    print('--------------------------------------------------------------')
    print('')

import time
import shapes
import numpy as np
import json
import math

n_blocks = 3
n_actions = n_blocks
max_cu = 3
max_cy = 2
max_s = 2
max_py = 2
max_co = 2

print('Program started')
sim.simxFinish(-1)  # just in case, close all opened connections
clientID = sim.simxStart('127.0.0.1', 19999, True, True, 5000, 5)  # Connect to CoppeliaSim
if clientID != -1:
    print('Connected to remote API server')

    # Now try to retrieve data in a blocking fashion (i.e. a service call):
    res, objs = sim.simxGetObjects(clientID, sim.sim_handle_all, sim.simx_opmode_blocking)
    if res == sim.simx_return_ok:
        print('Number of objects in the scene: ', len(objs))
    else:
        print('Remote API function call returned with error code: ', res)

    time.sleep(2)

    # Now retrieve streaming data (i.e. in a non-blocking fashion):
    startTime = time.time()

    black = (20 / 255, 20 / 255, 20 / 255)
    white = (248 / 255, 248 / 255, 248 / 255)
    red = (181 / 255, 37 / 255, 38 / 255)
    green = (0 / 255, 148 / 255, 60 / 255)
    yellow = (254 / 255, 182 / 255, 0 / 255)
    blue = (8 / 255, 145 / 255, 187 / 255)
    brown = (77 / 255, 38 / 255, 30 / 255)
    purple = (105 / 255, 52 / 255, 117 / 255)
    pink = (234 / 255, 152 / 255, 183 / 255)
    orange = (220 / 255, 65 / 255, 2 / 255)
    gray = (128 / 255, 128 / 255, 128 / 255)
    blueish = (28 / 255, 171 / 255, 174 / 255)
    greenish = (178 / 255, 179 / 255, 0)
    reddish = (132 / 255, 32 / 255, 57 / 255)
    yellowish = (234 / 255, 151 / 255, 0)
    brownish = (111 / 255, 88 / 255, 21 / 255)

    colors = [black, white, red, green, yellow, blue, brown, purple, pink, orange, gray, blueish, greenish, reddish,
              yellowish, brownish]

    even_list = [11]
    clean = False

    for j in even_list:
        print("trial: " + str(j))

        shapeslist = []
        # rands = np.random.randint(5, size=n_blocks)
        rands = [2, 3, 4]
        cu = 0
        cy = 0
        s = 0
        py = 0
        co = 0

        m = list(range(6))
        left = random.choice(m)
        front = random.choice(m)

        # sample shapes from the simulation
        for i in range(n_blocks):
            if rands[i] == 0:
                if cu < max_cu:
                    shapeslist.append(shapes.Shape(clientID, "Cuboid", cu))
                    cu += 1
                else:
                    rands[i] = 1
            if rands[i] == 1:
                if cy < max_cy:
                    shapeslist.append(shapes.Shape(clientID, "Cylinder", cy))
                    cy += 1
                else:
                    rands[i] = 2
            if rands[i] == 2:
                if s < max_s:
                    shapeslist.append(shapes.Shape(clientID, "Sphere", s))
                    s += 1
                else:
                    rands[i] = 3
            if rands[i] == 3:
                if py < max_py:
                    shapeslist.append(shapes.Shape(clientID, "Pyramid", py))
                    py += 1
                else:
                    rands[i] = 4
            if rands[i] == 4:
                if co < max_co:
                    shapeslist.append(shapes.Shape(clientID, "Cone", co))
                    co += 1
                else:
                    shapeslist.append(shapes.Shape(clientID, "Cuboid", cu))
                    cu += 1

        sample = []
        withoutAll = []

        # get lists which include all shapes except shape i respectively
        for n in range(n_blocks):
            withoutn = shapeslist.copy()
            withoutn.remove(shapeslist[n])
            withoutAll.append(withoutn)

        reshape = []
        arrangement = []

        order = list(range(n_blocks))

        # adapt size and color of each shape
        for shape in shapeslist:

            sim.simxStartSimulation(clientID, sim.simx_opmode_blocking)

            x = np.random.uniform(0.5, 1)
            y = x
            z = x

            xb = 1 / x
            yb = 1 / y
            zb = 1 / z

            rshape = [xb, yb, zb]
            reshape.append(rshape)

            sample_input = []
            sample_target = []

            r, g, b = random.choice(colors)
            fx = 0
            fy = 0

            shape.scale_shape(x, y, z)

            shape.set_color(r, g, b)
            sim.simxPauseSimulation(clientID, sim.simx_opmode_blocking)

        time.sleep(1)

        sim.simxStopSimulation(clientID, sim.simx_opmode_blocking)

        time.sleep(1)
        sim.simxStartSimulation(clientID, sim.simx_opmode_blocking)

        for a in range(n_actions):
            timestep = []

            # save previous state
            for s in range(len(shapeslist)):
                shape = shapeslist[s]

                properties = [list(shape.get_position_adapted()), list(shape.get_orientation_type_simple()),
                              list(shape.get_color())]
                size = shape.get_bounding_box()[0]
                properties.append([size])
                if shapes.out_of_bounds(shape.get_raw_position()):
                    properties.append(0)
                else:
                    properties.append(1)
                properties.append(shape.shape_type_numbered)

                timestep.append(properties)

            fx = 0
            fy = 0

            orientation_type = [1, 0, 0]

            rotation = [0, 0]
            if shapeslist[order[a]].shape_type_numbered == 2 or shapeslist[order[a]].shape_type_numbered == 1 or \
                    shapeslist[order[a]].shape_type_numbered == 4:
                rotation[0] = 1
                rotation[1] = 0
            else:
                rotation[0] = np.random.uniform(- 1, 1)
                rotation[1] = math.sin(math.acos(rotation[0]))
                if random.randint(0, 1) == 0:
                    rotation[1] = -rotation[1]

            left_right = float(np.random.uniform(-0.8, 0.8))
            front_back = float(np.random.uniform(-0.8, 0.8))
            up = 1

            if a == 1:
                left_right = left_right + float(np.random.uniform(-0.05, 0.05))
                front_back = front_back + float(np.random.uniform(-0.05, 0.05))

            if a == 2:
                left_right = left_right + float(np.random.uniform(-0.05, 0.05))
                front_back = front_back + float(np.random.uniform(-0.05, 0.05))

            shapeslist[order[a]].move_to(2, 2, [])

            time.sleep(1)

            shapeslist[order[a]].set_visual_orientation_simple([orientation_type[0], orientation_type[1],
                                                                orientation_type[2], rotation[0],
                                                                rotation[1]])
            # first set rotation
            orientation = list([orientation_type[0], orientation_type[1], orientation_type[2], rotation[0],
                                rotation[1]])
            # then set position
            position = shapeslist[order[a]].set_position([left_right, front_back, up], withoutAll[order[a]])

            time.sleep(1)
            sim.simxStartSimulation(clientID, sim.simx_opmode_blocking)
            time.sleep(2)

            sim.simxPauseSimulation(clientID, sim.simx_opmode_blocking)
            # save action
            properties = [[int(order[a])], list(position), list(orientation)]

            timestep.append(list(properties))

            arrangement.append(list(timestep))

        timestep = []
        # get resulting state
        for s in range(len(shapeslist)):
            shape = shapeslist[s]

            properties = []
            pos = shape.get_position_adapted()
            print("resulting position: " + str(pos))
            properties.append(list(pos))
            properties.append(list(shape.get_orientation_type_simple()))
            properties.append(list(shape.get_color()))
            properties.append([shape.get_bounding_box()[0]])
            if shapes.out_of_bounds(shape.get_raw_position()):
                properties.append(0)
            else:
                properties.append(1)
            properties.append(shape.shape_type_numbered)

            timestep.append(list(properties))

        input()

        arrangement.append(list(timestep))
        # save
        with open("test_arrangements/arrangement" + str(j) + ".json", 'w') as f:
            json.dump(list(arrangement), f, indent=2)
        sim.simxStartSimulation(clientID, sim.simx_opmode_blocking)
        # reset shapes
        for i in range(len(shapeslist)):
            shapeslist[i].scale_shape(reshape[i][0], reshape[i][1], reshape[i][2])
            shapeslist[i].turn_original_way_up()
            shapeslist[i].set_position_eval([i * 4 - 4, 3.3, 1])

        sim.simxStartSimulation(clientID, sim.simx_opmode_blocking)

    # Now send some data to CoppeliaSim in a non-blocking fashion:
    sim.simxAddStatusbarMessage(clientID, 'Hello CoppeliaSim!', sim.simx_opmode_oneshot)

    # Before closing the connection to CoppeliaSim, make sure that the last command sent out had time to arrive.
    # You can guarantee this with (for example):
    sim.simxGetPingTime(clientID)

    # Now close the connection to CoppeliaSim:
    sim.simxFinish(clientID)
else:
    print('Failed connecting to remote API server')
print('Program ended')