Repository: epignatelli/navix (language: python)

Issue:
`Environment.create` fails with `kwargs` already in the signature
Reproduce:

```
import navix as nx

nx.make("Navix-KeyCorridorS3R2-v0", observation_fn=nx.observations.rgb)
```


Agents walking through walls/lava?
Hi!

I have been training ppo agents in navix environments and I keep seeing that agents walk through walls/lava. I'm using version 0.6.14 of Navix. I'm using a PPO agent from PureJaxRL (purejaxrl/ppo_minigrid.py). I train the agent, and then load a network to render the behaviour, and I get the following:
```python
from ppo_minigrid import *
from navix import observations, rewards, terminations, register_env
import navix

config = {
    "LR": 2.5e-4,
    "NUM_ENVS": 16,
    "NUM_STEPS": 256,
    "TOTAL_TIMESTEPS": 2e6,
    "UPDATE_EPOCHS": 1,
    "NUM_MINIBATCHES": 8,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.05,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "Navix-LavaGapS7-v0",
    "ANNEAL_LR": True,
    "DEBUG": True,
}

rng = jax.random.PRNGKey(30)
train_jit = jax.jit(make_train(config))
out = train_jit(rng)

import matplotlib.pyplot as plt
import jax
import jax.numpy as jnp

env = NavixGymnaxWrapper("Navix-LavaGapS7-v0")
env = FlattenObservationWrapper(env)
env = LogWrapper(env)
key = jax.random.PRNGKey(30)
key, _key = jax.random.split(key)
obs, state = env.reset(_key, None)

def render(obs, title):
    plt.imshow(obs)
    plt.title(title)
    plt.axis('off')
    plt.savefig("image.png")

render(observations.rgb(state.env_state.state), "Initial observation")
network = ActorCritic(env.action_space(None).n, activation=config["ACTIVATION"])
params = out['runner_state'][0].params

for t in range(100):
    rng, _rng = jax.random.split(rng)
    pi, _ = network.apply(params, obs)
    action = pi.sample(seed=_rng)
    print(action)
    obs, state, reward, done, info = env.step(rng, state, action, None)
    render(observations.rgb(state.env_state.state), f"Step {t}")
    input('Continue...')
```

The initial observation:
![initial_obs](https://github.com/epignatelli/navix/assets/34233348/b9ce675f-e313-4388-a4c8-af74338912e3)

After 2 steps:
![t2](https://github.com/epignatelli/navix/assets/34233348/db92e443-afd6-45eb-ba92-6c96079fe21b)

After 3 steps:
![t3](https://github.com/epignatelli/navix/assets/34233348/8faeb066-6f43-4c77-a109-911d5c9dfb62)

I had the same with the DoorKey environments, where agents simply walk through the dividing wall. Any ideas of what is happening?

Thanks!

Daniel



Discussion / hints:

Hey @danieljarne! Thank you for spotting this -- it's definitely a bug introduced by some recent commit.
Let me have a look into it today and I'll get back to you asap.

Thank you for flagging this btw -- that's key at this development stage, and please do continue opening issues, should you find something not working properly or a missing feature.
I just had a look, and spotted a bug.
It should be fixed in #84, let me know if that works for you.

You can install the dev branch with:
```sh
pip install git+https://github.com/epignatelli/navix.git@fix-walkable
```

A quick not: walking through walls is not expected.
However, walking on lava is, as we use it to trigger events such as rewards or terminations.

Now `DoorKey` (and others) should correctly prevent that, and `LavaGap` has a default `on_lava_fall` termination condition.

Fix patch:
diff --git a/navix/_version.py b/navix/_version.py
index 85ee2ad..3812a70 100644
--- a/navix/_version.py
+++ b/navix/_version.py
@@ -18,5 +18,5 @@
 # under the License.
 
 
-__version__ = "0.6.16"
+__version__ = "0.6.17"
 __version_info__ = tuple(int(i) for i in __version__.split(".") if i.isdigit())
diff --git a/navix/environments/empty.py b/navix/environments/empty.py
index 6bd9a56..72805e7 100644
--- a/navix/environments/empty.py
+++ b/navix/environments/empty.py
@@ -116,10 +116,10 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
         height=8,
         width=8,
         random_start=False,
-        *args,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
         **kwargs,
     ),
 )
@@ -129,10 +129,10 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
         height=16,
         width=16,
         random_start=False,
-        *args,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
         **kwargs,
     ),
 )
@@ -142,10 +142,10 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
         height=5,
         width=5,
         random_start=True,
-        *args,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
         **kwargs,
     ),
 )
@@ -155,10 +155,10 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
         height=6,
         width=6,
         random_start=True,
-        *args,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
         **kwargs,
     ),
 )
@@ -168,10 +168,10 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
         height=8,
         width=8,
         random_start=True,
-        *args,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
         **kwargs,
     ),
 )
@@ -181,10 +181,10 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
         height=16,
         width=16,
         random_start=True,
-        *args,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
         **kwargs,
     ),
 )
diff --git a/navix/environments/key_corridor.py b/navix/environments/key_corridor.py
index 39ef5a2..c446e8a 100644
--- a/navix/environments/key_corridor.py
+++ b/navix/environments/key_corridor.py
@@ -149,72 +149,72 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
 register_env(
     "Navix-KeyCorridorS3R1-v0",
     lambda *args, **kwargs: KeyCorridor.create(
-        *args,
-        **kwargs,
         height=3,
         width=7,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
+        **kwargs,
     ),
 )
 register_env(
     "Navix-KeyCorridorS3R2-v0",
     lambda *args, **kwargs: KeyCorridor.create(
-        *args,
-        **kwargs,
         height=5,
         width=7,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
+        **kwargs,
     ),
 )
 register_env(
     "Navix-KeyCorridorS3R3-v0",
     lambda *args, **kwargs: KeyCorridor.create(
-        *args,
-        **kwargs,
         height=7,
         width=7,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
+        **kwargs,
     ),
 )
 register_env(
     "Navix-KeyCorridorS4R3-v0",
     lambda *args, **kwargs: KeyCorridor.create(
-        *args,
-        **kwargs,
         height=10,
         width=10,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
+        **kwargs,
     ),
 )
 register_env(
     "Navix-KeyCorridorS5R3-v0",
     lambda *args, **kwargs: KeyCorridor.create(
-        *args,
-        **kwargs,
         height=13,
         width=13,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
+        **kwargs,
     ),
 )
 register_env(
     "Navix-KeyCorridorS6R3-v0",
     lambda *args, **kwargs: KeyCorridor.create(
-        *args,
-        **kwargs,
         height=16,
         width=16,
         observation_fn=kwargs.pop("observation_fn", observations.symbolic),
         reward_fn=kwargs.pop("reward_fn", rewards.on_goal_reached),
         termination_fn=kwargs.pop("termination_fn", terminations.on_goal_reached),
+        *args,
+        **kwargs,
     ),
 )
diff --git a/navix/rendering/registry.py b/navix/rendering/registry.py
index 7872e52..91bb7dc 100644
--- a/navix/rendering/registry.py
+++ b/navix/rendering/registry.py
@@ -31,7 +31,7 @@
 SPRITES_DIR = os.path.normpath(
     os.path.join(__file__, "..", "..", "..", "assets", "sprites")
 )
-MIN_TILE_SIZE = 8
+MIN_TILE_SIZE = 32
 TILE_SIZE = MIN_TILE_SIZE
 
 


Test patch:
diff --git a/tests/test_environments.py b/tests/test_environments.py
index d88fb47..30bcfc2 100644
--- a/tests/test_environments.py
+++ b/tests/test_environments.py
@@ -5,7 +5,12 @@
 
 def test_room():
     def f():
-        env = nx.environments.Room.create(height=3, width=3, max_steps=8)
+        env = nx.environments.Room.create(
+            height=3,
+            width=3,
+            max_steps=8,
+            observation_fn=nx.observations.symbolic_first_person,
+        )
         key = jax.random.PRNGKey(4)
         reset = jax.jit(env._reset)
         step = jax.jit(env.step)
@@ -35,7 +40,12 @@ def f():
 
 def test_keydoor():
     def f():
-        env = nx.environments.DoorKey.create(height=5, width=10, max_steps=8)
+        env = nx.environments.DoorKey.create(
+            height=5,
+            width=10,
+            max_steps=8,
+            observation_fn=nx.observations.symbolic_first_person,
+        )
         key = jax.random.PRNGKey(1)
         reset = jax.jit(env._reset)
         step = jax.jit(env.step)