ARISE-Initiative
diff --git a/‎robomimic/envs/env_robosuite.py‎
Lines changed: 8 additions & 2 deletions b/‎robomimic/envs/env_robosuite.py‎
Lines changed: 8 additions & 2 deletions
@@ -7,12 +7,18 @@
 import numpy as np
 from copy import deepcopy
 
-import mujoco_py
 import robosuite
 
 import robomimic.utils.obs_utils as ObsUtils
 import robomimic.envs.env_base as EB
 
+# protect against missing mujoco-py module, since robosuite might be using mujoco-py or DM backend
+try:
+    import mujoco_py
+    MUJOCO_EXCEPTIONS = [mujoco_py.builder.MujocoException]
+except ImportError:
+    MUJOCO_EXCEPTIONS = []
+
 
 class EnvRobosuite(EB.EnvBase):
     """Wrapper class for robosuite environments (https://github.com/ARISE-Initiative/robosuite)"""
@@ -375,7 +381,7 @@ def rollout_exceptions(self):
         that the entire training run doesn't crash because of a bad policy that causes unstable
         simulation computations.
         """
-        return (mujoco_py.builder.MujocoException)
+        return tuple(MUJOCO_EXCEPTIONS)
 
     def __repr__(self):
         """