general improvements (#19)

src/xminigrid/__init__.py

@@ -2,7 +2,7 @@
 from .registration import make, register, registered_environments
 
 # TODO: add __all__
-__version__ = "0.7.0"
+__version__ = "0.7.1"
 
 # ---------- XLand-MiniGrid environments ----------
 

src/xminigrid/core/goals.py

@@ -10,6 +10,7 @@
 from .grid import equal, get_neighbouring_tiles, pad_along_axis
 
 MAX_GOAL_ENCODING_LEN = 4 + 1  # for idx
+NUM_GOALS = 15
 
 
 def check_goal(

src/xminigrid/core/rules.py

@@ -11,6 +11,7 @@
 from .grid import equal, get_neighbouring_tiles, pad_along_axis
 
 MAX_RULE_ENCODING_LEN = 6 + 1  # +1 for idx
+NUM_RULES = 12
 
 
 # this is very costly, will evaluate all rules under vmap. Submit a PR if you know how to do it better!

src/xminigrid/envs/minigrid/memory.py

@@ -109,7 +109,7 @@ def step(
         self, params: EnvParams, timestep: TimeStep[MemoryEnvCarry], action: IntOrArray
     ) -> TimeStep[MemoryEnvCarry]:
         # disabling pick_up action
-        action = jax.lax.select(jnp.equal(action, 3), 5, action)
+        action = jax.lax.select(jnp.equal(action, 3), jnp.asarray(5, dtype=jnp.uint8), action)
         new_grid, new_agent, _ = take_action(timestep.state.grid, timestep.state.agent, action)
 
         new_state = timestep.state.replace(grid=new_grid, agent=new_agent, step_num=timestep.state.step_num + 1)

src/xminigrid/manual_control.py

@@ -116,12 +116,10 @@ def start(self) -> None:
     def step(self, action: int) -> None:
         self.timestep = self._step(self.env_params, self.timestep, action)
         print(
-            "StepType: ",
-            self.timestep.step_type,
-            "Discount: ",
-            self.timestep.discount,
-            "Reward: ",
-            self.timestep.reward,
+            f"Step: {self.timestep.state.step_num} | ",
+            f"StepType: {self.timestep.step_type} | ",
+            f"Discount: {self.timestep.discount} | ",
+            f"Reward: {self.timestep.reward}",
         )
         self.render()
 
@@ -135,12 +133,10 @@ def reset(self) -> None:
         self.timestep = self._reset(self.env_params, reset_key)
         self.render()
         print(
-            "StepType: ",
-            self.timestep.step_type,
-            "Discount: ",
-            self.timestep.discount,
-            "Reward: ",
-            self.timestep.reward,
+            f"Step: {self.timestep.state.step_num} |",
+            f"StepType: {self.timestep.step_type} |",
+            f"Discount: {self.timestep.discount} |",
+            f"Reward: {self.timestep.reward}",
         )
 
     def key_handler(self, event: Event) -> None:

training/eval.py

@@ -15,14 +15,14 @@
 
 def main():
     orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
-    checkpoint = orbax_checkpointer.restore("../xland-minigrid-data/XLand-MiniGrid-R1-9x9-Trivial-v0-5B-gamma99")
+    checkpoint = orbax_checkpointer.restore("../xland-minigrid-data/checkpoints")
     config = checkpoint["config"]
     params = checkpoint["params"]
 
     env, env_params = xminigrid.make("XLand-MiniGrid-R1-9x9")
     env = GymAutoResetWrapper(env)
 
-    ruleset = xminigrid.load_benchmark("Trivial").get_ruleset(3)
+    ruleset = xminigrid.load_benchmark("trivial-1m").get_ruleset(3)
     env_params = env_params.replace(ruleset=ruleset)
 
     model = ActorCriticRNN(