Fix: add mujoco render arguments to init (#2891)

* fix: add render_mode getter to Wrappers

* fix: add render args to mujoco init

* reformat

* add type hints

gym/envs/mujoco/ant.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,9 +5,9 @@ from gym.envs.mujoco import mujoco_env
 
 
 class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self, "ant.xml", 5, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, "ant.xml", 5, mujoco_bindings="mujoco_py", **kwargs
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/ant_v3.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -13,7 +11,6 @@ DEFAULT_CAMERA_CONFIG = {
 class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="ant.xml",
         ctrl_cost_weight=0.5,
         contact_cost_weight=5e-4,
@@ -23,6 +20,7 @@ class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         contact_force_range=(-1.0, 1.0),
         reset_noise_scale=0.1,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -41,7 +39,9 @@ class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(self, xml_file, 5, mujoco_bindings="mujoco_py")
+        mujoco_env.MujocoEnv.__init__(
+            self, xml_file, 5, mujoco_bindings="mujoco_py", **kwargs
+        )
 
     @property
     def healthy_reward(self):

gym/envs/mujoco/ant_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -166,7 +164,6 @@ class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="ant.xml",
         ctrl_cost_weight=0.5,
         use_contact_forces=False,
@@ -177,6 +174,7 @@ class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         contact_force_range=(-1.0, 1.0),
         reset_noise_scale=0.1,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -197,7 +195,7 @@ class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(self, xml_file, 5, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, xml_file, 5, **kwargs)
 
     @property
     def healthy_reward(self):

gym/envs/mujoco/half_cheetah.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,13 +5,9 @@ from gym.envs.mujoco import mujoco_env
 
 
 class HalfCheetahEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self,
+            self, "half_cheetah.xml", 5, mujoco_bindings="mujoco_py", **kwargs
-            "half_cheetah.xml",
-            5,
-            render_mode=render_mode,
-            mujoco_bindings="mujoco_py",
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/half_cheetah_v3.py

@@ -1,7 +1,5 @@
 __credits__ = ["Rushiv Arora"]
 
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -15,12 +13,12 @@ DEFAULT_CAMERA_CONFIG = {
 class HalfCheetahEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="half_cheetah.xml",
         forward_reward_weight=1.0,
         ctrl_cost_weight=0.1,
         reset_noise_scale=0.1,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -35,7 +33,7 @@ class HalfCheetahEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         )
 
         mujoco_env.MujocoEnv.__init__(
-            self, xml_file, 5, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, xml_file, 5, mujoco_bindings="mujoco_py", **kwargs
         )
 
     def control_cost(self, action):

gym/envs/mujoco/half_cheetah_v4.py

@@ -1,7 +1,5 @@
 __credits__ = ["Rushiv Arora"]
 
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -151,11 +149,11 @@ class HalfCheetahEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         forward_reward_weight=1.0,
         ctrl_cost_weight=0.1,
         reset_noise_scale=0.1,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -169,9 +167,7 @@ class HalfCheetahEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(
+        mujoco_env.MujocoEnv.__init__(self, "half_cheetah.xml", 5, **kwargs)
-            self, "half_cheetah.xml", 5, render_mode=render_mode
-        )
 
     def control_cost(self, action):
         control_cost = self._ctrl_cost_weight * np.sum(np.square(action))

gym/envs/mujoco/hopper.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,9 +5,9 @@ from gym.envs.mujoco import mujoco_env
 
 
 class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self, "hopper.xml", 4, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, "hopper.xml", 4, mujoco_bindings="mujoco_py", **kwargs
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/hopper_v3.py

@@ -1,7 +1,5 @@
 __credits__ = ["Rushiv Arora"]
 
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -18,7 +16,6 @@ DEFAULT_CAMERA_CONFIG = {
 class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="hopper.xml",
         forward_reward_weight=1.0,
         ctrl_cost_weight=1e-3,
@@ -29,6 +26,7 @@ class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         healthy_angle_range=(-0.2, 0.2),
         reset_noise_scale=5e-3,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -50,7 +48,7 @@ class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         )
 
         mujoco_env.MujocoEnv.__init__(
-            self, xml_file, 4, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, xml_file, 4, mujoco_bindings="mujoco_py", **kwargs
         )
 
     @property

gym/envs/mujoco/hopper_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -142,7 +140,6 @@ class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         forward_reward_weight=1.0,
         ctrl_cost_weight=1e-3,
         healthy_reward=1.0,
@@ -152,6 +149,7 @@ class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         healthy_angle_range=(-0.2, 0.2),
         reset_noise_scale=5e-3,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -172,7 +170,7 @@ class HopperEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(self, "hopper.xml", 4, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, "hopper.xml", 4, **kwargs)
 
     @property
     def healthy_reward(self):

gym/envs/mujoco/humanoid.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -13,13 +11,9 @@ def mass_center(model, sim):
 
 
 class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self,
+            self, "humanoid.xml", 5, mujoco_bindings="mujoco_py", **kwargs
-            "humanoid.xml",
-            5,
-            render_mode=render_mode,
-            mujoco_bindings="mujoco_py",
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/humanoid_v3.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -22,7 +20,6 @@ def mass_center(model, sim):
 class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="humanoid.xml",
         forward_reward_weight=1.25,
         ctrl_cost_weight=0.1,
@@ -33,6 +30,7 @@ class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         healthy_z_range=(1.0, 2.0),
         reset_noise_scale=1e-2,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -51,7 +49,7 @@ class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         )
 
         mujoco_env.MujocoEnv.__init__(
-            self, xml_file, 5, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, xml_file, 5, mujoco_bindings="mujoco_py", **kwargs
         )
 
     @property

gym/envs/mujoco/humanoid_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -206,7 +204,6 @@ class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         forward_reward_weight=1.25,
         ctrl_cost_weight=0.1,
         healthy_reward=5.0,
@@ -214,6 +211,7 @@ class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         healthy_z_range=(1.0, 2.0),
         reset_noise_scale=1e-2,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -229,7 +227,7 @@ class HumanoidEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(self, "humanoid.xml", 5, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, "humanoid.xml", 5, **kwargs)
 
     @property
     def healthy_reward(self):

gym/envs/mujoco/humanoidstandup.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,13 +5,9 @@ from gym.envs.mujoco import mujoco_env
 
 
 class HumanoidStandupEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self,
+            self, "humanoidstandup.xml", 5, mujoco_bindings="mujoco_py", **kwargs
-            "humanoidstandup.xml",
-            5,
-            render_mode=render_mode,
-            mujoco_bindings="mujoco_py",
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/humanoidstandup_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -192,10 +190,8 @@ class HumanoidStandupEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     """
 
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
-        mujoco_env.MujocoEnv.__init__(
+        mujoco_env.MujocoEnv.__init__(self, "humanoidstandup.xml", 5, **kwargs)
-            self, "humanoidstandup.xml", 5, render_mode=render_mode
-        )
         utils.EzPickle.__init__(self)
 
     def _get_obs(self):

gym/envs/mujoco/inverted_double_pendulum.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,13 +5,13 @@ from gym.envs.mujoco import mujoco_env
 
 
 class InvertedDoublePendulumEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
             self,
             "inverted_double_pendulum.xml",
             5,
-            render_mode=render_mode,
             mujoco_bindings="mujoco_py",
+            **kwargs
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/inverted_double_pendulum_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -120,10 +118,8 @@ class InvertedDoublePendulumEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     """
 
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
-        mujoco_env.MujocoEnv.__init__(
+        mujoco_env.MujocoEnv.__init__(self, "inverted_double_pendulum.xml", 5, **kwargs)
-            self, "inverted_double_pendulum.xml", 5, render_mode=render_mode
-        )
         utils.EzPickle.__init__(self)
 
     def step(self, action):

gym/envs/mujoco/inverted_pendulum.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,14 +5,10 @@ from gym.envs.mujoco import mujoco_env
 
 
 class InvertedPendulumEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         utils.EzPickle.__init__(self)
         mujoco_env.MujocoEnv.__init__(
-            self,
+            self, "inverted_pendulum.xml", 2, mujoco_bindings="mujoco_py", **kwargs
-            "inverted_pendulum.xml",
-            2,
-            render_mode=render_mode,
-            mujoco_bindings="mujoco_py",
         )
 
     def step(self, a):

gym/envs/mujoco/inverted_pendulum_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -92,11 +90,9 @@ class InvertedPendulumEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     """
 
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         utils.EzPickle.__init__(self)
-        mujoco_env.MujocoEnv.__init__(
+        mujoco_env.MujocoEnv.__init__(self, "inverted_pendulum.xml", 2, **kwargs)
-            self, "inverted_pendulum.xml", 2, render_mode=render_mode
-        )
 
     def step(self, a):
         reward = 1.0

gym/envs/mujoco/mujoco_env.py

@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from functools import partial
 from os import path
 from typing import Optional
 
@@ -39,6 +40,10 @@ class MujocoEnv(gym.Env):
         model_path,
         frame_skip,
         render_mode: Optional[str] = None,
+        width: int = DEFAULT_SIZE,
+        height: int = DEFAULT_SIZE,
+        camera_id: Optional[int] = None,
+        camera_name: Optional[str] = None,
         mujoco_bindings="mujoco",
     ):
         if model_path.startswith("/"):
@@ -106,7 +111,14 @@ class MujocoEnv(gym.Env):
         self._set_action_space()
 
         self.render_mode = render_mode
-        self.renderer = Renderer(self.render_mode, self._render)
+        render_frame = partial(
+            self._render,
+            width=width,
+            height=height,
+            camera_name=camera_name,
+            camera_id=camera_id,
+        )
+        self.renderer = Renderer(self.render_mode, render_frame)
 
         action = self.action_space.sample()
         observation, _reward, done, _info = self.step(action)
@@ -206,15 +218,23 @@ class MujocoEnv(gym.Env):
 
     def render(
         self,
-        mode="human",
+        mode: str = "human",
-        width=DEFAULT_SIZE,
+        width: Optional[int] = None,
-        height=DEFAULT_SIZE,
+        height: Optional[int] = None,
-        camera_id=None,
+        camera_id: Optional[int] = None,
-        camera_name=None,
+        camera_name: Optional[str] = None,
     ):
         if self.render_mode is not None:
+            assert (
+                width is None
+                and height is None
+                and camera_id is None
+                and camera_name is None
+            ), "Unexpected argument for render. Specify render arguments at environment initialization."
             return self.renderer.get_renders()
         else:
+            width = width if width is not None else DEFAULT_SIZE
+            height = height if height is not None else DEFAULT_SIZE
             return self._render(
                 mode=mode,
                 width=width,
@@ -225,11 +245,11 @@ class MujocoEnv(gym.Env):
 
     def _render(
         self,
-        mode="human",
+        mode: str = "human",
-        width=DEFAULT_SIZE,
+        width: int = DEFAULT_SIZE,
-        height=DEFAULT_SIZE,
+        height: int = DEFAULT_SIZE,
-        camera_id=None,
+        camera_id: Optional[int] = None,
-        camera_name=None,
+        camera_name: Optional[str] = None,
     ):
         assert mode in self.metadata["render_modes"]
 

gym/envs/mujoco/pusher.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,10 +5,10 @@ from gym.envs.mujoco import mujoco_env
 
 
 class PusherEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         utils.EzPickle.__init__(self)
         mujoco_env.MujocoEnv.__init__(
-            self, "pusher.xml", 5, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, "pusher.xml", 5, mujoco_bindings="mujoco_py", **kwargs
         )
 
     def step(self, a):

gym/envs/mujoco/pusher_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -134,9 +132,9 @@ class PusherEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     """
 
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         utils.EzPickle.__init__(self)
-        mujoco_env.MujocoEnv.__init__(self, "pusher.xml", 5, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, "pusher.xml", 5, **kwargs)
 
     def step(self, a):
         vec_1 = self.get_body_com("object") - self.get_body_com("tips_arm")

gym/envs/mujoco/reacher.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,10 +5,10 @@ from gym.envs.mujoco import mujoco_env
 
 
 class ReacherEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         utils.EzPickle.__init__(self)
         mujoco_env.MujocoEnv.__init__(
-            self, "reacher.xml", 2, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, "reacher.xml", 2, mujoco_bindings="mujoco_py", **kwargs
         )
 
     def step(self, a):

gym/envs/mujoco/reacher_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -124,9 +122,9 @@ class ReacherEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     """
 
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         utils.EzPickle.__init__(self)
-        mujoco_env.MujocoEnv.__init__(self, "reacher.xml", 2, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, "reacher.xml", 2, **kwargs)
 
     def step(self, a):
         vec = self.get_body_com("fingertip") - self.get_body_com("target")

gym/envs/mujoco/swimmer.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,9 +5,9 @@ from gym.envs.mujoco import mujoco_env
 
 
 class SwimmerEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self, "swimmer.xml", 4, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, "swimmer.xml", 4, mujoco_bindings="mujoco_py", **kwargs
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/swimmer_v3.py

@@ -1,7 +1,5 @@
 __credits__ = ["Rushiv Arora"]
 
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -13,12 +11,12 @@ DEFAULT_CAMERA_CONFIG = {}
 class SwimmerEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="swimmer.xml",
         forward_reward_weight=1.0,
         ctrl_cost_weight=1e-4,
         reset_noise_scale=0.1,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -32,7 +30,7 @@ class SwimmerEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         )
 
         mujoco_env.MujocoEnv.__init__(
-            self, xml_file, 4, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, xml_file, 4, mujoco_bindings="mujoco_py", **kwargs
         )
 
     def control_cost(self, action):

gym/envs/mujoco/swimmer_v4.py

@@ -1,7 +1,5 @@
 __credits__ = ["Rushiv Arora"]
 
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -134,11 +132,11 @@ class SwimmerEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         forward_reward_weight=1.0,
         ctrl_cost_weight=1e-4,
         reset_noise_scale=0.1,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -151,7 +149,7 @@ class SwimmerEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(self, "swimmer.xml", 4, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, "swimmer.xml", 4, **kwargs)
 
     def control_cost(self, action):
         control_cost = self._ctrl_cost_weight * np.sum(np.square(action))

gym/envs/mujoco/walker2d.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -7,13 +5,9 @@ from gym.envs.mujoco import mujoco_env
 
 
 class Walker2dEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, render_mode: Optional[str] = None):
+    def __init__(self, **kwargs):
         mujoco_env.MujocoEnv.__init__(
-            self,
+            self, "walker2d.xml", 4, mujoco_bindings="mujoco_py", **kwargs
-            "walker2d.xml",
-            4,
-            render_mode=render_mode,
-            mujoco_bindings="mujoco_py",
         )
         utils.EzPickle.__init__(self)
 

gym/envs/mujoco/walker2d_v3.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -16,7 +14,6 @@ DEFAULT_CAMERA_CONFIG = {
 class Walker2dEnv(mujoco_env.MujocoEnv, utils.EzPickle):
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         xml_file="walker2d.xml",
         forward_reward_weight=1.0,
         ctrl_cost_weight=1e-3,
@@ -46,7 +43,7 @@ class Walker2dEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         )
 
         mujoco_env.MujocoEnv.__init__(
-            self, xml_file, 4, render_mode=render_mode, mujoco_bindings="mujoco_py"
+            self, xml_file, 4, mujoco_bindings="mujoco_py", **kwargs
         )
 
     @property

gym/envs/mujoco/walker2d_v4.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import numpy as np
 
 from gym import utils
@@ -160,7 +158,6 @@ class Walker2dEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
     def __init__(
         self,
-        render_mode: Optional[str] = None,
         forward_reward_weight=1.0,
         ctrl_cost_weight=1e-3,
         healthy_reward=1.0,
@@ -169,6 +166,7 @@ class Walker2dEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         healthy_angle_range=(-1.0, 1.0),
         reset_noise_scale=5e-3,
         exclude_current_positions_from_observation=True,
+        **kwargs
     ):
         utils.EzPickle.__init__(**locals())
 
@@ -187,7 +185,7 @@ class Walker2dEnv(mujoco_env.MujocoEnv, utils.EzPickle):
             exclude_current_positions_from_observation
         )
 
-        mujoco_env.MujocoEnv.__init__(self, "walker2d.xml", 4, render_mode=render_mode)
+        mujoco_env.MujocoEnv.__init__(self, "walker2d.xml", 4, **kwargs)
 
     @property
     def healthy_reward(self):