feat(jigsaw): Implement the Jigsaw env #147
Conversation
|
|
clement-bonnet
changed the title
Co-authored-by: Sasha <reallysasha@gmail.com>
Co-authored-by: Sasha <reallysasha@gmail.com>
Co-authored-by: Sasha <reallysasha@gmail.com>
Co-authored-by: Sasha <reallysasha@gmail.com>
| @pytest.fixture() | ||
| def piece_one_partially_placed(board_with_piece_one_placed: chex.Array) -> chex.Array: | ||
| """A 2D array of zeros where piece one has been placed partially correctly. | ||
| That is to say that there is overlap between where the piece has been placed and |
There was a problem hiding this comment.
Can you add the correct tabs throughout the codebase, please :)
| - env: snake # [bin_pack, cleaner, connector, cvrp, game_2048, job_shop, knapsack, maze, minesweeper, rubiks_cube, snake, tsp] | ||
| - env: jigsaw # [bin_pack, cleaner, connector, cvrp, game_2048, jigsaw, job_shop, knapsack, maze, minesweeper, rubiks_cube, snake, tsp] | ||
|
|
||
| agent: random # [random, a2c] |
| num_pieces: chex.Numeric # () | ||
| solved_board: chex.Array # (num_rows, num_cols) | ||
| pieces: chex.Array # (num_pieces, 3, 3) | ||
| action_mask: chex.Array # (num_pieces, num_rotations, num_rows-3, num_cols-3) |
There was a problem hiding this comment.
Are you calculating the action mask from scratch each step? If so you don't need it in the state
| from jumanji.environments.packing.jigsaw.types import State | ||
|
|
||
|
|
||
| class RewardFn(abc.ABC): |
There was a problem hiding this comment.
I think a protocol would make more sense here than a base class as it is just the callable? @clement-bonnet
| chosen_piece = rotate_piece(chosen_piece, rotation) | ||
|
|
||
| grid_piece = self._expand_piece_to_board(chosen_piece, row_idx, col_idx) | ||
| grid_mask_piece = self._get_ones_like_expanded_piece(grid_piece=grid_piece) |
There was a problem hiding this comment.
I think this works and gets rid of the extra method
| grid_mask_piece = self._get_ones_like_expanded_piece(grid_piece=grid_piece) | |
| grid_mask_piece = grid_piece == piece_idx |
| grids = batch_expand_piece_to_board(rotated_pieces, rows, cols) | ||
|
|
||
| batch_get_ones_like_expanded_piece = jax.vmap( | ||
| self._get_ones_like_expanded_piece, in_axes=(0) |
There was a problem hiding this comment.
If you delete _get_ones_like_expanded_piece then:
| self._get_ones_like_expanded_piece, in_axes=(0) | |
| lambda x: x != 0, in_axes=(0) |
|
There is a problem with the current formulation of the problem. In order to make the problem solvable in its current configuration would require the agent to have access to the solved board which makes the problem non-combinatorial. The way forward will be to rework Jigsaw as a new environment called FlatPack. This will be a 2D, discrete and flattened version of the BinPack problem with potential positive transfer to the Tetris environment since placed blocks will still interlock with each other. |
Implements the full Jigsaw environment with actor critic networks.
closes #143