From 4087203b0688414844e6999567a4d36f9b7accbe Mon Sep 17 00:00:00 2001
From: kngwyu <yuji.kngw.80s.revive@gmail.com>
Date: Tue, 20 Apr 2021 10:09:14 +0900
Subject: [PATCH] Formatting

---
 mujoco_maze/agent_model.py    | 9 +++------
 mujoco_maze/ant.py            | 4 +++-
 mujoco_maze/maze_env.py       | 9 ++++-----
 mujoco_maze/maze_env_utils.py | 7 ++++---
 mujoco_maze/swimmer.py        | 8 ++++++--
 mujoco_maze/websock_viewer.py | 1 -
 6 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/mujoco_maze/agent_model.py b/mujoco_maze/agent_model.py
index d209d78..9d50fe4 100644
--- a/mujoco_maze/agent_model.py
+++ b/mujoco_maze/agent_model.py
@@ -27,16 +27,13 @@ class AgentModel(ABC, MujocoEnv, EzPickle):
 
     @abstractmethod
     def _get_obs(self) -> np.ndarray:
-        """Returns the observation from the model.
-        """
+        """Returns the observation from the model."""
         pass
 
     def get_xy(self) -> np.ndarray:
-        """Returns the coordinate of the agent.
-        """
+        """Returns the coordinate of the agent."""
         pass
 
     def set_xy(self, xy: np.ndarray) -> None:
-        """Set the coordinate of the agent.
-        """
+        """Set the coordinate of the agent."""
         pass
diff --git a/mujoco_maze/ant.py b/mujoco_maze/ant.py
index 34356e1..01a9e20 100644
--- a/mujoco_maze/ant.py
+++ b/mujoco_maze/ant.py
@@ -82,7 +82,9 @@ class AntEnv(AgentModel):
 
     def reset_model(self):
         qpos = self.init_qpos + self.np_random.uniform(
-            size=self.model.nq, low=-0.1, high=0.1,
+            size=self.model.nq,
+            low=-0.1,
+            high=0.1,
         )
         qvel = self.init_qvel + self.np_random.randn(self.model.nv) * 0.1
 
diff --git a/mujoco_maze/maze_env.py b/mujoco_maze/maze_env.py
index 3a455d0..6b89d87 100644
--- a/mujoco_maze/maze_env.py
+++ b/mujoco_maze/maze_env.py
@@ -370,17 +370,16 @@ class MazeEnv(gym.Env):
             dtype=np.uint8,
         )
 
-    def render(self, mode="human", **kwargs) -> Any:
-        if self._websock_port is not None:
+    def render(self, mode="human", **kwargs) -> Optional[np.ndarray]:
+        if mode == "human" and self._websock_port is not None:
             if self._mj_offscreen_viewer is None:
-                from mujoco_py import MjRenderContextOffscreen as MjRenderOffscreen
+                from mujoco_py import MjRenderContextOffscreen as MjRCO
 
                 from mujoco_maze.websock_viewer import start_server
 
-                self._mj_offscreen_viewer = MjRenderOffscreen(self.wrapped_env.sim)
+                self._mj_offscreen_viewer = MjRCO(self.wrapped_env.sim)
                 self._websock_server_pipe = start_server(self._websock_port)
             self._websock_server_pipe.send(self._render_image())
-            return True
         else:
             return self.wrapped_env.render(mode, **kwargs)
 
diff --git a/mujoco_maze/maze_env_utils.py b/mujoco_maze/maze_env_utils.py
index 348b88c..f9edfe1 100644
--- a/mujoco_maze/maze_env_utils.py
+++ b/mujoco_maze/maze_env_utils.py
@@ -83,7 +83,9 @@ class MazeCell(Enum):
 
 class Line:
     def __init__(
-        self, p1: Union[Sequence[float], Point], p2: Union[Sequence[float], Point],
+        self,
+        p1: Union[Sequence[float], Point],
+        p2: Union[Sequence[float], Point],
     ) -> None:
         self.p1 = p1 if isinstance(p1, Point) else np.complex(*p1)
         self.p2 = p2 if isinstance(p2, Point) else np.complex(*p2)
@@ -141,8 +143,7 @@ class Collision:
 
 
 class CollisionDetector:
-    """For manual collision detection.
-    """
+    """For manual collision detection."""
 
     EPS: float = 0.05
     NEIGHBORS: List[Tuple[int, int]] = [[0, -1], [-1, 0], [0, 1], [1, 0]]
diff --git a/mujoco_maze/swimmer.py b/mujoco_maze/swimmer.py
index ef56e89..7d2c6de 100644
--- a/mujoco_maze/swimmer.py
+++ b/mujoco_maze/swimmer.py
@@ -55,10 +55,14 @@ class SwimmerEnv(AgentModel):
 
     def reset_model(self) -> np.ndarray:
         qpos = self.init_qpos + self.np_random.uniform(
-            low=-0.1, high=0.1, size=self.model.nq,
+            low=-0.1,
+            high=0.1,
+            size=self.model.nq,
         )
         qvel = self.init_qvel + self.np_random.uniform(
-            low=-0.1, high=0.1, size=self.model.nv,
+            low=-0.1,
+            high=0.1,
+            size=self.model.nv,
         )
 
         self.set_state(qpos, qvel)
diff --git a/mujoco_maze/websock_viewer.py b/mujoco_maze/websock_viewer.py
index 07b82d0..a355b3f 100644
--- a/mujoco_maze/websock_viewer.py
+++ b/mujoco_maze/websock_viewer.py
@@ -4,7 +4,6 @@ import multiprocessing as mp
 
 import fastapi
 import uvicorn
-
 from PIL import Image
 
 HTML = """