Fix adroit support

takuseno · Jan 6, 2025 · 972b043 · 972b043
1 parent abb6e83
commit 972b043
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/d3rlpy/datasets.py b/d3rlpy/datasets.py
@@ -441,15 +441,17 @@ def get_d4rl(
         )
 
         # remove incompatible wrappers
-        normalized_env = env.env.env.env  # type: ignore
-        assert isinstance(
-            normalized_env, (NormalizedBoxEnv, NormalizedBoxEnvFromUtils)
-        )
-        unwrapped_env: gym.Env[Any, Any] = normalized_env.wrapped_env
-        unwrapped_env.render_mode = render_mode  # overwrite
+        wrapped_env = env.env.env.env  # type: ignore
+        if isinstance(
+            wrapped_env, (NormalizedBoxEnv, NormalizedBoxEnvFromUtils)
+        ):
+            unwrapped_env: gym.Env[Any, Any] = wrapped_env.wrapped_env
+            unwrapped_env.render_mode = render_mode  # overwrite
+        else:
+            wrapped_env.env.render_mode = render_mode  # overwrite
 
         return dataset, TimeLimit(
-            normalized_env, max_episode_steps=max_episode_steps
+            wrapped_env, max_episode_steps=max_episode_steps
         )
     except ImportError as e:
         raise ImportError(