tantris_env.py

# the tiling environment, a grid world where you place tetris pieces

import random

# the dimension of the grid
L = 6

# hardcoding a few pieces, it's fine
PIECES = [
    [[1,1,1,1]], # horizontal line
    [[1], [1], [1], [1]], # vertical line
    [[1,1,1], [0,1,0]], # T
    [[1,0], [1,1], [1,0]], # sideways T poking right
    [[0,1,0], [1,1,1]], # upside down T
    [[0,1], [1,1], [0,1]], # sideways T poking left
    [[1,1], [1,1]], # square
]

# the logic go place a piece in the grid
def place_piece(grid, piece, x, y, hint = False):
    # deep copy a new grid
    new_grid = [row[:] for row in grid]
    # for each row in the piece
    for i in range(len(piece)):
        # for each column in the piece
        for j in range(len(piece[i])):
            # if the piece is not empty
            if piece[i][j] == 1:
                # check if the grid is out of bounds
                if x + j < 0 or x + j >= L or y + i < 0 or y + i >= L:
                    return False
                # check if the grid at the position is empty
                if grid[y+i][x+j] == 0:
                    # place the piece in the grid
                    to_put = 1 if not hint else str(PIECES.index(piece))
                    new_grid[y+i][x+j] = to_put
                else:
                    # if the grid is not empty, return False
                    return False
    return new_grid

# visualize the grid using ASCII characters
def visualize_grid(grid):
    for row in grid:
        row_to_print = ''
        for x in row:
            # check if the location is either 0, or 1, or a string
            if x == 0:
                row_to_print += '.'
            elif x == 1:
                row_to_print += '#'
            else:
                row_to_print += x
        print (row_to_print)
        
# visualize the ground-truth program
def visualize_program(program):
    grid = [[0 for x in range(L)] for y in range(L)]
    for piece_id, x, y in program:
        grid = place_piece(grid, PIECES[piece_id], x, y, hint = True)
    visualize_grid(grid)


# the execution of the program
# program is a list of tuples, where the tuple is (piece_id, x, y)
def exe(program, save_hint = False):
    # create a grid
    grid = [[0 for x in range(L)] for y in range(L)]
    # for each piece, x, y in the program
    for piece_id, x, y in program:
        # place the piece in the grid
        grid = place_piece(grid, PIECES[piece_id], x, y, hint = save_hint)
        # if the grid is False, execution failed
        if grid is False:
            return False
    # return the grid
    return grid

# keep trying until we can place a piece
# gives up early if cannot place a piece
def place_a_piece(grid, num_tries=20):
    if num_tries == 0:
        return False
    # randomly choose a piece to place
    piece_id = random.randint(0, len(PIECES) - 1)
    # randomly choose a location
    x = random.randint(0, L - 1)
    y = random.randint(0, L - 1)
    # place the piece
    new_grid = place_piece(grid, PIECES[piece_id], x, y)
    if new_grid == False:
        return place_a_piece(grid, num_tries=num_tries-1)
    return (piece_id, x, y), new_grid

# generate some legal programs
def generate_program():
    program = []
    grid = [[0 for x in range(L)] for y in range(L)]
    # randomly choose between 3 to 6 pieces
    num_pieces = random.randint(3, 6)
    # for each piece
    for i in range(num_pieces):
        result = place_a_piece(grid)
        if result == False:
            return generate_program()
        # place a piece
        (piece_id, x, y), grid = result
        # add the piece to the program
        program.append((piece_id, x, y))
    return program

def gen_data_pair():
    prog = generate_program()
    # executing the program
    grid = exe(prog, save_hint = False)
    prog_str = repr(prog).replace(' ', '')
    grid_str = repr(grid).replace(' ', '')
    return grid_str, prog_str

if __name__ == '__main__':
    prog = [(0, 0, 0), (1, 0, 2), (2, 3, 1), (3,3,3)]
    grid = exe(prog, save_hint = True)
    visualize_grid(grid)

    # generating a random data point
    grid_str, prog_str = gen_data_pair()
    print (grid_str)
    print (prog_str)
    # visualize the grid
    print ("grid interpreted")
    grid = eval(grid_str)
    visualize_grid(grid)
    # visualize the program
    print ("program visualized onto grid")
    visualize_program(eval(prog_str))

    # start counting time
    import time
    start = time.time()
    for i in range(1000):
        grid_str, prog_str = gen_data_pair()
        # visualize_program(eval(prog_str))
    # count time now
    end = time.time()
    print ("time:", end - start)

    # 0.276077 seconds for 1000 data points, how long for a million? use code to find out
    # 0.276077 seconds for 1000 data points, how long for a million? use code to find out