From 1c10a124dc895dc556875c4bb39bb4901e6d14ca Mon Sep 17 00:00:00 2001
From: Kallinteris Andreas
 <30759571+Kallinteris-Andreas@users.noreply.github.com>
Date: Wed, 24 Apr 2024 01:08:25 +0300
Subject: [PATCH] Update mujoco_env.py

---
 gymnasium/envs/mujoco/mujoco_env.py | 80 ++++++++++-------------------
 1 file changed, 27 insertions(+), 53 deletions(-)

diff --git a/gymnasium/envs/mujoco/mujoco_env.py b/gymnasium/envs/mujoco/mujoco_env.py
index 7e76db23a..70a7aaf15 100644
--- a/gymnasium/envs/mujoco/mujoco_env.py
+++ b/gymnasium/envs/mujoco/mujoco_env.py
@@ -1,4 +1,4 @@
-Jfrom os import path
+from os import path
 from typing import Any, Dict, Optional, Tuple, Union
 
 import numpy as np
@@ -34,20 +34,42 @@ def expand_model_path(model_path: str) -> str:
     return fullpath
 
 
-class BaseMujocoEnv(gym.Env[NDArray[np.float64], NDArray[np.float32]]):
-    """Superclass for all MuJoCo environments."""
+class MujocoEnv(BaseMujocoEnv):
+    """Superclass for MuJoCo based environments."""
 
     def __init__(
         self,
-        model_path,
-        frame_skip,
+        model_path: str,
+        frame_skip: int,
         observation_space: Optional[Space],
         render_mode: Optional[str] = None,
         width: int = DEFAULT_SIZE,
         height: int = DEFAULT_SIZE,
         camera_id: Optional[int] = None,
         camera_name: Optional[str] = None,
+        default_camera_config: Optional[Dict[str, Union[float, int]]] = None,
+        max_geom: int = 1000,
+        visual_options: Dict[int, bool] = {},
     ):
+        """Base abstract class for mujoco based environments.
+
+        Args:
+            model_path: Path to the MuJoCo Model.
+            frame_skip: Number of MuJoCo simulation steps per gym `step()`.
+            observation_space: The observation space of the environment.
+            render_mode: The `render_mode` used.
+            width: The width of the render window.
+            height: The height of the render window.
+            camera_id: The camera ID used.
+            camera_name: The name of the camera used (can not be used in conjunction with `camera_id`).
+            default_camera_config: configuration for rendering camera.
+            max_geom: max number of rendered geometries.
+            visual_options: render flag options.
+
+        Raises:
+            OSError: when the `model_path` does not exist.
+            error.DependencyNotInstalled: When `mujoco` is not installed.
+        """
         self.fullpath = expand_model_path(model_path)
 
         self.width = width
@@ -77,54 +99,6 @@ def __init__(
         self.camera_name = camera_name
         self.camera_id = camera_id
 
-
-class MujocoEnv(BaseMujocoEnv):
-    """Superclass for MuJoCo based environments."""
-
-    def __init__(
-        self,
-        model_path,
-        frame_skip,
-        observation_space: Optional[Space],
-        render_mode: Optional[str] = None,
-        width: int = DEFAULT_SIZE,
-        height: int = DEFAULT_SIZE,
-        camera_id: Optional[int] = None,
-        camera_name: Optional[str] = None,
-        default_camera_config: Optional[Dict[str, Union[float, int]]] = None,
-        max_geom: int = 1000,
-        visual_options: Dict[int, bool] = {},
-    ):
-        """Base abstract class for mujoco based environments.
-
-        Args:
-            model_path: Path to the MuJoCo Model.
-            frame_skip: Number of MuJoCo simulation steps per gym `step()`.
-            observation_space: The observation space of the environment.
-            render_mode: The `render_mode` used.
-            width: The width of the render window.
-            height: The height of the render window.
-            camera_id: The camera ID used.
-            camera_name: The name of the camera used (can not be used in conjunction with `camera_id`).
-            default_camera_config: configuration for rendering camera.
-            max_geom: max number of rendered geometries.
-            visual_options: render flag options.
-
-        Raises:
-            OSError: when the `model_path` does not exist.
-            error.DependencyNotInstalled: When `mujoco` is not installed.
-        """
-        super().__init__(
-            model_path,
-            frame_skip,
-            observation_space,
-            render_mode,
-            width,
-            height,
-            camera_id,
-            camera_name,
-        )
-
         from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
 
         self.mujoco_renderer = MujocoRenderer(