improving baselines: extended obs, better embedder, faster with bf16 (#…

…32)

* extended obs wip

* extended obs wip

* extended obs with rules wip

* additional comment

* fix ruff action

* fix render wrapper

* updated baselines with dir obs

* embedder

* updated configs with more SPS

* fix typo

* added bf16 support

.gitignore

@@ -174,4 +174,4 @@ g
 # will remove later
 scripts/*testing*
 configs/wip
-scripts/eval_example.py
+scripts/eval_example.py

configs/meta.yaml

@@ -6,7 +6,7 @@ num_envs: 16384
 num_steps_per_env: 12800
 num_steps_per_update: 256
 update_epochs: 1
-num_minibatches: 32
+num_minibatches: 16
 eval_num_envs: 16384
 eval_num_episodes: 25
 train_seed: 5

configs/single.yaml

@@ -6,6 +6,6 @@ total_timesteps: 1_000_000_000
 num_envs: 16384
 num_steps: 256
 update_epochs: 1
-num_minibatches: 8
+num_minibatches: 16
 eval_episodes: 512
 train_seed: 5

src/xminigrid/__init__.py

@@ -2,11 +2,11 @@
 from .registration import make, register, registered_environments
 
 # TODO: add __all__
-__version__ = "0.8.0"
+__version__ = "0.9.0"
 
 # ---------- XLand-MiniGrid environments ----------
 
-# WARN: TMP, only for FPS measurements
+# WARN: TMP, only for FPS measurements, will remove later
 # register(
 #     id="MiniGrid-1Rules",
 #     entry_point="xminigrid.envs.xland_tmp:XLandMiniGrid",
@@ -79,9 +79,6 @@
 #     width=64,
 # )
 
-
-# TODO: reconsider grid sizes and time limits after the benchmarks are generated.
-#  Should be enough space for initial tiles even in the hardest setting
 register(
     id="XLand-MiniGrid-R1-9x9",
     entry_point="xminigrid.envs.xland:XLandMiniGrid",