Spaces:
Sleeping
Sleeping
| import numpy as np | |
| from grip_env.pieces import PentominoPiece, COLOURS, COLOUR_NAMES, PIECE_NAMES | |
| from utils import layout_utils | |
| import math | |
| # Setting the layout fixed for now i.e start positions of each piece are at every 5x5 grid | |
| class BoardLayout(): | |
| ''' | |
| This class is used to generate a random layout of pentomino pieces on a board. | |
| Args: | |
| board_size: The number of grids in one dimension ofthe Pentomino Board | |
| num_pieces: The number of pieces to be placed on the board, including the target piece | |
| shapes: A list of the pentomino shapes to be selected from | |
| colours: A list of the pentomino colours to be selected from | |
| seed: The seed for the random number generator | |
| ''' | |
| def __init__(self, board_size: int, num_pieces: int, shapes: np.array, colours: np.array, seed: int): | |
| self.board_size = board_size | |
| self.num_pieces = num_pieces | |
| self.shapes = shapes | |
| self.colours = colours | |
| self.mapped_regions = layout_utils.map_regions(self.board_size) | |
| np.random.seed(seed) | |
| def set_start_positions(self) -> np.array: | |
| ''' | |
| Get start positions of everything on the board | |
| Args: | |
| regions: A list defining the regions where the piece will be spawned (['top', 'top left', 'right',...]) | |
| If None, then use all possible regions including the center grid | |
| Returns: | |
| all_start_positions - [[ax, ay], [p1x, p1y], [p2x, p2y], ....] | |
| The starting positions of agents and all the pieces (top left corner of 5x5 grid) | |
| ''' | |
| # # Set the starting position at the center of the grid where the gripper will be spawned | |
| # center_sq = math.ceil((self.board_size)/2) | |
| # agent_start_pos = np.array([center_sq, center_sq], dtype=np.int64) # Get start position of agent | |
| # # Use this location to check for overlaps for new pieces generated | |
| # all_start_positions = np.array([agent_start_pos]) # Initialize with agent start position, so atleast one step is taken | |
| max_tries = 100 # Maximum number of tries | |
| tries = 0 # Counter for tries | |
| flag = True | |
| while flag: | |
| all_start_positions = [] # Initialize with empty list - Piece can be spawned on center gird as well, overlapping with agent | |
| # Select a random start position for each piece | |
| tries += 1 # Increment the try counter | |
| if tries > max_tries: # Check if max tries exceeded | |
| print("Max tries exceeded - Restart Board Layout - Try increasing the board size or reducing the number of piecess") | |
| break # Exit the main while loop | |
| spawn_choices = [[x, y] for x in range(self.board_size) for y in range(self.board_size)] # Get possible spawn locations across the board | |
| for i in range(self.num_pieces): | |
| random_choice = np.random.randint(0, len(spawn_choices)) # Select a random index | |
| piece_start_pos = spawn_choices[random_choice] # Random grid mark in the specified region | |
| # Draw randomly, until a valid value is found | |
| # This ensures no overlaps between pieces and center grid (central 3x3 will always be empty) | |
| while not layout_utils.valid(self.board_size, piece_start_pos, all_start_positions): | |
| # Remove invalid starting position and select a start position again | |
| spawn_choices.remove(piece_start_pos) | |
| if not spawn_choices: # Check if all positions are exhausted | |
| flag = False | |
| break # Exit the inner while loop | |
| random_choice = np.random.randint(0, len(spawn_choices)) | |
| piece_start_pos = spawn_choices[random_choice] | |
| all_start_positions.append(piece_start_pos) | |
| if not flag: | |
| break | |
| # The search space is not exhausted and all pieces have been spawned successfully | |
| if flag: | |
| break | |
| # else try again | |
| assert len(all_start_positions) == self.num_pieces, "Number of pieces spawned is not equal to the number of pieces specified" | |
| return all_start_positions | |
| def set_board_layout(self, target_shape=None, target_colour=None, level=None): | |
| # Get all start positions for all pieces on the board | |
| all_start_positions = self.set_start_positions() | |
| # Set agent start position at the center of the board | |
| center_sq = math.ceil((self.board_size)/2) | |
| agent_start_pos = np.array([center_sq, center_sq], dtype=np.int64) # Get start position of agent | |
| grid_info = [] | |
| available_shapes = list(self.shapes) # List of available shapes | |
| available_colours = list(self.colours) # List of available colours | |
| for i in range(len(all_start_positions)): | |
| piece_position = all_start_positions[i] | |
| # Select a random shape from the available shapes | |
| piece_shape = np.random.choice(available_shapes) | |
| # Select a random colour from the available colours | |
| colour_name = np.random.choice(available_colours) | |
| # Get target_pos | |
| if i == 0: | |
| target_pos = piece_position | |
| if target_shape: | |
| piece_shape = target_shape # Overwrite target shape if specified | |
| if target_colour: | |
| colour_name = target_colour # Overwrite target colour if specified | |
| if level == "easy" or level == "sample": | |
| available_shapes.remove(piece_shape) # Remove the selected shape from the available shapes | |
| available_colours.remove(colour_name) # Remove the selected colour from the available colours | |
| piece_rotation = 0 # No rotation | |
| elif level == "medium": | |
| # Introduce rotation for medium level | |
| available_shapes.remove(piece_shape) # Remove the selected shape from the available shapes | |
| available_colours.remove(colour_name) # Remove the selected colour from the available colours | |
| piece_rotation = np.random.randint(0, 4) # Random rotation | |
| else: | |
| # Hard level, allow same shape or colour repitition, based on randomness | |
| random_value = np.random.randint(0, 2) | |
| if random_value: | |
| available_colours.remove(colour_name) # Remove the selected colour from the available colours | |
| else: | |
| available_shapes.remove(piece_shape) # Remove the selected shape from the available shapes | |
| piece_rotation = np.random.randint(0, 4) # Random rotation | |
| piece = PentominoPiece(piece_shape, piece_rotation, piece_position) | |
| piece_grids = piece.get_grid_locations() | |
| piece_region = layout_utils.get_region(piece_position, self.mapped_regions) | |
| piece_data = { | |
| "piece_grids": piece_grids, | |
| "piece_colour": colour_name, | |
| "colour_value": COLOURS[colour_name], | |
| "start_position": piece_position, | |
| "piece_shape": piece_shape, | |
| "piece_rotation": piece_rotation, | |
| "piece_region": piece_region | |
| } | |
| grid_info.append(piece_data) | |
| return agent_start_pos, target_pos, grid_info | |
| if __name__ == '__main__': | |
| board1 = BoardLayout(board_size=18, num_pieces=4, shapes=PIECE_NAMES, colours=COLOUR_NAMES, seed=640) | |
| agent_start_pos, target_pos, info = board1.set_board_layout( | |
| target_shape = 'P', | |
| target_colour = 'red', | |
| level = 'easy') | |