evanthebouncy · August 2, 2022 18:40
diff --git a/tantris_env.py b/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
	# 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