Merge pull request #13 from KraHsu/fix/play-viewer-crashes

Fix viewer crashes on first play and on window close

Author: KraHsu

src/genelab/cache.py

@@ -14,5 +14,6 @@ def ensure_project_cache() -> None:
 
     XDG_CACHE_DIR.mkdir(parents=True, exist_ok=True)
     MATPLOTLIB_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    (XDG_CACHE_DIR / "genesis").mkdir(parents=True, exist_ok=True)
     os.environ.setdefault("XDG_CACHE_HOME", str(XDG_CACHE_DIR))
     os.environ.setdefault("MPLCONFIGDIR", str(MATPLOTLIB_CACHE_DIR))

src/genelab/rl/runner.py

@@ -166,6 +166,7 @@ def play_task(
     env = _build_env(env_cfg)
     wrapped = RslRlVecEnvWrapper(env, clip_actions=None)
 
+    import genesis as gs  # type: ignore[import-not-found]
     import torch
 
     action_shape = (env.num_envs, wrapped.num_actions)
@@ -213,7 +214,12 @@ def _random_policy(_obs: Any) -> "torch.Tensor":
         while True:
             with torch.inference_mode():
                 actions = policy(obs)
-            obs, _, _, _ = wrapped.step(actions)
+            try:
+                obs, _, _, _ = wrapped.step(actions)
+            except gs.GenesisException as exc:
+                if str(exc) != "Viewer closed.":
+                    raise
+                break
             step += 1
             if max_steps is not None and step >= max_steps:
                 break