Repository: PKU-Alignment/omnisafe (language: python)

Issue:
How to change the render mode?
## Questions
It seems that the environment cannot modify its rendering mode. We tested two ways and both failed.

Like the new way in gymnasium library:
`env = safety_gymnasium.make(env_name, render='rgb_array')` 
which gets 
`TypeError: __init__() got an unexpected keyword argument 'render'`

Or the old way in gym library:
`env.render(render_mode='rgb_array')`
which gets
`TypeError: env_render_passive_checker() got an unexpected keyword argument 'render_mode'`.


## Checklist

- [x] I have checked that there is no similar issue in the repo. (**required**)
- [x] I have read the [documentation](https://omnisafe.readthedocs.io). (**required**)



Discussion / hints:

Fix patch:
diff --git a/examples/vis_safety_gymnasium.py b/examples/vis_safety_gymnasium.py
index 88c39eec2..f368490b2 100644
--- a/examples/vis_safety_gymnasium.py
+++ b/examples/vis_safety_gymnasium.py
@@ -1,13 +1,36 @@
+import argparse
+
 import safety_gymnasium
 
 
-env_name = 'SafetyPointGoal1-v0'
-env = safety_gymnasium.make(env_name)
+def run_random(env_name):
+    env = safety_gymnasium.make(env_name, render_mode='human')
+    # env.seed(0)
+    obs, _ = env.reset()
+    terminled = False
+    ep_ret = 0
+    ep_cost = 0
+    while True:
+        if terminled:
+            print('Episode Return: %.3f \t Episode Cost: %.3f' % (ep_ret, ep_cost))
+            ep_ret, ep_cost = 0, 0
+            obs, _ = env.reset()
+        assert env.observation_space.contains(obs)
+        act = env.action_space.sample()
+        assert env.action_space.contains(act)
+        # Use the environment's built_in max_episode_steps
+        if hasattr(env, '_max_episode_steps'):
+            max_ep_len = env._max_episode_steps
+
+        obs, reward, cost, terminled, truncated, info = env.step(act)
+
+        ep_ret += reward
+        ep_cost += cost
+
 
-obs, info = env.reset()
-terminated = False
+if __name__ == '__main__':
 
-while not terminated:
-    act = env.action_space.sample()
-    obs, reward, cost, terminated, truncated, info = env.step(act)
-    env.render()
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--env', default='SafetyPointGoal2-v0')
+    args = parser.parse_args()
+    run_random(args.env)
diff --git a/omnisafe/envs/Safety_Gymnasium/examples/safety_gym_v2_vision.py b/omnisafe/envs/Safety_Gymnasium/examples/safety_gym_v2_vision.py
new file mode 100644
index 000000000..10b1fb467
--- /dev/null
+++ b/omnisafe/envs/Safety_Gymnasium/examples/safety_gym_v2_vision.py
@@ -0,0 +1,51 @@
+import argparse
+import os
+
+# import gymnasium
+import safety_gymnasium
+from gymnasium.utils.save_video import save_video
+
+
+WORKDIR = os.path.abspath('.')
+DIR = os.path.join(WORKDIR, 'omnisafe/envs/Safety_Gymnasium/examples', 'cached_test_vision_video')
+
+
+def run_random(env_name):
+    env = safety_gymnasium.make(env_name)
+    # env.seed(0)
+    obs, _ = env.reset()
+    terminled = False
+    ep_ret = 0
+    ep_cost = 0
+    render_list = []
+    for i in range(1001):
+        if terminled:
+            print('Episode Return: %.3f \t Episode Cost: %.3f' % (ep_ret, ep_cost))
+            ep_ret, ep_cost = 0, 0
+            obs, _ = env.reset()
+            save_video(
+                frames=render_list,
+                video_folder=DIR,
+                name_prefix=f'test_vision_output',
+                fps=30,
+            )
+            render_list = []
+        assert env.observation_space.contains(obs)
+        act = env.action_space.sample()
+        assert env.action_space.contains(act)
+        # Use the environment's built_in max_episode_steps
+        if hasattr(env, '_max_episode_steps'):
+            max_ep_len = env._max_episode_steps
+        render_list.append(obs['vision'])
+        obs, reward, cost, terminled, truncated, info = env.step(act)
+
+        ep_ret += reward
+        ep_cost += cost
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--env', default='SafetyCarGoal0Vision-v0')
+    args = parser.parse_args()
+    run_random(args.env)
diff --git a/omnisafe/envs/Safety_Gymnasium/examples/test_safety_gym_v2.py b/omnisafe/envs/Safety_Gymnasium/examples/vis_safety_gym_v2.py
similarity index 85%
rename from omnisafe/envs/Safety_Gymnasium/examples/test_safety_gym_v2.py
rename to omnisafe/envs/Safety_Gymnasium/examples/vis_safety_gym_v2.py
index 3d7925da3..b62a45b07 100644
--- a/omnisafe/envs/Safety_Gymnasium/examples/test_safety_gym_v2.py
+++ b/omnisafe/envs/Safety_Gymnasium/examples/vis_safety_gym_v2.py
@@ -1,11 +1,10 @@
 import argparse
 
-# import gymnasium
 import safety_gymnasium
 
 
 def run_random(env_name):
-    env = safety_gymnasium.make(env_name)
+    env = safety_gymnasium.make(env_name, render_mode='human')
     # env.seed(0)
     obs, _ = env.reset()
     terminled = False
@@ -27,12 +26,11 @@ def run_random(env_name):
 
         ep_ret += reward
         ep_cost += cost
-        env.render()
 
 
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--env', default='SafetyPointGoal2-v0')
+    parser.add_argument('--env', default='SafetyPointGoal0-v0')
     args = parser.parse_args()
     run_random(args.env)
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/__init__.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/__init__.py
index c2318369f..e67ef98eb 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/__init__.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/__init__.py
@@ -29,7 +29,7 @@
     },
 }
 
-MAKE_VISION_ENVIRONMENTS = False
+MAKE_VISION_ENVIRONMENTS = True
 
 # ========================================#
 # Helper Class for Easy Gym Registration #
@@ -81,7 +81,7 @@ def register(self, name='', config={}):
 
             if MAKE_VISION_ENVIRONMENTS:
                 # Vision: note, these environments are experimental! Correct behavior not guaranteed
-                vision_env_name = f'{self.prefix}-{robot_name}{self.name + name}Vision-{VERSION}'
+                vision_env_name = f'{self.prefix}{robot_name}{self.name + name}Vision-{VERSION}'
                 vision_config = {
                     'world': {},
                     'task': {
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/base_task.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/base_task.py
index acc8d9aaa..29c31aefd 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/base_task.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/base_task.py
@@ -44,14 +44,17 @@ def __init__(
         # otherwise, raise a python exception.
         # TODO: randomize starting joint positions
 
+        self.observe_vision = False
+        self.observation_flatten = True
+
         # # Vision observation parameters
         # TODO JJM
         self.vision_size = (
             60,
             40,
         )  # Size (width, height) of vision observation; gets flipped internally to (rows, cols) format
-        # self.vision_render = True  # Render vision observation in the viewer
-        # self.vision_render_size = (300, 200)  # Size to render the vision in the viewer
+        self.vision_render = True  # Render vision observation in the viewer
+        self.vision_render_size = (300, 200)  # Size to render the vision in the viewer
 
         # Lidar observation parameters
         self.lidar_num_bins = 16  # Bins (around a full circle) for lidar sensing
@@ -354,8 +357,8 @@ def obs_vision(self):
         # Get a render context so we can
         rows, cols = self.vision_size
         width, height = cols, rows
-        vision = self.sim.render(width, height, camera_name='vision', mode='offscreen')
-        return np.array(vision, dtype='float32') / 255
+        vision = self.engine.render(width, height, mode='rgb_array', camera_name='vision', cost={})
+        return vision
 
     def obs_lidar(self, positions, group):
         """
@@ -477,16 +480,16 @@ def build_sensor_observation_space(self):
         # if self.observe_sensors:
         for sensor in self.sensors_obs:  # Explicitly listed sensors
             dim = self.robot.sensor_dim[sensor]
-            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (dim,), dtype=np.float32)
+            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (dim,), dtype=np.float64)
         # Velocities don't have wraparound effects that rotational positions do
         # Wraparounds are not kind to neural networks
         # Whereas the angle 2*pi is very close to 0, this isn't true in the network
         # In theory the network could learn this, but in practice we simplify it
         # when the sensors_angle_components switch is enabled.
         for sensor in self.robot.hinge_vel_names:
-            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32)
+            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float64)
         for sensor in self.robot.ballangvel_names:
-            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (3,), dtype=np.float32)
+            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (3,), dtype=np.float64)
         # Angular positions have wraparound effects, so output something more friendly
         if self.sensors_angle_components:
             # Single joints are turned into sin(x), cos(x) pairs
@@ -494,7 +497,7 @@ def build_sensor_observation_space(self):
             # Since for angles, small perturbations in angle give small differences in sin/cos
             for sensor in self.robot.hinge_pos_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (2,), dtype=np.float32
+                    -np.inf, np.inf, (2,), dtype=np.float64
                 )
             # Quaternions are turned into 3x3 rotation matrices
             # Quaternions have a wraparound issue in how they are normalized,
@@ -508,18 +511,18 @@ def build_sensor_observation_space(self):
             # Instead we use a 3x3 rotation matrix, which if normalized, smoothly varies as well.
             for sensor in self.robot.ballquat_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (3, 3), dtype=np.float32
+                    -np.inf, np.inf, (3, 3), dtype=np.float64
                 )
         else:
             # Otherwise include the sensor without any processing
             # TODO: comparative study of the performance with and without this feature.
             for sensor in self.robot.hinge_pos_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (1,), dtype=np.float32
+                    -np.inf, np.inf, (1,), dtype=np.float64
                 )
             for sensor in self.robot.ballquat_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (4,), dtype=np.float32
+                    -np.inf, np.inf, (4,), dtype=np.float64
                 )
 
         return obs_space_dict
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/builder.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/builder.py
index c1efd9518..fd77dfea7 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/builder.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/builder.py
@@ -37,13 +37,6 @@ class Builder(gymnasium.Env, gymnasium.utils.EzPickle):
         'render_lidar_size': 0.025,
         'render_lidar_offset_init': 0.5,
         'render_lidar_offset_delta': 0.06,
-        # Vision observation parameters
-        'vision_size': (
-            60,
-            40,
-        ),  # Size (width, height) of vision observation; gets flipped internally to (rows, cols) format
-        'vision_render': True,  # Render vision observation in the viewer
-        'vision_render_size': (300, 200),  # Size to render the vision in the viewer
         # Frameskip is the number of physics simulation steps per environment step
         # Frameskip is sampled as a binomial distribution
         # For deterministic steps, set frameskip_binom_p = 1.0 (always take max frameskip)
@@ -51,7 +44,7 @@ class Builder(gymnasium.Env, gymnasium.utils.EzPickle):
         'frameskip_binom_p': 1.0,  # Probability of trial return (controls distribution)
     }
 
-    def __init__(self, config={}):
+    def __init__(self, config={}, **kwargs):
         # First, parse configuration. Important note: LOTS of stuff happens in
         # parse, and many attributes of the class get set through setattr. If you
         # are trying to track down where an attribute gets initially set, and
@@ -62,12 +55,13 @@ def __init__(self, config={}):
         self.get_config(config)
         self.task_id = config['task']['task_id']
         self.seed()
-        print('xxx')
 
         self._setup_simulation()
 
         self.done = True
 
+        self.render_mode = kwargs.get('render_mode', None)
+
     def get_config(self, config):
         """Parse a config dict - see self.DEFAULT for description"""
         world_config = config['world']
@@ -134,7 +128,6 @@ def reset(self, seed=None, options=None):
         self.first_reset = False  # Built our first world successfully
 
         # Return an observation
-
         return (self.task.obs(), info)
 
     def world_xy(self, pos):
@@ -227,11 +220,17 @@ def step(self, action):
         terminaled = self.done
         truncated = False
 
+        if self.render_mode == 'human':
+            self.render()
         return self.task.obs(), reward, cost, terminaled, truncated, info
 
-    def render(self, mode='human', camera_id=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT):
+    def render(self, camera_id=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT):
+        assert self.render_mode, 'Please specify the render mode when you make env.'
+        assert (
+            not self.task.observe_vision
+        ), 'When you use vision envs, you should not call this function explicitly.'
         return self.engine.render(
-            mode=mode, camera_id=camera_id, width=width, height=height, cost=self._cost
+            mode=self.render_mode, camera_id=camera_id, width=width, height=height, cost=self._cost
         )
 
     @property
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/engine.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/engine.py
index 419b6dda3..f0487bbc3 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/engine.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/engine.py
@@ -27,7 +27,7 @@ def __init__(self, task, world_config={}, task_config={}):
         self.task = task
 
         self.robot = Robot(self.robot_base)
-        self.action_space = gymnasium.spaces.Box(-1, 1, (self.robot.nu,), dtype=np.float32)
+        self.action_space = gymnasium.spaces.Box(-1, 1, (self.robot.nu,), dtype=np.float64)
         self.placements = self.task.build_placements_dict()
         # self.world = self.get_world()
         self.clear()
@@ -295,29 +295,36 @@ def render_swap_callback(self):
         #     self.viewer.draw_pixels(self.save_obs_vision, 0, 0)
         pass
 
-    def render(self, width, height, mode='human', camera_id=None, camera_name=None, cost={}):
+    def render(self, width, height, mode, camera_id=None, camera_name=None, cost={}):
         """Render the environment to the screen"""
         self.model.vis.global_.offwidth = width
         self.model.vis.global_.offheight = height
 
-        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.')
+        if mode in {
+            'rgb_array',
+            'depth_array',
+        }:
 
-        if mode == 'rgb_array':
-            self._get_viewer(mode)
-            self.viewer._hide_overlay = True
-            camera_id = 3
-            self.viewer.cam.fixedcamid = camera_id  # self.model.camera_name2id(mode)
-            self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
-        elif mode == 'depth_array':
-            self._get_viewer(mode)
-            self.viewer._hide_overlay = True
-            camera_id = 3
-            self.viewer.cam.fixedcamid = camera_id  # self.model.camera_name2id(mode)
-            self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
-        elif mode == 'human':
+            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_id = 3
+
+            if camera_id is None:
+                camera_id = mujoco.mj_name2id(
+                    self.model,
+                    mujoco.mjtObj.mjOBJ_CAMERA,
+                    camera_name,
+                )
+
+                self._get_viewer(mode).render(camera_id=camera_id)
+
+        if mode == 'human':
             self._get_viewer(mode)
-            self.viewer.cam.fixedcamid = -1
             self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
 
         self.viewer.render_swap_callback = self.render_swap_callback
@@ -325,17 +332,6 @@ def render(self, width, height, mode='human', camera_id=None, camera_name=None,
         self.viewer.vopt.geomgroup[:] = 1
         # self.viewer.update_sim(self.sim)
 
-        if camera_id is not None:
-            # Update camera if desired
-            self.viewer.cam.fixedcamid = camera_id
-        elif camera_name is not None:
-            camera_id = mujoco.mj_name2id(
-                self.model,
-                mujoco.mjtObj.mjOBJ_CAMERA,
-                camera_name,
-            )
-            self.viewer.cam.fixedcamid = camera_id
-
         # Lidar markers
         if self.render_lidar_markers:
             offset = self.render_lidar_offset_init  # Height offset for successive lidar indicators
@@ -428,8 +424,8 @@ def render(self, width, height, mode='human', camera_id=None, camera_name=None,
     def _get_viewer(
         self, mode
     ) -> Union[
-        'gym.envs.mujoco.mujoco_rendering.Viewer',
-        'gym.envs.mujoco.mujoco_rendering.RenderContextOffscreen',
+        'gymnasium.envs.mujoco.mujoco_rendering.Viewer',
+        'gymnasium.envs.mujoco.mujoco_rendering.RenderContextOffscreen',
     ]:
         self.viewer = self._viewers.get(mode)
         if self.viewer is None:
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level0.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level0.py
index 9bde5fe43..3d5e199e7 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level0.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level0.py
@@ -145,7 +145,7 @@ def build_world_config(self, layout):
 
         # if not self.observe_vision:
         #    world_config['render_context'] = -1  # Hijack this so we don't create context
-        # world_config['observe_vision'] = self.observe_vision
+        world_config['observe_vision'] = self.observe_vision
 
         # Extra geoms (immovable objects) to add to the scene
         world_config['geoms'] = {}
@@ -166,28 +166,30 @@ def build_observation_space(self):
 
         # if self.observe_goal_lidar:
         obs_space_dict['goal_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
         # if self.buttons_num and self.observe_buttons:
         obs_space_dict['buttons_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
-        # if self.observe_vision:
-        #     width, height = self.vision_size
-        #     rows, cols = height, width
-        #     self.vision_size = (rows, cols)
-        #     obs_space_dict['vision'] = gym.spaces.Box(0, 1.0, self.vision_size + (3,), dtype=np.float32)
+        if self.observe_vision:
+            width, height = self.vision_size
+            rows, cols = height, width
+            self.vision_size = (rows, cols)
+            obs_space_dict['vision'] = gymnasium.spaces.Box(
+                0, 255, self.vision_size + (3,), dtype=np.uint8
+            )
 
         # Flatten it ourselves
         self.obs_space_dict = obs_space_dict
-        # if self.observation_flatten:
-        self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
-        self.observation_space = gymnasium.spaces.Box(
-            -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
-        )
-        # else:
-        #     self.observation_space = gym.spaces.Dict(obs_space_dict)
+        if self.observation_flatten:
+            self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
+            self.observation_space = gymnasium.spaces.Box(
+                -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
+            )
+        else:
+            self.observation_space = gymnasium.spaces.Dict(obs_space_dict)
 
     def obs(self):
         """Return the observation of our agent"""
@@ -206,18 +208,17 @@ def obs(self):
         else:
             obs['buttons_lidar'] = np.zeros(self.lidar_num_bins)
 
-        # if self.observe_vision:
-        #     obs['vision'] = self.obs_vision()
-
-        # if self.observation_flatten:
-        flat_obs = np.zeros(self.obs_flat_size)
-        offset = 0
-        for k in sorted(self.obs_space_dict.keys()):
-            k_size = np.prod(obs[k].shape)
-            flat_obs[offset : offset + k_size] = obs[k].flat
-            offset += k_size
-        obs = flat_obs
-        assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
+        if self.observe_vision:
+            obs['vision'] = self.obs_vision()
+        if self.observation_flatten:
+            flat_obs = np.zeros(self.obs_flat_size)
+            offset = 0
+            for k in sorted(self.obs_space_dict.keys()):
+                k_size = np.prod(obs[k].shape)
+                flat_obs[offset : offset + k_size] = obs[k].flat
+                offset += k_size
+            obs = flat_obs
+            assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
         return obs
 
     def buttons_timer_tick(self):
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level1.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level1.py
index e41f9ac7e..0dc02500d 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level1.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/button/button_level1.py
@@ -92,7 +92,7 @@ def build_world_config(self, layout):
 
         # if not self.observe_vision:
         #    world_config['render_context'] = -1  # Hijack this so we don't create context
-        # world_config['observe_vision'] = self.observe_vision
+        world_config['observe_vision'] = self.observe_vision
 
         # Extra objects to add to the scene
         world_config['objects'] = {}
@@ -139,38 +139,41 @@ def build_observation_space(self):
 
         # if self.observe_goal_lidar:
         obs_space_dict['goal_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.observe_hazards:
         obs_space_dict['hazards_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.gremlins_num and self.observe_gremlins:
         obs_space_dict['gremlins_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.buttons_num and self.observe_buttons:
         obs_space_dict['buttons_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
-        # if self.observe_vision:
-        #     width, height = self.vision_size
-        #     rows, cols = height, width
-        #     self.vision_size = (rows, cols)
-        #     obs_space_dict['vision'] = gym.spaces.Box(0, 1.0, self.vision_size + (3,), dtype=np.float32)
+        if self.observe_vision:
+            width, height = self.vision_size
+            rows, cols = height, width
+            self.vision_size = (rows, cols)
+            obs_space_dict['vision'] = gymnasium.spaces.Box(
+                0, 255, self.vision_size + (3,), dtype=np.uint8
+            )
+
         # Flatten it ourselves
         self.obs_space_dict = obs_space_dict
-        # if self.observation_flatten:
-        self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
-        self.observation_space = gymnasium.spaces.Box(
-            -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
-        )
-        # else:
-        #     self.observation_space = gym.spaces.Dict(obs_space_dict)
+        if self.observation_flatten:
+            self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
+            self.observation_space = gymnasium.spaces.Box(
+                -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
+            )
+        else:
+            self.observation_space = gymnasium.spaces.Dict(obs_space_dict)
 
     def obs(self):
         """Return the observation of our agent"""
@@ -195,15 +198,15 @@ def obs(self):
         else:
             obs['buttons_lidar'] = np.zeros(self.lidar_num_bins)
 
-        # if self.observe_vision:
-        #     obs['vision'] = self.obs_vision()
-        # if self.observation_flatten:
-        flat_obs = np.zeros(self.obs_flat_size)
-        offset = 0
-        for k in sorted(self.obs_space_dict.keys()):
-            k_size = np.prod(obs[k].shape)
-            flat_obs[offset : offset + k_size] = obs[k].flat
-            offset += k_size
-        obs = flat_obs
-        assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
+        if self.observe_vision:
+            obs['vision'] = self.obs_vision()
+        if self.observation_flatten:
+            flat_obs = np.zeros(self.obs_flat_size)
+            offset = 0
+            for k in sorted(self.obs_space_dict.keys()):
+                k_size = np.prod(obs[k].shape)
+                flat_obs[offset : offset + k_size] = obs[k].flat
+                offset += k_size
+            obs = flat_obs
+            assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
         return obs
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level0.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level0.py
index 0f5f5c950..cd39e6fdb 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level0.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level0.py
@@ -144,18 +144,26 @@ def build_observation_space(self):
 
         # if self.observe_goal_lidar:
         obs_space_dict['goal_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
+        if self.observe_vision:
+            width, height = self.vision_size
+            rows, cols = height, width
+            self.vision_size = (rows, cols)
+            obs_space_dict['vision'] = gymnasium.spaces.Box(
+                0, 255, self.vision_size + (3,), dtype=np.uint8
+            )
+
         # Flatten it ourselves
         self.obs_space_dict = obs_space_dict
-        # if self.observation_flatten:
-        self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
-        self.observation_space = gymnasium.spaces.Box(
-            -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
-        )
-        # else:
-        #     self.observation_space = gym.spaces.Dict(obs_space_dict)
+        if self.observation_flatten:
+            self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
+            self.observation_space = gymnasium.spaces.Box(
+                -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
+            )
+        else:
+            self.observation_space = gymnasium.spaces.Dict(obs_space_dict)
 
     def obs(self):
         """Return the observation of our agent"""
@@ -187,13 +195,15 @@ def obs(self):
             for sensor in self.robot.ballquat_names:
                 obs[sensor] = self.world.get_sensor(sensor)
 
-        # if self.observation_flatten:
-        flat_obs = np.zeros(self.obs_flat_size)
-        offset = 0
-        for k in sorted(self.obs_space_dict.keys()):
-            k_size = np.prod(obs[k].shape)
-            flat_obs[offset : offset + k_size] = obs[k].flat
-            offset += k_size
-        obs = flat_obs
-        assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
+        if self.observe_vision:
+            obs['vision'] = self.obs_vision()
+        if self.observation_flatten:
+            flat_obs = np.zeros(self.obs_flat_size)
+            offset = 0
+            for k in sorted(self.obs_space_dict.keys()):
+                k_size = np.prod(obs[k].shape)
+                flat_obs[offset : offset + k_size] = obs[k].flat
+                offset += k_size
+            obs = flat_obs
+            assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
         return obs
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level1.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level1.py
index 76fb4eaec..8e5bfff27 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level1.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/goal/goal_level1.py
@@ -103,16 +103,16 @@ def build_observation_space(self):
 
         for sensor in self.sensors_obs:  # Explicitly listed sensors
             dim = self.robot.sensor_dim[sensor]
-            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (dim,), dtype=np.float32)
+            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (dim,), dtype=np.float64)
         # Velocities don't have wraparound effects that rotational positions do
         # Wraparounds are not kind to neural networks
         # Whereas the angle 2*pi is very close to 0, this isn't true in the network
         # In theory the network could learn this, but in practice we simplify it
         # when the sensors_angle_components switch is enabled.
         for sensor in self.robot.hinge_vel_names:
-            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32)
+            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float64)
         for sensor in self.robot.ballangvel_names:
-            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (3,), dtype=np.float32)
+            obs_space_dict[sensor] = gymnasium.spaces.Box(-np.inf, np.inf, (3,), dtype=np.float64)
         # Angular positions have wraparound effects, so output something more friendly
         if self.sensors_angle_components:
             # Single joints are turned into sin(x), cos(x) pairs
@@ -120,7 +120,7 @@ def build_observation_space(self):
             # Since for angles, small perturbations in angle give small differences in sin/cos
             for sensor in self.robot.hinge_pos_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (2,), dtype=np.float32
+                    -np.inf, np.inf, (2,), dtype=np.float64
                 )
             # Quaternions are turned into 3x3 rotation matrices
             # Quaternions have a wraparound issue in how they are normalized,
@@ -134,43 +134,51 @@ def build_observation_space(self):
             # Instead we use a 3x3 rotation matrix, which if normalized, smoothly varies as well.
             for sensor in self.robot.ballquat_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (3, 3), dtype=np.float32
+                    -np.inf, np.inf, (3, 3), dtype=np.float64
                 )
         else:
             # Otherwise include the sensor without any processing
             # TODO: comparative study of the performance with and without this feature.
             for sensor in self.robot.hinge_pos_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (1,), dtype=np.float32
+                    -np.inf, np.inf, (1,), dtype=np.float64
                 )
             for sensor in self.robot.ballquat_names:
                 obs_space_dict[sensor] = gymnasium.spaces.Box(
-                    -np.inf, np.inf, (4,), dtype=np.float32
+                    -np.inf, np.inf, (4,), dtype=np.float64
                 )
 
         # if self.observe_goal_lidar:
         obs_space_dict['goal_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.observe_hazards:
         obs_space_dict['hazards_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
         # if self.observe_vases:
         obs_space_dict['vases_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
+        if self.observe_vision:
+            width, height = self.vision_size
+            rows, cols = height, width
+            self.vision_size = (rows, cols)
+            obs_space_dict['vision'] = gymnasium.spaces.Box(
+                0, 255, self.vision_size + (3,), dtype=np.uint8
+            )
+
         # Flatten it ourselves
         self.obs_space_dict = obs_space_dict
-        # if self.observation_flatten:
-        self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
-        self.observation_space = gymnasium.spaces.Box(
-            -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
-        )
-        # else:
-        #     self.observation_space = gym.spaces.Dict(obs_space_dict)
+        if self.observation_flatten:
+            self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
+            self.observation_space = gymnasium.spaces.Box(
+                -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
+            )
+        else:
+            self.observation_space = gymnasium.spaces.Dict(obs_space_dict)
 
     def obs(self):
         """Return the observation of our agent"""
@@ -187,13 +195,15 @@ def obs(self):
         # if self.observe_vases:
         obs['vases_lidar'] = self.obs_lidar(self.vases_pos, GROUP['vase'])
 
-        # if self.observation_flatten:
-        flat_obs = np.zeros(self.obs_flat_size)
-        offset = 0
-        for k in sorted(self.obs_space_dict.keys()):
-            k_size = np.prod(obs[k].shape)
-            flat_obs[offset : offset + k_size] = obs[k].flat
-            offset += k_size
-        obs = flat_obs
-        assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
+        if self.observe_vision:
+            obs['vision'] = self.obs_vision()
+        if self.observation_flatten:
+            flat_obs = np.zeros(self.obs_flat_size)
+            offset = 0
+            for k in sorted(self.obs_space_dict.keys()):
+                k_size = np.prod(obs[k].shape)
+                flat_obs[offset : offset + k_size] = obs[k].flat
+                offset += k_size
+            obs = flat_obs
+            assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
         return obs
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level0.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level0.py
index e06ee5e04..ee581aa9f 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level0.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level0.py
@@ -164,29 +164,31 @@ def build_observation_space(self):
         # obs_space_dict['box_compass'] = gym.spaces.Box(-1.0, 1.0, (self.compass_shape,), dtype=np.float32)
         # if self.observe_box_lidar:
         obs_space_dict['box_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.observe_goal_lidar:
         obs_space_dict['goal_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
-        # if self.observe_vision:
-        #     width, height = self.vision_size
-        #     rows, cols = height, width
-        #     self.vision_size = (rows, cols)
-        #     obs_space_dict['vision'] = gym.spaces.Box(0, 1.0, self.vision_size + (3,), dtype=np.float32)
+        if self.observe_vision:
+            width, height = self.vision_size
+            rows, cols = height, width
+            self.vision_size = (rows, cols)
+            obs_space_dict['vision'] = gymnasium.spaces.Box(
+                0, 255, self.vision_size + (3,), dtype=np.uint8
+            )
 
         # Flatten it ourselves
         self.obs_space_dict = obs_space_dict
-        # if self.observation_flatten:
-        self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
-        self.observation_space = gymnasium.spaces.Box(
-            -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
-        )
-        # else:
-        #     self.observation_space = gym.spaces.Dict(obs_space_dict)
+        if self.observation_flatten:
+            self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
+            self.observation_space = gymnasium.spaces.Box(
+                -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
+            )
+        else:
+            self.observation_space = gymnasium.spaces.Dict(obs_space_dict)
 
     def obs(self):
         """Return the observation of our agent"""
@@ -204,18 +206,17 @@ def obs(self):
 
         obs.update(self.get_sensor_obs())
 
-        # if self.observe_vision:
-        #     obs['vision'] = self.obs_vision()
-
-        # if self.observation_flatten:
-        flat_obs = np.zeros(self.obs_flat_size)
-        offset = 0
-        for k in sorted(self.obs_space_dict.keys()):
-            k_size = np.prod(obs[k].shape)
-            flat_obs[offset : offset + k_size] = obs[k].flat
-            offset += k_size
-        obs = flat_obs
-        assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
+        if self.observe_vision:
+            obs['vision'] = self.obs_vision()
+        if self.observation_flatten:
+            flat_obs = np.zeros(self.obs_flat_size)
+            offset = 0
+            for k in sorted(self.obs_space_dict.keys()):
+                k_size = np.prod(obs[k].shape)
+                flat_obs[offset : offset + k_size] = obs[k].flat
+                offset += k_size
+            obs = flat_obs
+            assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
         return obs
 
     def dist_box(self):
diff --git a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level1.py b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level1.py
index 76b72a946..83309ad13 100644
--- a/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level1.py
+++ b/omnisafe/envs/Safety_Gymnasium/safety_gymnasium/envs/safety_gym_v2/tasks/push/push_level1.py
@@ -133,38 +133,41 @@ def build_observation_space(self):
         #     obs_space_dict['box_compass'] = gym.spaces.Box(-1.0, 1.0, (self.compass_shape,), dtype=np.float32)
         # if self.observe_box_lidar:
         obs_space_dict['box_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.observe_goal_lidar:
         obs_space_dict['goal_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.observe_hazards:
         obs_space_dict['hazards_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
         # if self.pillars_num and self.observe_pillars:
         obs_space_dict['pillars_lidar'] = gymnasium.spaces.Box(
-            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float32
+            0.0, 1.0, (self.lidar_num_bins,), dtype=np.float64
         )
 
-        # if self.observe_vision:
-        #     width, height = self.vision_size
-        #     rows, cols = height, width
-        #     self.vision_size = (rows, cols)
-        #     obs_space_dict['vision'] = gym.spaces.Box(0, 1.0, self.vision_size + (3,), dtype=np.float32)
+        if self.observe_vision:
+            width, height = self.vision_size
+            rows, cols = height, width
+            self.vision_size = (rows, cols)
+            obs_space_dict['vision'] = gymnasium.spaces.Box(
+                0, 255, self.vision_size + (3,), dtype=np.uint8
+            )
+
         # Flatten it ourselves
         self.obs_space_dict = obs_space_dict
-        # if self.observation_flatten:
-        self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
-        self.observation_space = gymnasium.spaces.Box(
-            -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
-        )
-        # else:
-        #     self.observation_space = gym.spaces.Dict(obs_space_dict)
+        if self.observation_flatten:
+            self.obs_flat_size = sum([np.prod(i.shape) for i in self.obs_space_dict.values()])
+            self.observation_space = gymnasium.spaces.Box(
+                -np.inf, np.inf, (self.obs_flat_size,), dtype=np.float64
+            )
+        else:
+            self.observation_space = gymnasium.spaces.Dict(obs_space_dict)
 
     def obs(self):
         """Return the observation of our agent"""
@@ -188,16 +191,15 @@ def obs(self):
         # if self.pillars_num and self.observe_pillars:
         obs['pillars_lidar'] = self.obs_lidar(self.pillars_pos, GROUP['pillar'])
 
-        # if self.observe_vision:
-        #     obs['vision'] = self.obs_vision()
-
-        # if self.observation_flatten:
-        flat_obs = np.zeros(self.obs_flat_size)
-        offset = 0
-        for k in sorted(self.obs_space_dict.keys()):
-            k_size = np.prod(obs[k].shape)
-            flat_obs[offset : offset + k_size] = obs[k].flat
-            offset += k_size
-        obs = flat_obs
-        assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
+        if self.observe_vision:
+            obs['vision'] = self.obs_vision()
+        if self.observation_flatten:
+            flat_obs = np.zeros(self.obs_flat_size)
+            offset = 0
+            for k in sorted(self.obs_space_dict.keys()):
+                k_size = np.prod(obs[k].shape)
+                flat_obs[offset : offset + k_size] = obs[k].flat
+                offset += k_size
+            obs = flat_obs
+            assert self.observation_space.contains(obs), f'Bad obs {obs} {self.observation_space}'
         return obs


Test patch:
diff --git a/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2.py b/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2.py
index a5444d6ba..61bb74c56 100644
--- a/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2.py
+++ b/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2.py
@@ -9,7 +9,7 @@
 def test_off_policy(agent_id, env_id, level):
     """test_env"""
     env_name = 'Safety' + agent_id + env_id + level + '-v0'
-    env = safety_gymnasium.make(env_name)
+    env = safety_gymnasium.make(env_name, render_mode='rgb_array')
     obs, _ = env.reset()
     terminled = False
     ep_ret = 0
@@ -28,6 +28,5 @@ def test_off_policy(agent_id, env_id, level):
             max_ep_len = env._max_episode_steps
 
         obs, reward, cost, terminled, truncated, info = env.step(act)
-
         ep_ret += reward
         ep_cost += cost
diff --git a/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2_vision.py b/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2_vision.py
new file mode 100644
index 000000000..1d2eac4dd
--- /dev/null
+++ b/omnisafe/envs/Safety_Gymnasium/tests/test_safety_gym_v2_vision.py
@@ -0,0 +1,32 @@
+import safety_gymnasium
+
+import helpers
+
+
+@helpers.parametrize(
+    agent_id=['Point', 'Car'], env_id=['Goal', 'Push', 'Button'], level=['0', '1', '2']
+)
+def test_off_policy(agent_id, env_id, level):
+    """test_vision_env"""
+    env_name = 'Safety' + agent_id + env_id + level + 'Vision' + '-v0'
+    env = safety_gymnasium.make(env_name)
+
+    obs, _ = env.reset()
+    terminled = False
+    ep_ret = 0
+    ep_cost = 0
+    for step in range(1000):
+        if terminled:
+            print('Episode Return: %.3f \t Episode Cost: %.3f' % (ep_ret, ep_cost))
+            ep_ret, ep_cost = 0, 0
+            obs, _ = env.reset()
+        assert env.observation_space.contains(obs)
+        act = env.action_space.sample()
+        assert env.action_space.contains(act)
+        # Use the environment's built_in max_episode_steps
+        if hasattr(env, '_max_episode_steps'):
+            max_ep_len = env._max_episode_steps
+        obs, reward, cost, terminled, truncated, info = env.step(act)
+
+        ep_ret += reward
+        ep_cost += cost