From a7e7e8c9bdde2cb51d99e0493b0abfedbcdf2840 Mon Sep 17 00:00:00 2001
From: Gaiejj <524339208@qq.com>
Date: Thu, 2 May 2024 14:06:32 +0800
Subject: [PATCH] refactor: refactor making warning

---
 omnisafe/algorithms/algo_wrapper.py   |  3 +--
 omnisafe/envs/__init__.py             |  5 +----
 omnisafe/envs/safety_isaac_gym_env.py | 13 +++++++++++--
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/omnisafe/algorithms/algo_wrapper.py b/omnisafe/algorithms/algo_wrapper.py
index db8d25a73..b50504e1b 100644
--- a/omnisafe/algorithms/algo_wrapper.py
+++ b/omnisafe/algorithms/algo_wrapper.py
@@ -144,8 +144,7 @@ def _init_checks(self) -> None:
         assert self.cfgs.train_cfgs.parallel > 0, 'parallel must be greater than 0!'
         assert (
             self.env_id in support_envs()
-        ), f"{self.env_id} doesn't exist. Please choose from {support_envs()}.\
-            \nIf you are using Safe Isaac Gym environments, please install Isaac Gym first."
+        ), f"{self.env_id} doesn't exist. Please choose from {support_envs()}."
 
     def _init_algo(self) -> None:
         """Initialize the algorithm."""
diff --git a/omnisafe/envs/__init__.py b/omnisafe/envs/__init__.py
index 852377d11..645bf7208 100644
--- a/omnisafe/envs/__init__.py
+++ b/omnisafe/envs/__init__.py
@@ -23,7 +23,4 @@
 from omnisafe.envs.mujoco_env import MujocoEnv
 from omnisafe.envs.safety_gymnasium_env import SafetyGymnasiumEnv
 from omnisafe.envs.safety_gymnasium_modelbased import SafetyGymnasiumModelBased
-
-
-with suppress(ImportError):
-    from omnisafe.envs.safety_isaac_gym_env import SafetyIsaacGymEnv
+from omnisafe.envs.safety_isaac_gym_env import SafetyIsaacGymEnv
diff --git a/omnisafe/envs/safety_isaac_gym_env.py b/omnisafe/envs/safety_isaac_gym_env.py
index 6ea4e0b74..e9d7592e9 100644
--- a/omnisafe/envs/safety_isaac_gym_env.py
+++ b/omnisafe/envs/safety_isaac_gym_env.py
@@ -23,7 +23,13 @@
 
 from omnisafe.envs.core import CMDP, env_register
 from omnisafe.typing import DEVICE_CPU
-from omnisafe.utils.isaac_gym_utils import make_isaac_gym_env
+
+
+ISAAC_GYM_AVAILABLE = True
+try:
+    from omnisafe.utils.isaac_gym_utils import make_isaac_gym_env
+except ImportError:
+    ISAAC_GYM_AVAILABLE = False
 
 
 @env_register
@@ -64,7 +70,10 @@ def __init__(
         super().__init__(env_id)
         self._num_envs = num_envs
         self._device = torch.device(device)
-        self._env = make_isaac_gym_env(env_id=env_id, device=device, num_envs=num_envs)
+        if ISAAC_GYM_AVAILABLE:
+            self._env = make_isaac_gym_env(env_id=env_id, device=device, num_envs=num_envs)
+        else:
+            raise ImportError('Please install Isaac Gym to use Safe Isaac Gym!')
         self._action_space = self._env.action_space
         self._observation_space = self._env.observation_space
         self.need_evaluation = False