fix evaluation in standalone (#14)

* fix evaluation in standalone ppo
dunnolab · Mar 24, 2024 · 501a636 · 501a636
1 parent 1f5bc0c
commit 501a636
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/examples/train_meta_standalone.ipynb b/examples/train_meta_standalone.ipynb
@@ -727,6 +727,13 @@
     "env, env_params = xminigrid.make(config.env_id)\n",
     "env = GymAutoResetWrapper(env)\n",
     "\n",
+    "# enabling image observations if needed\n",
+    "if config.img_obs:\n",
+    "    from xminigrid.experimental.img_obs import RGBImgObservationWrapper\n",
+    "\n",
+    "    env = RGBImgObservationWrapper(env)\n",
+    "    \n",
+    "\n",
     "ruleset = xminigrid.load_benchmark(config.benchmark_id).get_ruleset(ruleset_id=0)\n",
     "env_params = env_params.replace(ruleset=ruleset)\n",
     "\n",
@@ -739,6 +746,7 @@
     "    rnn_hidden_dim=config.rnn_hidden_dim,\n",
     "    rnn_num_layers=config.rnn_num_layers,\n",
     "    head_hidden_dim=config.head_hidden_dim,\n",
+    "    img_obs=config.img_obs,\n",
     ")\n",
     "\n",
     "# jitting all functions\n",

diff --git a/examples/train_single_standalone.ipynb b/examples/train_single_standalone.ipynb
@@ -690,6 +690,19 @@
     "env, env_params = xminigrid.make(config.env_id)\n",
     "env = GymAutoResetWrapper(env)\n",
     "\n",
+    "# for single-task XLand environments\n",
+    "if config.benchmark_id is not None:\n",
+    "    assert \"XLand-MiniGrid\" in config.env_id, \"Benchmarks should be used only with XLand environments.\"\n",
+    "    assert config.ruleset_id is not None, \"Ruleset ID should be specified for benchmarks usage.\"\n",
+    "    benchmark = xminigrid.load_benchmark(config.benchmark_id)\n",
+    "    env_params = env_params.replace(ruleset=benchmark.get_ruleset(config.ruleset_id))\n",
+    "\n",
+    "# enabling image observations if needed\n",
+    "if config.img_obs:\n",
+    "    from xminigrid.experimental.img_obs import RGBImgObservationWrapper\n",
+    "\n",
+    "    env = RGBImgObservationWrapper(env)\n",
+    "\n",
     "# you can use train_state from the final runner_state also\n",
     "# we just demo here how to do it if you loaded params from the checkpoint\n",
     "params = train_info[\"runner_state\"][1].params\n",
@@ -699,6 +712,7 @@
     "    rnn_hidden_dim=config.rnn_hidden_dim,\n",
     "    rnn_num_layers=config.rnn_num_layers,\n",
     "    head_hidden_dim=config.head_hidden_dim,\n",
+    "    img_obs=config.img_obs,\n",
     ")\n",
     "\n",
     "# jitting all functions\n",