Replace mujoco_py with dm_control.

Author: Ilya Kostrikov

gym/envs/mujoco/mujoco_env.py

@@ -9,10 +9,10 @@
 import gym
 
 try:
-    import mujoco_py
+    import dm_control.mujoco as dm_mujoco
 except ImportError as e:
     raise error.DependencyNotInstalled(
-        "{}. (HINT: you need to install mujoco_py, and also perform the setup instructions here: https://github.com/openai/mujoco-py/.)".format(
+        "{}. (HINT: you need to install dm_control)".format(
             e
         )
     )
@@ -51,8 +51,8 @@ def __init__(self, model_path, frame_skip):
         if not path.exists(fullpath):
             raise OSError(f"File {fullpath} does not exist")
         self.frame_skip = frame_skip
-        self.model = mujoco_py.load_model_from_path(fullpath)
-        self.sim = mujoco_py.MjSim(self.model)
+        self.sim = dm_mujoco.Physics.from_xml_path(fullpath)
+        self.model = self.sim.model
         self.data = self.sim.data
         self.viewer = None
         self._viewers = {}
@@ -111,11 +111,10 @@ def reset(self, seed: Optional[int] = None):
 
     def set_state(self, qpos, qvel):
         assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,)
-        old_state = self.sim.get_state()
-        new_state = mujoco_py.MjSimState(
-            old_state.time, qpos, qvel, old_state.act, old_state.udd_state
-        )
-        self.sim.set_state(new_state)
+        state = self.sim.get_state()
+        state[:self.model.nq] = qpos
+        state[self.model.nq:self.model.nq+self.model.nv] = qvel
+        self.sim.set_state(state)
         self.sim.forward()
 
     @property
@@ -138,55 +137,18 @@ def render(
         camera_id=None,
         camera_name=None,
     ):
-        if mode == "rgb_array" or mode == "depth_array":
-            if camera_id is not None and camera_name is not None:
-                raise ValueError(
-                    "Both `camera_id` and `camera_name` cannot be"
-                    " specified at the same time."
-                )
-
-            no_camera_specified = camera_name is None and camera_id is None
-            if no_camera_specified:
-                camera_name = "track"
-
-            if camera_id is None and camera_name in self.model._camera_name2id:
-                camera_id = self.model.camera_name2id(camera_name)
-
-            self._get_viewer(mode).render(width, height, camera_id=camera_id)
-
         if mode == "rgb_array":
-            # window size used for old mujoco-py:
-            data = self._get_viewer(mode).read_pixels(width, height, depth=False)
-            # original image is upside-down, so flip it
-            return data[::-1, :, :]
-        elif mode == "depth_array":
-            self._get_viewer(mode).render(width, height)
-            # window size used for old mujoco-py:
-            # Extract depth part of the read_pixels() tuple
-            data = self._get_viewer(mode).read_pixels(width, height, depth=True)[1]
-            # original image is upside-down, so flip it
-            return data[::-1, :]
-        elif mode == "human":
-            self._get_viewer(mode).render()
+            camera_id = camera_id or 0
+            return self.sim.render(height, width, camera_id)
+        else:
+            raise NotImplemented
 
     def close(self):
         if self.viewer is not None:
             # self.viewer.finish()
             self.viewer = None
             self._viewers = {}
 
-    def _get_viewer(self, mode):
-        self.viewer = self._viewers.get(mode)
-        if self.viewer is None:
-            if mode == "human":
-                self.viewer = mujoco_py.MjViewer(self.sim)
-            elif mode == "rgb_array" or mode == "depth_array":
-                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
-
-            self.viewer_setup()
-            self._viewers[mode] = self.viewer
-        return self.viewer
-
     def get_body_com(self, body_name):
         return self.data.get_body_xpos(body_name)
 

gym/envs/mujoco/pusher.py

@@ -2,8 +2,6 @@
 from gym import utils
 from gym.envs.mujoco import mujoco_env
 
-import mujoco_py
-
 
 class PusherEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(self):

requirements.txt

@@ -1,7 +1,7 @@
 ale-py~=0.7
 opencv-python>=3.0
 box2d-py==2.3.5
-mujoco_py>=1.50, <2.0
+dm_control>=0.0.403778684
 scipy>=1.4.1
 numpy>=1.18.0
 pyglet>=1.4.0

setup.py

@@ -14,7 +14,7 @@
     "accept-rom-license": ["autorom[accept-rom-license]~=0.4.2"],
     "box2d": ["box2d-py==2.3.5", "pyglet>=1.4.0"],
     "classic_control": ["pyglet>=1.4.0"],
-    "mujoco": ["mujoco_py>=1.50, <2.0"],
+    "mujoco": ["dm_control>=0.0.403778684"],
     "toy_text": ["scipy>=1.4.1"],
     "other": ["lz4>=3.1.0", "opencv-python>=3.0"],
 }

tests/envs/spec_list.py

@@ -11,7 +11,7 @@
 skip_mujoco = not (os.environ.get("MUJOCO_KEY"))
 if not skip_mujoco:
     try:
-        import mujoco_py
+        import dm_control
     except ImportError:
         skip_mujoco = True