From a0009ed9a208ff9864a5c1368c04c273bb20dd06 Mon Sep 17 00:00:00 2001
From: Avnish Narayan <38871737+avnishn@users.noreply.github.com>
Date: Thu, 22 Jul 2021 17:42:05 -0700
Subject: [PATCH] Add easier rendering options, modify near plane, add better
 camera angles (#348)

---
 .../assets_v2/objects/assets/xyz_base.xml     |  3 ++-
 .../envs/assets_v2/scene/basic_scene.xml      |  2 +-
 .../envs/assets_v2/scene/basic_scene_b.xml    |  2 +-
 .../envs/assets_v2/scene/basic_scene_c.xml    |  2 +-
 metaworld/envs/mujoco/mujoco_env.py           | 25 ++++++++++---------
 5 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/metaworld/envs/assets_v2/objects/assets/xyz_base.xml b/metaworld/envs/assets_v2/objects/assets/xyz_base.xml
index 66d13dd4a..3511c6d31 100644
--- a/metaworld/envs/assets_v2/objects/assets/xyz_base.xml
+++ b/metaworld/envs/assets_v2/objects/assets/xyz_base.xml
@@ -15,7 +15,8 @@
 
       <camera pos="0 0.5 1.5" name="topview" />
       <camera name="corner" mode="fixed" pos="-1.1 -0.4 0.6" xyaxes="-1 1 0 -0.2 -0.2 -1"/>
-      <camera name="corner2" mode="fixed" pos="1.1 -0.4 0.6" euler="3.9 2.3 3.7"/>
+      <camera name="corner2" fovy="60" mode="fixed" pos="1.3 -0.2 1.1" euler="3.9 2.3 0.6"/>
+      <camera name="corner3" fovy="45" mode="fixed" pos="0.9 0 1.5" euler="3.5 2.7 1"/>
       <!--<geom name="floor" type="plane" pos="0 0 -.9" size="10 10 10"-->
             <!--rgba="0 0 0 1" contype="15" conaffinity="15" />-->
       <!--<geom name="tableTop" type="box" pos="0 0.6 -0.45" size="0.4 0.2 0.45"
diff --git a/metaworld/envs/assets_v2/scene/basic_scene.xml b/metaworld/envs/assets_v2/scene/basic_scene.xml
index eade314ae..20f9281df 100644
--- a/metaworld/envs/assets_v2/scene/basic_scene.xml
+++ b/metaworld/envs/assets_v2/scene/basic_scene.xml
@@ -24,7 +24,7 @@
     </asset>
 
     <visual>
-        <map fogstart="1.5" fogend="5" force="0.1" znear="0.1"/>
+        <map fogstart="1.5" fogend="5" force="0.1" znear="0.01"/>
         <quality shadowsize="4096" offsamples="4"/>
 
         <headlight ambient="0.4 0.4 0.4"/>
diff --git a/metaworld/envs/assets_v2/scene/basic_scene_b.xml b/metaworld/envs/assets_v2/scene/basic_scene_b.xml
index e11384008..7964a97fa 100644
--- a/metaworld/envs/assets_v2/scene/basic_scene_b.xml
+++ b/metaworld/envs/assets_v2/scene/basic_scene_b.xml
@@ -18,7 +18,7 @@
   </asset>
 
   <visual>
-      <map fogstart="1.5" fogend="5" force="0.1" znear="0.1"/>
+      <map fogstart="1.5" fogend="5" force="0.1" znear="0.01"/>
       <quality shadowsize ="4096" offsamples="4"/>
 
       <headlight ambient="0.4 0.4 0.4"/>
diff --git a/metaworld/envs/assets_v2/scene/basic_scene_c.xml b/metaworld/envs/assets_v2/scene/basic_scene_c.xml
index 1332b2a37..53ad304cc 100644
--- a/metaworld/envs/assets_v2/scene/basic_scene_c.xml
+++ b/metaworld/envs/assets_v2/scene/basic_scene_c.xml
@@ -18,7 +18,7 @@
   </asset>
 
   <visual>
-      <map fogstart="1.5" fogend="5" force="0.1" znear="0.1"/>
+      <map fogstart="1.5" fogend="5" force="0.1" znear="0.01"/>
       <quality shadowsize ="4096" offsamples="4"/>
 
       <headlight ambient="0.4 0.4 0.4"/>
diff --git a/metaworld/envs/mujoco/mujoco_env.py b/metaworld/envs/mujoco/mujoco_env.py
index 0c41e6591..47dd2a58d 100644
--- a/metaworld/envs/mujoco/mujoco_env.py
+++ b/metaworld/envs/mujoco/mujoco_env.py
@@ -36,9 +36,9 @@ class MujocoEnv(gym.Env, abc.ABC):
      - Do not automatically set the observation/action space.
     """
 
-    max_path_length = 150
+    max_path_length = 500
 
-    def __init__(self, model_path, frame_skip, rgb_array_res=(640, 480)):
+    def __init__(self, model_path, frame_skip):
         if not path.exists(model_path):
             raise IOError("File %s does not exist" % model_path)
 
@@ -48,7 +48,6 @@ def __init__(self, model_path, frame_skip, rgb_array_res=(640, 480)):
         self.data = self.sim.data
         self.viewer = None
         self._viewers = {}
-        self._rgb_array_res = rgb_array_res
 
         self.metadata = {
             'render.modes': ['human'],
@@ -119,17 +118,19 @@ def do_simulation(self, ctrl, n_frames=None):
                 warnings.warn(str(err), category=RuntimeWarning)
                 self._did_see_sim_exception = True
 
-    def render(self, mode='human'):
-        if mode == 'human':
-            self._get_viewer(mode).render()
-        elif mode == 'rgb_array':
+    def render(self, offscreen=False, camera_name="corner2", resolution=(640, 480)):
+        assert_string = ("camera_name should be one of ",
+                "corner3, corner, corner2, topview, gripperPOV, behindGripper")
+        assert camera_name in {"corner3", "corner", "corner2", 
+            "topview", "gripperPOV", "behindGripper"}, assert_string
+        if not offscreen:
+            self._get_viewer('human').render()
+        else:
             return self.sim.render(
-                *self._rgb_array_res,
+                *resolution,
                 mode='offscreen',
-                camera_name='topview'
-            )[:, :, ::-1]
-        else:
-            raise ValueError("mode can only be either 'human' or 'rgb_array'")
+                camera_name=camera_name
+            )
 
     def close(self):
         if self.viewer is not None: