MuZero.action_space_size(::Type{TicTacToe}) = 9
MuZero.is_terminal(game::TicTacToe) = game |> MuZero.state |> is_terminal
MuZero.get_state(g::TicTacToe) = MuZero.state(g)
state(g::TicTacToe) = g.states[end]

function make_tictactoe_config()
    make_board_game_config(
        TicTacToe, 
        network_type=Incremental,
        max_moves=9, 
        root_dirichlet_alpha=0.3, 
        action_space_size=9, # Number of all possible actions in the game 
        support_size = 10, # Value and reward are scaled (with almost sqrt) and encoded on a vector with a range of -support_size to support_size. Choose it so that support_size <= sqrt(max(abs(discounted reward)))
        observation_shape = (3, 3, 3),
        stacked_observations = 0, # Number of previous observations and previous actions to add to the current observation
        hidden_state_size = 32, # The 1-dimensional size of the encoded game state
        checkpoint_interval = 100, #  Number of training steps before using the model for self-playing
        num_actors = 3,
        weight_decay = 1e-3,
        lr_init=0.002f0,
        save_interval = 10000,
        num_unroll_steps=1,
    )
end

MuZero.num_cols(::Type{TicTacToeState}) = 3
MuZero.num_rows(::Type{TicTacToeState}) = 3
MuZero.getstatetype(::Type{TicTacToe}) = TicTacToeState