Restructure mbrl.util package

mohakbhardwaj · Oct 8, 2021 · 1e07015 · 1e07015
1 parent cfde559
commit 1e07015
Show file tree

Hide file tree

Showing 7 changed files with 65 additions and 464 deletions.
diff --git a/mbrl/diagnostics/eval_model_on_dataset.py b/mbrl/diagnostics/eval_model_on_dataset.py
@@ -12,7 +12,6 @@
 
 import mbrl.util
 import mbrl.util.common
-import mbrl.util.mujoco
 
 
 class DatasetEvaluator:

diff --git a/mbrl/diagnostics/visualize_model_preds.py b/mbrl/diagnostics/visualize_model_preds.py
@@ -16,7 +16,6 @@
 import mbrl.models
 import mbrl.planning
 import mbrl.util.common
-import mbrl.util.mujoco
 
 VisData = Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]
 

diff --git a/mbrl/util/__init__.py b/mbrl/util/__init__.py
@@ -48,11 +48,11 @@ def create_handler(cfg: Union[Dict, omegaconf.ListConfig, omegaconf.DictConfig])
 
     target = cfg.overrides.env_cfg.get("_target_")
     if "pybulletgym" in target:
-        from mbrl.util.pybullet_handler import PybulletEnvHandler
+        from mbrl.util.pybullet import PybulletEnvHandler
 
         return PybulletEnvHandler()
     elif "mujoco" in target:
-        from mbrl.util.mujoco_handler import MujocoEnvHandler
+        from mbrl.util.mujoco import MujocoEnvHandler
 
         return MujocoEnvHandler()
     else:
@@ -84,12 +84,12 @@ def create_handler_from_str(env_name: str):
     if "dmcontrol___" in env_name:
         raise NotImplementedError
     elif "pybulletgym___" in env_name:
-        from .pybullet_handler import PybulletEnvHandler
+        from mbrl.util.pybullet import PybulletEnvHandler
 
         return PybulletEnvHandler()
     elif "gym___" in env_name:
         raise NotImplementedError
     else:
-        from .mujoco_handler import MujocoEnvHandler
+        from mbrl.util.mujoco import MujocoEnvHandler
 
         return MujocoEnvHandler()
diff --git a/mbrl/util/env_handler.py → mbrl/util/env.py b/mbrl/util/env_handler.py → mbrl/util/env.py