Fix bug in navigation task

[jysdoran]

When the goal probability is one the whole array should be true instead of false to allow for rewards

[epignatelli]

Thanks, @jysdoran, for spotting the issue! Really appreciated the time you dedicated to contribute.

I found another bug while reviewing this -- added some comments.

Would you mind adding also a unittest to make sure this remains fixed in the future?
I put together a quick one but feel free to improve in it if you find a better one.

I can pick it up if you don't have time!

def test_navigation():
    """Unittest for https://github.com/epignatelli/navix/pull/47"""
    height = 10
    width = 10
    grid = jnp.zeros((height - 2, width - 2), dtype=jnp.int32)
    grid = jnp.pad(grid, 1, mode="constant", constant_values=-1)

    players = Player(
        position=jnp.asarray((1, 1)), direction=jnp.asarray(0), pocket=EMPTY_POCKET_ID
    )
    goals = Goal(position=jnp.asarray([(1, 1), (1, 1)]), probability=jnp.asarray([0.0, 0.0]))
    keys = Key(position=jnp.asarray((2, 2)), id=jnp.asarray(0))
    doors = Door(
        position=jnp.asarray([(1, 5), (1, 6)]),
        direction=jnp.asarray((0, 2)),
        requires=jnp.asarray((0, 0)),
        open=jnp.asarray((False, True)),
    )

    entities = {
        Entities.PLAYER: players[None],
        Entities.GOAL: goals,
        Entities.KEY: keys[None],
        Entities.DOOR: doors,
    }

    state = nx.entities.State(
        key=jax.random.PRNGKey(0),
        grid=grid,
        cache=nx.graphics.RenderingCache.init(grid),
        entities=entities,
    )
    action = jnp.asarray(0)
    reward = nx.tasks.navigation(state, action, state)
    assert jnp.array_equal(reward, jnp.asarray(0.0))

[epignatelli]

Thanks @jysdoran for the PR again, I just pushed the other fixes, add the unitest and bumped the version to 0.3.11.

The fix should come online in a few minutes.