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

Issue:
[BUG] Why can I make Gymnasium environments when using Safety-Gymnasium?
### Required prerequisites

- [X] I have read the documentation <https://omnisafe.readthedocs.io>.
- [X] I have searched the [Issue Tracker](https://github.com/PKU-MARL/omnisafe/issues) and [Discussions](https://github.com/PKU-MARL/omnisafe/discussions) that this hasn't already been reported. (+1 or comment there if it has.)
- [X] Consider asking first in a [Discussion](https://github.com/PKU-MARL/omnisafe/discussions/new).

### What version of OmniSafe are you using?

0.01

### System information

3.8.15 (default, Nov 24 2022, 15:19:38) 
[GCC 11.2.0] linux
0.0.1

### Problem description

When I test safety-gymnasium, I wrongly input 'Humanoid-v4' as env_id, but it works and something goes run.


### Reproducible example code

The Python snippets:

```python
import argparse

import safety_gymnasium


def run_random(env_name):
    """Random run."""
    env = safety_gymnasium.make(env_name, render_mode='human')
    obs, _ = env.reset()
    # Use below to specify seed.
    # obs, _ = env.reset(seed=0)
    terminated, truncated = False, False
    ep_ret, ep_cost = 0, 0
    while True:
        if terminated or truncated:
            print(f'Episode Return: {ep_ret} \t Episode Cost: {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'):  # pylint: disable=unused-variable
            max_ep_len = env._max_episode_steps  # pylint: disable=unused-variable,protected-access
        # pylint: disable-next=unused-variable
        obs, reward, cost, terminated, truncated, info = env.step(act)

        ep_ret += reward
        ep_cost += cost


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--env', default='Humanoid-v4')
    args = parser.parse_args()
    run_random(args.env)

```

Command lines:

```bash

```

Extra dependencies:

```text

```

Steps to reproduce:

1. run code above.


### Traceback

```pytb
Traceback (most recent call last):
  File "~/omnisafe/envs/safety_gymnasium/examples/env.py", line 53, in <module>
    run_random(args.env)
  File "~/omnisafe/envs/safety_gymnasium/examples/env.py", line 42, in run_random
    obs, reward, cost, terminated, truncated, info = env.step(act)
  File "~/omnisafe/envs/safety-gymnasium/safety_gymnasium/wrappers/time_limit.py", line 45, in step
    observation, reward, cost, terminated, truncated, info = self.env.step(action)
ValueError: not enough values to unpack (expected 6, got 5)
```


### Expected behavior

_No response_

### Additional context

_No response_


Discussion / hints:

Fix patch:
diff --git a/.pylintrc b/.pylintrc
index 5bd5b63ec..dbe86bd5c 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -184,7 +184,10 @@ good-names=i,
            kl,
            Jc,
            Fvp,
-           xHx
+           xHx,
+           id,
+           e,
+           f,
 
 # Good variable names regexes, separated by a comma. If names match any regex,
 # they will always be accepted
diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt
index a7ddd068a..ff9093fa4 100644
--- a/docs/source/spelling_wordlist.txt
+++ b/docs/source/spelling_wordlist.txt
@@ -340,5 +340,10 @@ Shie
 Mannor
 https
 neurips
+boolean
+autoreset
+eg
+bools
+dtypes
 vectorized
 async
diff --git a/envs/safety-gymnasium/examples/env.py b/envs/safety-gymnasium/examples/env.py
index 648341ed2..5cfd868e0 100644
--- a/envs/safety-gymnasium/examples/env.py
+++ b/envs/safety-gymnasium/examples/env.py
@@ -22,7 +22,7 @@
 def run_random(env_name):
     """Random run."""
     env = safety_gymnasium.make(env_name, render_mode='human')
-    obs, _ = env.reset()
+    obs, info = env.reset()  # pylint: disable=unused-variable
     # Use below to specify seed.
     # obs, _ = env.reset(seed=0)
     terminated, truncated = False, False
@@ -31,7 +31,7 @@ def run_random(env_name):
         if terminated or truncated:
             print(f'Episode Return: {ep_ret} \t Episode Cost: {ep_cost}')
             ep_ret, ep_cost = 0, 0
-            obs, _ = env.reset()
+            obs, info = env.reset()  # pylint: disable=unused-variable
         assert env.observation_space.contains(obs)
         act = env.action_space.sample()
         assert env.action_space.contains(act)
@@ -48,6 +48,6 @@ def run_random(env_name):
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--env', default='SafetyRacecarCircle0-v0')
+    parser.add_argument('--env', default='SafetyAntVelocity-v4')
     args = parser.parse_args()
     run_random(args.env)
diff --git a/envs/safety-gymnasium/examples/vision_env.py b/envs/safety-gymnasium/examples/vision_env.py
index dd88d8d48..bb8d72fab 100644
--- a/envs/safety-gymnasium/examples/vision_env.py
+++ b/envs/safety-gymnasium/examples/vision_env.py
@@ -27,7 +27,7 @@
 def run_random(env_name):
     """Random run."""
     env = safety_gymnasium.make(env_name)
-    obs, _ = env.reset()
+    obs, info = env.reset()  # pylint: disable=unused-variable
     # Use below to specify seed.
     # obs, _ = env.reset(seed=0)
     terminated, truncated = False, False
@@ -37,7 +37,7 @@ def run_random(env_name):
         if terminated or truncated:
             print(f'Episode Return: {ep_ret} \t Episode Cost: {ep_cost}')
             ep_ret, ep_cost = 0, 0
-            obs, _ = env.reset()
+            obs, info = env.reset()  # pylint: disable=unused-variable
             save_video(
                 frames=render_list,
                 video_folder=DIR,
@@ -62,6 +62,6 @@ def run_random(env_name):
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--env', default='SafetyCarGoal0Vision-v0')
+    parser.add_argument('--env', default='SafetyCarRun0Vision-v0')
     args = parser.parse_args()
     run_random(args.env)
diff --git a/envs/safety-gymnasium/safety_gymnasium/__init__.py b/envs/safety-gymnasium/safety_gymnasium/__init__.py
index 344f9ef14..308759b5e 100644
--- a/envs/safety-gymnasium/safety_gymnasium/__init__.py
+++ b/envs/safety-gymnasium/safety_gymnasium/__init__.py
@@ -25,7 +25,7 @@
 VERSION = 'v0'
 
 ROBOT_NAMES = ('Point', 'Car', 'Racecar')
-ROBOT_XMLS = {name: f'xmls/{name.lower()}.xml' for name in ROBOT_NAMES}
+ROBOT_XMLS = {name: f'assets/xmls/{name.lower()}.xml' for name in ROBOT_NAMES}
 BASE_SENSORS = ['accelerometer', 'velocimeter', 'gyro', 'magnetometer']
 EXTRA_SENSORS = {
     'Doggo': [
@@ -150,3 +150,45 @@ def combine(tasks, agents, max_episode_steps):
 
 run_tasks = {'Run0': {}}
 combine(run_tasks, robot_configs, max_episode_steps=500)
+
+
+# Safety Velocity
+# ----------------------------------------
+register(
+    id='SafetyHalfCheetahVelocity-v4',
+    entry_point='safety_gymnasium.tasks.safety_velocity.safety_half_cheetah_velocity:SafetyHalfCheetahVelocityEnv',
+    max_episode_steps=1000,
+    reward_threshold=4800.0,
+)
+
+register(
+    id='SafetyHopperVelocity-v4',
+    entry_point='safety_gymnasium.tasks.safety_velocity.safety_hopper_velocity:SafetyHopperVelocityEnv',
+    max_episode_steps=1000,
+    reward_threshold=3800.0,
+)
+
+register(
+    id='SafetySwimmerVelocity-v4',
+    entry_point='safety_gymnasium.tasks.safety_velocity.safety_swimmer_velocity:SafetySwimmerVelocityEnv',
+    max_episode_steps=1000,
+    reward_threshold=360.0,
+)
+
+register(
+    id='SafetyWalker2dVelocity-v4',
+    max_episode_steps=1000,
+    entry_point='safety_gymnasium.tasks.safety_velocity.safety_walker2d_velocity:SafetyWalker2dVelocityEnv',
+)
+
+register(
+    id='SafetyAntVelocity-v4',
+    entry_point='safety_gymnasium.tasks.safety_velocity.safety_ant_velocity:SafetyAntVelocityEnv',
+    max_episode_steps=1000,
+    reward_threshold=6000.0,
+)
+register(
+    id='SafetyHumanoidVelocity-v4',
+    entry_point='safety_gymnasium.tasks.safety_velocity.safety_humanoid_velocity:SafetyHumanoidVelocityEnv',
+    max_episode_steps=1000,
+)
diff --git a/envs/safety-gymnasium/safety_gymnasium/assets/robot.py b/envs/safety-gymnasium/safety_gymnasium/assets/robot.py
index ad8590923..11b85ac4a 100644
--- a/envs/safety-gymnasium/safety_gymnasium/assets/robot.py
+++ b/envs/safety-gymnasium/safety_gymnasium/assets/robot.py
@@ -33,7 +33,7 @@ class Robot:  # pylint: disable=too-many-instance-attributes
     placements: list = None  # Robot placements list (defaults to full extents)
     locations: list = field(default_factory=list)  # Explicitly place robot XY coordinate
     keepout: float = 0.4  # Needs to be set to match the robot XML used
-    base: str = 'xmls/car.xml'  # Which robot XML to use as the base
+    base: str = 'assets/xmls/car.xml'  # Which robot XML to use as the base
     rot: float = None  # Override robot starting angle
 
     def __post_init__(self, path):
diff --git a/envs/safety-gymnasium/safety_gymnasium/xmls/car.xml b/envs/safety-gymnasium/safety_gymnasium/assets/xmls/car.xml
similarity index 100%
rename from envs/safety-gymnasium/safety_gymnasium/xmls/car.xml
rename to envs/safety-gymnasium/safety_gymnasium/assets/xmls/car.xml
diff --git a/envs/safety-gymnasium/safety_gymnasium/xmls/doggo.xml b/envs/safety-gymnasium/safety_gymnasium/assets/xmls/doggo.xml
similarity index 100%
rename from envs/safety-gymnasium/safety_gymnasium/xmls/doggo.xml
rename to envs/safety-gymnasium/safety_gymnasium/assets/xmls/doggo.xml
diff --git a/envs/safety-gymnasium/safety_gymnasium/xmls/point.xml b/envs/safety-gymnasium/safety_gymnasium/assets/xmls/point.xml
similarity index 100%
rename from envs/safety-gymnasium/safety_gymnasium/xmls/point.xml
rename to envs/safety-gymnasium/safety_gymnasium/assets/xmls/point.xml
diff --git a/envs/safety-gymnasium/safety_gymnasium/xmls/racecar.xml b/envs/safety-gymnasium/safety_gymnasium/assets/xmls/racecar.xml
similarity index 99%
rename from envs/safety-gymnasium/safety_gymnasium/xmls/racecar.xml
rename to envs/safety-gymnasium/safety_gymnasium/assets/xmls/racecar.xml
index 9105cb8bb..228d37524 100644
--- a/envs/safety-gymnasium/safety_gymnasium/xmls/racecar.xml
+++ b/envs/safety-gymnasium/safety_gymnasium/assets/xmls/racecar.xml
@@ -16,7 +16,7 @@
 
 <mujoco>
     <size njmax="3000" nconmax="1000"/>
-    <compiler angle="radian" meshdir="../assets/meshes" autolimits="true" />
+    <compiler angle="radian" meshdir="../meshes" autolimits="true" />
     <option integrator="RK4" timestep="0.004"/>
     <default>
         <geom condim="6" density="5" rgba="1 0 0 1" size=".05" type="sphere"/>
diff --git a/envs/safety-gymnasium/safety_gymnasium/bases/mujoco_env.py b/envs/safety-gymnasium/safety_gymnasium/bases/mujoco_env.py
deleted file mode 100644
index a42ad0f95..000000000
--- a/envs/safety-gymnasium/safety_gymnasium/bases/mujoco_env.py
+++ /dev/null
@@ -1,429 +0,0 @@
-# Copyright 2022 OmniSafe Team. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-# pylint: disable=all
-
-import os
-from typing import Optional, Union
-
-import gymnasium
-import numpy as np
-from gymnasium import error, logger, spaces
-from gymnasium.spaces import Space
-
-
-try:
-    import mujoco_py
-except ImportError as e:
-    MUJOCO_PY_IMPORT_ERROR = e
-else:
-    MUJOCO_PY_IMPORT_ERROR = None
-
-try:
-    import mujoco
-except ImportError as e:
-    MUJOCO_IMPORT_ERROR = e
-else:
-    MUJOCO_IMPORT_ERROR = None
-
-
-DEFAULT_SIZE = 480
-
-
-class BaseMujocoEnv(gymnasium.Env):
-    """Superclass for all MuJoCo environments."""
-
-    def __init__(
-        self,
-        model_path,
-        frame_skip,
-        observation_space: Space,
-        render_mode: Optional[str] = None,
-        width: int = DEFAULT_SIZE,
-        height: int = DEFAULT_SIZE,
-        camera_id: Optional[int] = None,
-        camera_name: Optional[str] = None,
-    ):
-        if model_path.startswith('/'):
-            self.fullpath = model_path
-        else:
-            self.fullpath = os.path.join(os.path.dirname(__file__), 'assets', model_path)
-        if not os.path.exists(self.fullpath):
-            raise OSError(f'File {self.fullpath} does not exist')
-
-        self.width = width
-        self.height = height
-        self._initialize_simulation()  # may use width and height
-
-        self.init_qpos = self.data.qpos.ravel().copy()
-        self.init_qvel = self.data.qvel.ravel().copy()
-        self._viewers = {}
-
-        self.frame_skip = frame_skip
-
-        self.viewer = None
-
-        assert self.metadata['render_modes'] == [
-            'human',
-            'rgb_array',
-            'depth_array',
-        ], self.metadata['render_modes']
-        assert (
-            int(np.round(1.0 / self.dt)) == self.metadata['render_fps']
-        ), f'Expected value: {int(np.round(1.0 / self.dt))}, Actual value: {self.metadata["render_fps"]}'
-
-        self.observation_space = observation_space
-        self._set_action_space()
-
-        self.render_mode = render_mode
-        self.camera_name = camera_name
-        self.camera_id = camera_id
-
-    def _set_action_space(self):
-        bounds = self.model.actuator_ctrlrange.copy().astype(np.float32)
-        low, high = bounds.T
-        self.action_space = spaces.Box(low=low, high=high, dtype=np.float32)
-        return self.action_space
-
-    # methods to override:
-    # ----------------------------
-
-    def reset_model(self):
-        """
-        Reset the robot degrees of freedom (qpos and qvel).
-        Implement this in each subclass.
-        """
-        raise NotImplementedError
-
-    def viewer_setup(self):
-        """
-        This method is called when the viewer is initialized.
-        Optionally implement this method, if you need to tinker with camera position and so forth.
-        """
-
-    def _initialize_simulation(self):
-        """
-        Initialize MuJoCo simulation data structures mjModel and mjData.
-        """
-        raise NotImplementedError
-
-    def _reset_simulation(self):
-        """
-        Reset MuJoCo simulation data structures, mjModel and mjData.
-        """
-        raise NotImplementedError
-
-    def _step_mujoco_simulation(self, ctrl, n_frames):
-        """
-        Step over the MuJoCo simulation.
-        """
-        raise NotImplementedError
-
-    def render(self):
-        """
-        Render a frame from the MuJoCo simulation as specified by the render_mode.
-        """
-        raise NotImplementedError
-
-    # -----------------------------
-
-    def reset(
-        self,
-        *,
-        seed: Optional[int] = None,
-        options: Optional[dict] = None,
-    ):
-        super().reset(seed=seed)
-
-        self._reset_simulation()
-
-        ob = self.reset_model()
-        if self.render_mode == 'human':
-            self.render()
-        return ob, {}
-
-    def set_state(self, qpos, qvel):
-        """
-        Set the joints position qpos and velocity qvel of the model. Override this method depending on the MuJoCo bindings used.
-        """
-        assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,)
-
-    @property
-    def dt(self):
-        return self.model.opt.timestep * self.frame_skip
-
-    def do_simulation(self, ctrl, n_frames):
-        """
-        Step the simulation n number of frames and applying a control action.
-        """
-        # Check control input is contained in the action space
-        if np.array(ctrl).shape != self.action_space.shape:
-            raise ValueError('Action dimension mismatch')
-        self._step_mujoco_simulation(ctrl, n_frames)
-
-    def close(self):
-        if self.viewer is not None:
-            self.viewer = None
-            self._viewers = {}
-
-    def get_body_com(self, body_name):
-        """Return the cartesian position of a body frame"""
-        raise NotImplementedError
-
-    def state_vector(self):
-        """Return the position and velocity joint states of the model"""
-        return np.concatenate([self.data.qpos.flat, self.data.qvel.flat])
-
-
-class MuJocoPyEnv(BaseMujocoEnv):
-    def __init__(
-        self,
-        model_path: str,
-        frame_skip: int,
-        observation_space: Space,
-        render_mode: Optional[str] = None,
-        width: int = DEFAULT_SIZE,
-        height: int = DEFAULT_SIZE,
-        camera_id: Optional[int] = None,
-        camera_name: Optional[str] = None,
-    ):
-        if MUJOCO_PY_IMPORT_ERROR is not None:
-            raise error.DependencyNotInstalled(
-                f'{MUJOCO_PY_IMPORT_ERROR}. (HINT: you need to install mujoco_py, and also perform the setup instructions here: https://github.com/openai/mujoco-py/.)'
-            )
-
-        logger.warn(
-            'This version of the mujoco environments depends '
-            'on the mujoco-py bindings, which are no longer maintained '
-            'and may stop working. Please upgrade to the v4 versions of '
-            'the environments (which depend on the mujoco python bindings instead), unless '
-            'you are trying to precisely replicate previous works).'
-        )
-
-        super().__init__(
-            model_path,
-            frame_skip,
-            observation_space,
-            render_mode,
-            width,
-            height,
-            camera_id,
-            camera_name,
-        )
-
-    def _initialize_simulation(self):
-        self.model = mujoco_py.load_model_from_path(self.fullpath)
-        self.sim = mujoco_py.MjSim(self.model)
-        self.data = self.sim.data
-
-    def _reset_simulation(self):
-        self.sim.reset()
-
-    def set_state(self, qpos, qvel):
-        super().set_state(qpos, qvel)
-        state = self.sim.get_state()
-        state = mujoco_py.MjSimState(state.time, qpos, qvel, state.act, state.udd_state)
-        self.sim.set_state(state)
-        self.sim.forward()
-
-    def _step_mujoco_simulation(self, ctrl, n_frames):
-        self.sim.data.ctrl[:] = ctrl
-
-        for _ in range(n_frames):
-            self.sim.step()
-
-    def render(self):
-        width, height = self.width, self.height
-        camera_name, camera_id = self.camera_name, self.camera_id
-        if self.render_mode in {'rgb_array', 'depth_array'}:
-            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_name = 'track'
-
-            if camera_id is None and camera_name in self.model._camera_name2id:
-                if camera_name in self.model._camera_name2id:
-                    camera_id = self.model.camera_name2id(camera_name)
-
-                self._get_viewer(self.render_mode).render(width, height, camera_id=camera_id)
-
-        if self.render_mode == 'rgb_array':
-            data = self._get_viewer(self.render_mode).read_pixels(width, height, depth=False)
-            # original image is upside-down, so flip it
-            return data[::-1, :, :]
-        elif self.render_mode == 'depth_array':
-            self._get_viewer(self.render_mode).render(width, height)
-            # Extract depth part of the read_pixels() tuple
-            data = self._get_viewer(self.render_mode).read_pixels(width, height, depth=True)[1]
-            # original image is upside-down, so flip it
-            return data[::-1, :]
-        elif self.render_mode == 'human':
-            self._get_viewer(self.render_mode).render()
-
-    def _get_viewer(
-        self, mode
-    ) -> Union['mujoco_py.MjViewer', 'mujoco_py.MjRenderContextOffscreen']:
-        self.viewer = self._viewers.get(mode)
-        if self.viewer is None:
-            if mode == 'human':
-                self.viewer = mujoco_py.MjViewer(self.sim)
-
-            elif mode in {'rgb_array', 'depth_array'}:
-                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
-            else:
-                raise AttributeError(
-                    f"Unknown mode: {mode}, expected modes: {self.metadata['render_modes']}"
-                )
-
-            self.viewer_setup()
-            self._viewers[mode] = self.viewer
-
-        return self.viewer
-
-    def get_body_com(self, body_name):
-        return self.data.get_body_xpos(body_name)
-
-
-class MujocoEnv(BaseMujocoEnv):
-    """Superclass for MuJoCo environments."""
-
-    def __init__(
-        self,
-        model_path,
-        frame_skip,
-        observation_space: Space,
-        render_mode: Optional[str] = None,
-        width: int = DEFAULT_SIZE,
-        height: int = DEFAULT_SIZE,
-        camera_id: Optional[int] = None,
-        camera_name: Optional[str] = None,
-    ):
-        if MUJOCO_IMPORT_ERROR is not None:
-            raise error.DependencyNotInstalled(
-                f'{MUJOCO_IMPORT_ERROR}. (HINT: you need to install mujoco)'
-            )
-        super().__init__(
-            model_path,
-            frame_skip,
-            observation_space,
-            render_mode,
-            width,
-            height,
-            camera_id,
-            camera_name,
-        )
-
-    def _initialize_simulation(self):
-        self.model = mujoco.MjModel.from_xml_path(self.fullpath)
-        # MjrContext will copy model.vis.global_.off* to con.off*
-        self.model.vis.global_.offwidth = self.width
-        self.model.vis.global_.offheight = self.height
-        self.data = mujoco.MjData(self.model)
-
-    def _reset_simulation(self):
-        mujoco.mj_resetData(self.model, self.data)
-
-    def set_state(self, qpos, qvel):
-        super().set_state(qpos, qvel)
-        self.data.qpos[:] = np.copy(qpos)
-        self.data.qvel[:] = np.copy(qvel)
-        if self.model.na == 0:
-            self.data.act[:] = None
-        mujoco.mj_forward(self.model, self.data)  # pylint: disable=no-member
-
-    def _step_mujoco_simulation(self, ctrl, n_frames):
-        self.data.ctrl[:] = ctrl
-
-        mujoco.mj_step(self.model, self.data, nstep=self.frame_skip)
-
-        # As of MuJoCo 2.0, force-related quantities like cacc are not computed
-        # unless there's a force sensor in the model.
-        # See https://github.com/openai/gym/issues/1541
-        mujoco.mj_rnePostConstraint(self.model, self.data)
-
-    def render(self):
-        if self.render_mode in {
-            'rgb_array',
-            'depth_array',
-        }:
-            camera_id = self.camera_id
-            camera_name = self.camera_name
-
-            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_name = 'track'
-
-            if camera_id is None:
-                camera_id = mujoco.mj_name2id(
-                    self.model,
-                    mujoco.mjtObj.mjOBJ_CAMERA,
-                    camera_name,
-                )
-
-                self._get_viewer(self.render_mode).render(camera_id=camera_id)
-
-        if self.render_mode == 'rgb_array':
-            data = self._get_viewer(self.render_mode).read_pixels(depth=False)
-            # original image is upside-down, so flip it
-            return data[::-1, :, :]
-        elif self.render_mode == 'depth_array':
-            self._get_viewer(self.render_mode).render()
-            # Extract depth part of the read_pixels() tuple
-            data = self._get_viewer(self.render_mode).read_pixels(depth=True)[1]
-            # original image is upside-down, so flip it
-            return data[::-1, :]
-        elif self.render_mode == 'human':
-            self._get_viewer(self.render_mode).render()
-
-    def close(self):
-        if self.viewer is not None:
-            self.viewer.close()
-        super().close()
-
-    def _get_viewer(
-        self, mode
-    ) -> Union[
-        'gymnasium.envs.mujoco.mujoco_rendering.Viewer',
-        'gymnasium.envs.mujoco.mujoco_rendering.RenderContextOffscreen',
-    ]:
-        self.viewer = self._viewers.get(mode)
-        if self.viewer is None:
-            if mode == 'human':
-                from gymnasium.envs.mujoco.mujoco_rendering import Viewer
-
-                self.viewer = Viewer(self.model, self.data)
-            elif mode in {'rgb_array', 'depth_array'}:
-                from gymnasium.envs.mujoco.mujoco_rendering import RenderContextOffscreen
-
-                self.viewer = RenderContextOffscreen(self.model, self.data)
-            else:
-                raise AttributeError(
-                    f"Unexpected mode: {mode}, expected modes: {self.metadata['render_modes']}"
-                )
-
-            self.viewer_setup()
-            self._viewers[mode] = self.viewer
-        return self.viewer
-
-    def get_body_com(self, body_name):
-        return self.data.body(body_name).xpos
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/circle/circle.py b/envs/safety-gymnasium/safety_gymnasium/tasks/circle/circle.py
deleted file mode 100644
index 99bb9f7e3..000000000
--- a/envs/safety-gymnasium/safety_gymnasium/tasks/circle/circle.py
+++ /dev/null
@@ -1,104 +0,0 @@
-# Copyright 2022 OmniSafe Team. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Circle."""
-
-import mujoco
-import numpy as np
-from safety_gymnasium.assets.geoms import Circle, Sigwalls
-from safety_gymnasium.bases import BaseTask
-
-
-class CircleLevel0(BaseTask):
-    """A robot want to loop around the boundary of circle, while avoid going outside the boundary."""
-
-    def __init__(self, config):
-        super().__init__(config=config)
-
-        self.num_steps = 500
-
-        robot_placements_square = 0.8
-        self.robot.placements = [
-            (
-                -robot_placements_square,
-                -robot_placements_square,
-                robot_placements_square,
-                robot_placements_square,
-            )
-        ]
-        self.robot.keepout = 0
-
-        self.lidar_max_dist = 6
-
-        self.reward_factor: float = (
-            1e-1  # Reward for circle goal (complicated formula depending on pos and vel)
-        )
-
-        self.add_geoms(Circle(), Sigwalls())
-
-        self.specific_agent_config()
-
-    def calculate_cost(self):
-        """There are costs only when agent go out of the boundary."""
-        # pylint: disable-next=no-member
-        mujoco.mj_forward(self.model, self.data)  # Ensure positions and contacts are correct
-        cost = {}
-
-        cost['cost_out_of_boundary'] = self.robot_pos[0] > 1.125 or self.robot_pos[0] < -1.125
-
-        # Sum all costs into single total cost
-        cost['cost'] = sum(v for k, v in cost.items() if k.startswith('cost_'))
-        return cost
-
-    def calculate_reward(self):
-        """The agent should loop around the boundary of circle."""
-        reward = 0.0
-        # Circle environment reward
-        robot_com = self.world.robot_com()
-        robot_vel = self.world.robot_vel()
-        x, y, _ = robot_com  # pylint: disable=invalid-name
-        u, v, _ = robot_vel  # pylint: disable=invalid-name
-        radius = np.sqrt(x**2 + y**2)
-        # pylint: disable-next=no-member
-        reward += (
-            # pylint: disable-next=no-member
-            ((-u * y + v * x) / radius)
-            / (1 + np.abs(radius - self.circle.radius))  # pylint: disable=no-member
-        ) * self.reward_factor
-        return reward
-
-    def specific_agent_config(self):
-        pass
-
-    def specific_reset(self):
-        pass
-
-    def specific_step(self):
-        pass
-
-    def build_goal(self):
-        pass
-
-    def update_world(self):
-        pass
-
-    @property
-    def goal_achieved(self):
-        """Weather the goal of task is achieved."""
-        return False
-
-    @property
-    def circle_pos(self):
-        """Helper to get circle position from layout."""
-        return [[0, 0, 0]]
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/__init__.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/__init__.py
new file mode 100644
index 000000000..edaffa3b5
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Safe velocity tasks."""
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_ant_velocity.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_ant_velocity.py
new file mode 100644
index 000000000..11b679cbe
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_ant_velocity.py
@@ -0,0 +1,79 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Ant environment with a safety constraint on velocity."""
+
+import numpy as np
+from gymnasium.envs.mujoco.ant_v4 import AntEnv
+from safety_gymnasium.utils.task_utils import add_velocity_marker, clear_viewer
+
+
+class SafetyAntVelocityEnv(AntEnv):
+    """Ant environment with a safety constraint on velocity."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._velocity_threshold = 2.5745
+        self.model.light(0).castshadow = False
+
+    def step(self, action):  # pylint: disable=too-many-locals
+        xy_position_before = self.get_body_com('torso')[:2].copy()
+        self.do_simulation(action, self.frame_skip)
+        xy_position_after = self.get_body_com('torso')[:2].copy()
+
+        xy_velocity = (xy_position_after - xy_position_before) / self.dt
+        x_velocity, y_velocity = xy_velocity
+
+        forward_reward = x_velocity
+        healthy_reward = self.healthy_reward
+
+        rewards = forward_reward + healthy_reward
+
+        costs = ctrl_cost = self.control_cost(action)
+
+        terminated = self.terminated
+        observation = self._get_obs()
+        info = {
+            'reward_forward': forward_reward,
+            'reward_ctrl': -ctrl_cost,
+            'reward_survive': healthy_reward,
+            'x_position': xy_position_after[0],
+            'y_position': xy_position_after[1],
+            'distance_from_origin': np.linalg.norm(xy_position_after, ord=2),
+            'x_velocity': x_velocity,
+            'y_velocity': y_velocity,
+            'forward_reward': forward_reward,
+        }
+        if self._use_contact_forces:
+            contact_cost = self.contact_cost
+            costs += contact_cost
+            info['reward_ctrl'] = -contact_cost
+
+        reward = rewards - costs
+
+        velocity = np.sqrt(x_velocity**2 + y_velocity**2)
+        cost = velocity > self._velocity_threshold
+
+        if self.viewer:
+            clear_viewer(self.viewer)
+            add_velocity_marker(
+                viewer=self.viewer,
+                pos=self.get_body_com('torso')[:3].copy(),
+                vel=velocity,
+                cost=cost,
+                velocity_threshold=self._velocity_threshold,
+            )
+        if self.render_mode == 'human':
+            self.render()
+        return observation, reward, cost, terminated, False, info
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_half_cheetah_velocity.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_half_cheetah_velocity.py
new file mode 100644
index 000000000..d5e0a5e8b
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_half_cheetah_velocity.py
@@ -0,0 +1,62 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""HalfCheetah environment with a safety constraint on velocity."""
+
+from gymnasium.envs.mujoco.half_cheetah_v4 import HalfCheetahEnv
+from safety_gymnasium.utils.task_utils import add_velocity_marker, clear_viewer
+
+
+class SafetyHalfCheetahVelocityEnv(HalfCheetahEnv):
+    """HalfCheetah environment with a safety constraint on velocity."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._velocity_threshold = 2.8795
+        self.model.light(0).castshadow = False
+
+    def step(self, action):
+        x_position_before = self.data.qpos[0]
+        self.do_simulation(action, self.frame_skip)
+        x_position_after = self.data.qpos[0]
+        x_velocity = (x_position_after - x_position_before) / self.dt
+
+        ctrl_cost = self.control_cost(action)
+
+        forward_reward = self._forward_reward_weight * x_velocity
+
+        observation = self._get_obs()
+        reward = forward_reward - ctrl_cost
+        terminated = False
+        info = {
+            'x_position': x_position_after,
+            'x_velocity': x_velocity,
+            'reward_run': forward_reward,
+            'reward_ctrl': -ctrl_cost,
+        }
+
+        cost = x_velocity > self._velocity_threshold
+
+        if self.viewer:
+            clear_viewer(self.viewer)
+            add_velocity_marker(
+                viewer=self.viewer,
+                pos=self.get_body_com('torso')[:3].copy(),
+                vel=x_velocity,
+                cost=cost,
+                velocity_threshold=self._velocity_threshold,
+            )
+        if self.render_mode == 'human':
+            self.render()
+        return observation, reward, cost, terminated, False, info
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_hopper_velocity.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_hopper_velocity.py
new file mode 100644
index 000000000..8c516ea76
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_hopper_velocity.py
@@ -0,0 +1,64 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Hopper environment with a safety constraint on velocity."""
+
+from gymnasium.envs.mujoco.hopper_v4 import HopperEnv
+from safety_gymnasium.utils.task_utils import add_velocity_marker, clear_viewer
+
+
+class SafetyHopperVelocityEnv(HopperEnv):
+    """Hopper environment with a safety constraint on velocity."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._velocity_threshold = 0.37315
+        self.model.light(0).castshadow = False
+
+    def step(self, action):
+        x_position_before = self.data.qpos[0]
+        self.do_simulation(action, self.frame_skip)
+        x_position_after = self.data.qpos[0]
+        x_velocity = (x_position_after - x_position_before) / self.dt
+
+        ctrl_cost = self.control_cost(action)
+
+        forward_reward = self._forward_reward_weight * x_velocity
+        healthy_reward = self.healthy_reward
+
+        rewards = forward_reward + healthy_reward
+        costs = ctrl_cost
+
+        observation = self._get_obs()
+        reward = rewards - costs
+        terminated = self.terminated
+        info = {
+            'x_position': x_position_after,
+            'x_velocity': x_velocity,
+        }
+
+        cost = x_velocity > self._velocity_threshold
+
+        if self.viewer:
+            clear_viewer(self.viewer)
+            add_velocity_marker(
+                viewer=self.viewer,
+                pos=self.get_body_com('torso')[:3].copy(),
+                vel=x_velocity,
+                cost=cost,
+                velocity_threshold=self._velocity_threshold,
+            )
+        if self.render_mode == 'human':
+            self.render()
+        return observation, reward, cost, terminated, False, info
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_humanoid_velocity.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_humanoid_velocity.py
new file mode 100644
index 000000000..b06aafb9b
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_humanoid_velocity.py
@@ -0,0 +1,74 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Humanoid environment with a safety constraint on velocity."""
+
+import numpy as np
+from gymnasium.envs.mujoco.humanoid_v4 import HumanoidEnv, mass_center
+from safety_gymnasium.utils.task_utils import add_velocity_marker, clear_viewer
+
+
+class SafetyHumanoidVelocityEnv(HumanoidEnv):
+    """Humanoid environment with a safety constraint on velocity."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._velocity_threshold = 2.3475
+        self.model.light(0).castshadow = False
+
+    def step(self, action):  # pylint: disable=too-many-locals
+        xy_position_before = mass_center(self.model, self.data)
+        self.do_simulation(action, self.frame_skip)
+        xy_position_after = mass_center(self.model, self.data)
+
+        xy_velocity = (xy_position_after - xy_position_before) / self.dt
+        x_velocity, y_velocity = xy_velocity
+
+        ctrl_cost = self.control_cost(action)
+
+        forward_reward = self._forward_reward_weight * x_velocity
+        healthy_reward = self.healthy_reward
+
+        rewards = forward_reward + healthy_reward
+
+        observation = self._get_obs()
+        reward = rewards - ctrl_cost
+        terminated = self.terminated
+        info = {
+            'reward_linvel': forward_reward,
+            'reward_quadctrl': -ctrl_cost,
+            'reward_alive': healthy_reward,
+            'x_position': xy_position_after[0],
+            'y_position': xy_position_after[1],
+            'distance_from_origin': np.linalg.norm(xy_position_after, ord=2),
+            'x_velocity': x_velocity,
+            'y_velocity': y_velocity,
+            'forward_reward': forward_reward,
+        }
+
+        velocity = np.sqrt(x_velocity**2 + y_velocity**2)
+        cost = velocity > self._velocity_threshold
+
+        if self.viewer:
+            clear_viewer(self.viewer)
+            add_velocity_marker(
+                viewer=self.viewer,
+                pos=self.get_body_com('torso')[:3].copy(),
+                vel=velocity,
+                cost=cost,
+                velocity_threshold=self._velocity_threshold,
+            )
+        if self.render_mode == 'human':
+            self.render()
+        return observation, reward, cost, terminated, False, info
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_swimmer_velocity.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_swimmer_velocity.py
new file mode 100644
index 000000000..2f11b68bd
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_swimmer_velocity.py
@@ -0,0 +1,69 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Swimmer environment with a safety constraint on velocity."""
+
+import numpy as np
+from gymnasium.envs.mujoco.swimmer_v4 import SwimmerEnv
+from safety_gymnasium.utils.task_utils import add_velocity_marker, clear_viewer
+
+
+class SafetySwimmerVelocityEnv(SwimmerEnv):
+    """Swimmer environment with a safety constraint on velocity."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._velocity_threshold = 0.04845
+        self.model.light(0).castshadow = False
+
+    def step(self, action):
+        xy_position_before = self.data.qpos[0:2].copy()
+        self.do_simulation(action, self.frame_skip)
+        xy_position_after = self.data.qpos[0:2].copy()
+
+        xy_velocity = (xy_position_after - xy_position_before) / self.dt
+        x_velocity, y_velocity = xy_velocity
+
+        forward_reward = self._forward_reward_weight * x_velocity
+
+        ctrl_cost = self.control_cost(action)
+
+        observation = self._get_obs()
+        reward = forward_reward - ctrl_cost
+        info = {
+            'reward_fwd': forward_reward,
+            'reward_ctrl': -ctrl_cost,
+            'x_position': xy_position_after[0],
+            'y_position': xy_position_after[1],
+            'distance_from_origin': np.linalg.norm(xy_position_after, ord=2),
+            'x_velocity': x_velocity,
+            'y_velocity': y_velocity,
+            'forward_reward': forward_reward,
+        }
+
+        velocity = np.sqrt(x_velocity**2 + y_velocity**2)
+        cost = velocity > self._velocity_threshold
+
+        if self.viewer:
+            clear_viewer(self.viewer)
+            add_velocity_marker(
+                viewer=self.viewer,
+                pos=self.get_body_com('torso')[:3].copy(),
+                vel=velocity,
+                cost=cost,
+                velocity_threshold=self._velocity_threshold,
+            )
+        if self.render_mode == 'human':
+            self.render()
+        return observation, reward, cost, False, False, info
diff --git a/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_walker2d_velocity.py b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_walker2d_velocity.py
new file mode 100644
index 000000000..ff7c14f51
--- /dev/null
+++ b/envs/safety-gymnasium/safety_gymnasium/tasks/safety_velocity/safety_walker2d_velocity.py
@@ -0,0 +1,64 @@
+# Copyright 2022 OmniSafe Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Walker2d environment with a safety constraint on velocity."""
+
+from gymnasium.envs.mujoco.walker2d_v4 import Walker2dEnv
+from safety_gymnasium.utils.task_utils import add_velocity_marker, clear_viewer
+
+
+class SafetyWalker2dVelocityEnv(Walker2dEnv):
+    """Walker2d environment with a safety constraint on velocity."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._velocity_threshold = 1.7075
+        self.model.light(0).castshadow = False
+
+    def step(self, action):
+        x_position_before = self.data.qpos[0]
+        self.do_simulation(action, self.frame_skip)
+        x_position_after = self.data.qpos[0]
+        x_velocity = (x_position_after - x_position_before) / self.dt
+
+        ctrl_cost = self.control_cost(action)
+
+        forward_reward = self._forward_reward_weight * x_velocity
+        healthy_reward = self.healthy_reward
+
+        rewards = forward_reward + healthy_reward
+        costs = ctrl_cost
+
+        observation = self._get_obs()
+        reward = rewards - costs
+        terminated = self.terminated
+        info = {
+            'x_position': x_position_after,
+            'x_velocity': x_velocity,
+        }
+
+        cost = x_velocity > self._velocity_threshold
+
+        if self.viewer:
+            clear_viewer(self.viewer)
+            add_velocity_marker(
+                viewer=self.viewer,
+                pos=self.get_body_com('torso')[:3].copy(),
+                vel=x_velocity,
+                cost=cost,
+                velocity_threshold=self._velocity_threshold,
+            )
+        if self.render_mode == 'human':
+            self.render()
+        return observation, reward, cost, terminated, False, info
diff --git a/envs/safety-gymnasium/safety_gymnasium/utils/passive_env_checker.py b/envs/safety-gymnasium/safety_gymnasium/utils/passive_env_checker.py
index 9a1ae9dd7..80c6db381 100644
--- a/envs/safety-gymnasium/safety_gymnasium/utils/passive_env_checker.py
+++ b/envs/safety-gymnasium/safety_gymnasium/utils/passive_env_checker.py
@@ -13,222 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 """A set of functions for passively checking environment implementations."""
-# pylint: disable=all
-
-import inspect
-from functools import partial
-from typing import Callable
 
 import numpy as np
-from gymnasium import Space, error, logger, spaces
-
-
-def _check_box_observation_space(observation_space: spaces.Box):
-    """Checks that a :class:`Box` observation space is defined in a sensible way.
-
-    Args:
-        observation_space: A box observation space
-    """
-    # Check if the box is an image
-    if (len(observation_space.shape) == 3 and observation_space.shape[0] != 1) or (
-        len(observation_space.shape) == 4 and observation_space.shape[0] == 1
-    ):
-        if observation_space.dtype != np.uint8:
-            logger.warn(
-                f'It seems a Box observation space is an image but the `dtype` is not `np.uint8`, actual type: {observation_space.dtype}. '
-                'If the Box observation space is not an image, we recommend flattening the observation to have only a 1D vector.'
-            )
-        if np.any(observation_space.low != 0) or np.any(observation_space.high != 255):
-            logger.warn(
-                'It seems a Box observation space is an image but the lower and upper bounds are not [0, 255]. '
-                f'Actual lower bound: {np.min(observation_space.low)}, upper bound: {np.max(observation_space.high)}. '
-                'Generally, CNN policies assume observations are within that range, so you may encounter an issue if the observation values are not.'
-            )
-
-    if len(observation_space.shape) not in [1, 3]:
-        if not (len(observation_space.shape) == 2 and observation_space.shape[0] == 1):
-            logger.warn(
-                'A Box observation space has an unconventional shape (neither an image, nor a 1D vector). '
-                'We recommend flattening the observation to have only a 1D vector or use a custom policy to properly process the data. '
-                f'Actual observation shape: {observation_space.shape}'
-            )
-
-    assert (
-        observation_space.low.shape == observation_space.shape
-    ), f'The Box observation space shape and low shape have different shapes, low shape: {observation_space.low.shape}, box shape: {observation_space.shape}'
-    assert (
-        observation_space.high.shape == observation_space.shape
-    ), f'The Box observation space shape and high shape have have different shapes, high shape: {observation_space.high.shape}, box shape: {observation_space.shape}'
-
-    if np.any(observation_space.low == observation_space.high):
-        logger.warn(
-            'A Box observation space maximum and minimum values are equal. '
-            f'Actual equal coordinates: {[x for x in zip(*np.where(observation_space.low == observation_space.high))]}'
-        )
-    elif np.any(observation_space.high < observation_space.low):
-        logger.warn(
-            'A Box observation space low value is greater than a high value. '
-            f'Actual less than coordinates: {[x for x in zip(*np.where(observation_space.high < observation_space.low))]}'
-        )
-
-
-def _check_box_action_space(action_space: spaces.Box):
-    """Checks that a :class:`Box` action space is defined in a sensible way.
-
-    Args:
-        action_space: A box action space
-    """
-    assert (
-        action_space.low.shape == action_space.shape
-    ), f'The Box action space shape and low shape have have different shapes, low shape: {action_space.low.shape}, box shape: {action_space.shape}'
-    assert (
-        action_space.high.shape == action_space.shape
-    ), f'The Box action space shape and high shape have different shapes, high shape: {action_space.high.shape}, box shape: {action_space.shape}'
-
-    if np.any(action_space.low == action_space.high):
-        logger.warn(
-            'A Box action space maximum and minimum values are equal. '
-            f'Actual equal coordinates: {[x for x in zip(*np.where(action_space.low == action_space.high))]}'
-        )
-    elif np.any(action_space.high < action_space.low):
-        logger.warn(
-            'A Box action space low value is greater than a high value. '
-            f'Actual less than coordinates: {[x for x in zip(*np.where(action_space.high < action_space.low))]}'
-        )
-
-
-def check_space(space: Space, space_type: str, check_box_space_fn: Callable[[spaces.Box], None]):
-    """A passive check of the environment action space that should not affect the environment."""
-    if not isinstance(space, spaces.Space):
-        raise AssertionError(
-            f'{space_type} space does not inherit from `gymnasium.spaces.Space`, actual type: {type(space)}'
-        )
-
-    elif isinstance(space, spaces.Box):
-        check_box_space_fn(space)
-    elif isinstance(space, spaces.Discrete):
-        assert (
-            0 < space.n
-        ), f"Discrete {space_type} space's number of elements must be positive, actual number of elements: {space.n}"
-        assert (
-            space.shape == ()
-        ), f"Discrete {space_type} space's shape should be empty, actual shape: {space.shape}"
-    elif isinstance(space, spaces.MultiDiscrete):
-        assert (
-            space.shape == space.nvec.shape
-        ), f"Multi-discrete {space_type} space's shape must be equal to the nvec shape, space shape: {space.shape}, nvec shape: {space.nvec.shape}"
-        assert np.all(
-            0 < space.nvec
-        ), f"Multi-discrete {space_type} space's all nvec elements must be greater than 0, actual nvec: {space.nvec}"
-    elif isinstance(space, spaces.MultiBinary):
-        assert np.all(
-            0 < np.asarray(space.shape)
-        ), f"Multi-binary {space_type} space's all shape elements must be greater than 0, actual shape: {space.shape}"
-    elif isinstance(space, spaces.Tuple):
-        assert 0 < len(space.spaces), f'An empty Tuple {space_type} space is not allowed.'
-        for subspace in space.spaces:
-            check_space(subspace, space_type, check_box_space_fn)
-    elif isinstance(space, spaces.Dict):
-        assert 0 < len(space.spaces.keys()), f'An empty Dict {space_type} space is not allowed.'
-        for subspace in space.values():
-            check_space(subspace, space_type, check_box_space_fn)
-
-
-check_observation_space = partial(
-    check_space,
-    space_type='observation',
-    check_box_space_fn=_check_box_observation_space,
-)
-check_action_space = partial(
-    check_space, space_type='action', check_box_space_fn=_check_box_action_space
-)
-
-
-def check_obs(obs, observation_space: spaces.Space, method_name: str):
-    """Check that the observation returned by the environment correspond to the declared one.
-
-    Args:
-        obs: The observation to check
-        observation_space: The observation space of the observation
-        method_name: The method name that generated the observation
-    """
-    pre = f'The obs returned by the `{method_name}()` method'
-    if isinstance(observation_space, spaces.Discrete):
-        if not isinstance(obs, (np.int64, int)):
-            logger.warn(f'{pre} should be an int or np.int64, actual type: {type(obs)}')
-    elif isinstance(observation_space, spaces.Box):
-        if observation_space.shape != ():
-            if not isinstance(obs, np.ndarray):
-                logger.warn(f'{pre} was expecting a numpy array, actual type: {type(obs)}')
-            elif obs.dtype != observation_space.dtype:
-                logger.warn(
-                    f'{pre} was expecting numpy array dtype to be {observation_space.dtype}, actual type: {obs.dtype}'
-                )
-    elif isinstance(observation_space, (spaces.MultiBinary, spaces.MultiDiscrete)):
-        if not isinstance(obs, np.ndarray):
-            logger.warn(f'{pre} was expecting a numpy array, actual type: {type(obs)}')
-    elif isinstance(observation_space, spaces.Tuple):
-        if not isinstance(obs, tuple):
-            logger.warn(f'{pre} was expecting a tuple, actual type: {type(obs)}')
-        assert len(obs) == len(
-            observation_space.spaces
-        ), f'{pre} length is not same as the observation space length, obs length: {len(obs)}, space length: {len(observation_space.spaces)}'
-        for sub_obs, sub_space in zip(obs, observation_space.spaces):
-            check_obs(sub_obs, sub_space, method_name)
-    elif isinstance(observation_space, spaces.Dict):
-        assert isinstance(obs, dict), f'{pre} must be a dict, actual type: {type(obs)}'
-        assert (
-            obs.keys() == observation_space.spaces.keys()
-        ), f'{pre} observation keys is not same as the observation space keys, obs keys: {list(obs.keys())}, space keys: {list(observation_space.spaces.keys())}'
-        for space_key in observation_space.spaces.keys():
-            check_obs(obs[space_key], observation_space[space_key], method_name)
-
-    try:
-        if obs not in observation_space:
-            logger.warn(f'{pre} is not within the observation space.')
-    except Exception as e:
-        logger.warn(f'{pre} is not within the observation space with exception: {e}')
-
-
-def env_reset_passive_checker(env, **kwargs):
-    """A passive check of the `Env.reset` function investigating the returning reset information and returning the data unchanged."""
-    signature = inspect.signature(env.reset)
-    if 'seed' not in signature.parameters and 'kwargs' not in signature.parameters:
-        logger.warn(
-            'Future gymnasium versions will require that `Env.reset` can be passed a `seed` instead of using `Env.seed` for resetting the environment random number generator.'
-        )
-    else:
-        seed_param = signature.parameters.get('seed')
-        # Check the default value is None
-        if seed_param is not None and seed_param.default is not None:
-            logger.warn(
-                'The default seed argument in `Env.reset` should be `None`, otherwise the environment will by default always be deterministic. '
-                f'Actual default: {seed_param}'
-            )
-
-    if 'options' not in signature.parameters and 'kwargs' not in signature.parameters:
-        logger.warn(
-            'Future gymnasium versions will require that `Env.reset` can be passed `options` to allow the environment initialisation to be passed additional information.'
-        )
-
-    # Checks the result of env.reset with kwargs
-    result = env.reset(**kwargs)
-
-    if not isinstance(result, tuple):
-        logger.warn(
-            f'The result returned by `env.reset()` was not a tuple of the form `(obs, info)`, where `obs` is a observation and `info` is a dictionary containing additional information. Actual type: `{type(result)}`'
-        )
-    elif len(result) != 2:
-        logger.warn(
-            'The result returned by `env.reset()` should be `(obs, info)` by default, , where `obs` is a observation and `info` is a dictionary containing additional information.'
-        )
-    else:
-        obs, info = result
-        check_obs(obs, env.observation_space, 'reset')
-        assert isinstance(
-            info, dict
-        ), f'The second element returned by `env.reset()` was not a dictionary, actual type: {type(info)}'
-    return result
+from gymnasium import error, logger
+from gymnasium.utils.passive_env_checker import check_obs
 
 
 def env_step_passive_checker(env, action):
@@ -261,14 +49,26 @@ def env_step_passive_checker(env, action):
             )
     else:
         raise error.Error(
-            f'Expected `Env.step` to return a four or five element tuple, actual number of elements returned: {len(result)}.'
+            f'Expected `Env.step` to return a four or five element tuple,\
+            actual number of elements returned: {len(result)}.'
         )
 
     check_obs(obs, env.observation_space, 'step')
+    check_reward_cost(reward=reward, cost=cost)
+
+    assert isinstance(
+        info, dict
+    ), f'The `info` returned by `step()` must be a python dictionary, actual type: {type(info)}'
+
+    return result
 
+
+def check_reward_cost(reward, cost):
+    """Check out the type and the value of the reward and cost."""
     if not (np.issubdtype(type(reward), np.integer) or np.issubdtype(type(reward), np.floating)):
         logger.warn(
-            f'The reward returned by `step()` must be a float, int, np.integer or np.floating, actual type: {type(reward)}'
+            f'The reward returned by `step()` must be a float,\
+            int, np.integer or np.floating, actual type: {type(reward)}'
         )
     else:
         if np.isnan(reward):
@@ -276,111 +76,13 @@ def env_step_passive_checker(env, action):
         if np.isinf(reward):
             logger.warn('The reward is an inf value.')
 
-    assert isinstance(
-        info, dict
-    ), f'The `info` returned by `step()` must be a python dictionary, actual type: {type(info)}'
-
-    return result
-
-
-def _check_render_return(render_mode, render_return):
-    """Produces warning if `render_return` doesn't match `render_mode`."""
-    if render_mode == 'human':
-        if render_return is not None:
-            logger.warn(f'Human rendering should return `None`, got {type(render_return)}')
-    elif render_mode == 'rgb_array':
-        if not isinstance(render_return, np.ndarray):
-            logger.warn(
-                f'RGB-array rendering should return a numpy array, got {type(render_return)}'
-            )
-        else:
-            if render_return.dtype != np.uint8:
-                logger.warn(
-                    f'RGB-array rendering should return a numpy array with dtype uint8, got {render_return.dtype}'
-                )
-            if render_return.ndim != 3:
-                logger.warn(
-                    f'RGB-array rendering should return a numpy array with three axes, got {render_return.ndim}'
-                )
-            if render_return.ndim == 3 and render_return.shape[2] != 3:
-                logger.warn(
-                    f'RGB-array rendering should return a numpy array in which the last axis has three dimensions, got {render_return.shape[2]}'
-                )
-    elif render_mode == 'depth_array':
-        if not isinstance(render_return, np.ndarray):
-            logger.warn(
-                f'Depth-array rendering should return a numpy array, got {type(render_return)}'
-            )
-        elif render_return.ndim != 2:
-            logger.warn(
-                f'Depth-array rendering should return a numpy array with two axes, got {render_return.ndim}'
-            )
-    elif render_mode in ['ansi', 'ascii']:
-        if not isinstance(render_return, str):
-            logger.warn(f'ANSI/ASCII rendering should produce a string, got {type(render_return)}')
-    elif render_mode.endswith('_list'):
-        if not isinstance(render_return, list):
-            logger.warn(
-                f'Render mode `{render_mode}` should produce a list, got {type(render_return)}'
-            )
-        else:
-            base_render_mode = render_mode[: -len('_list')]
-            for item in render_return:
-                _check_render_return(
-                    base_render_mode, item
-                )  # Check that each item of the list matches the base render mode
-
-
-def env_render_passive_checker(env):
-    """A passive check of the `Env.render` that the declared render modes/fps in the metadata of the environment is declared."""
-    render_modes = env.metadata.get('render_modes')
-    if render_modes is None:
+    if not (np.issubdtype(type(cost), np.integer) or np.issubdtype(type(cost), np.floating)):
         logger.warn(
-            "No render modes was declared in the environment (env.metadata['render_modes'] is None or not defined), you may have trouble when calling `.render()`."
+            f'The reward returned by `step()` must be a float,\
+            int, np.integer or np.floating, actual type: {type(cost)}'
         )
     else:
-        if not isinstance(render_modes, (list, tuple)):
-            logger.warn(
-                f'Expects the render_modes to be a sequence (i.e. list, tuple), actual type: {type(render_modes)}'
-            )
-        elif not all(isinstance(mode, str) for mode in render_modes):
-            logger.warn(
-                f'Expects all render modes to be strings, actual types: {[type(mode) for mode in render_modes]}'
-            )
-
-        render_fps = env.metadata.get('render_fps')
-        # We only require `render_fps` if rendering is actually implemented
-        if len(render_modes) > 0:
-            if render_fps is None:
-                logger.warn(
-                    "No render fps was declared in the environment (env.metadata['render_fps'] is None or not defined), rendering may occur at inconsistent fps."
-                )
-            else:
-                if not (
-                    np.issubdtype(type(render_fps), np.integer)
-                    or np.issubdtype(type(render_fps), np.floating)
-                ):
-                    logger.warn(
-                        f"Expects the `env.metadata['render_fps']` to be an integer or a float, actual type: {type(render_fps)}"
-                    )
-                else:
-                    assert (
-                        render_fps > 0
-                    ), f"Expects the `env.metadata['render_fps']` to be greater than zero, actual value: {render_fps}"
-
-        # env.render is now an attribute with default None
-        if len(render_modes) == 0:
-            assert (
-                env.render_mode is None
-            ), f'With no render_modes, expects the Env.render_mode to be None, actual value: {env.render_mode}'
-        else:
-            assert env.render_mode is None or env.render_mode in render_modes, (
-                'The environment was initialized successfully however with an unsupported render mode. '
-                f'Render mode: {env.render_mode}, modes: {render_modes}'
-            )
-
-    result = env.render()
-    if env.render_mode is not None:
-        _check_render_return(env.render_mode, result)
-
-    return result
+        if np.isnan(cost):
+            logger.warn('The reward is a NaN value.')
+        if np.isinf(cost):
+            logger.warn('The reward is an inf value.')
diff --git a/envs/safety-gymnasium/safety_gymnasium/utils/registration.py b/envs/safety-gymnasium/safety_gymnasium/utils/registration.py
index 8a1b7b1bb..ff7fa7e7e 100644
--- a/envs/safety-gymnasium/safety_gymnasium/utils/registration.py
+++ b/envs/safety-gymnasium/safety_gymnasium/utils/registration.py
@@ -12,436 +12,42 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# pylint: disable=all
+"""Register and make environments."""
 
-import contextlib
 import copy
-import difflib
 import importlib
 import importlib.util
-import re
-import sys
-import warnings
-from dataclasses import dataclass, field
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union, overload
-from typing_extensions import Literal  # Python 3.8+
-
-import numpy as np
-from gymnasium.wrappers import AutoResetWrapper, HumanRendering, RenderCollection
-from gymnasium.wrappers.compatibility import EnvCompatibility
-from safety_gymnasium.wrappers import OrderEnforcing, TimeLimit
-from safety_gymnasium.wrappers.env_checker import PassiveEnvChecker
-
-
-if sys.version_info < (3, 10):
-    import importlib_metadata as metadata  # type: ignore
-else:
-    import importlib.metadata as metadata
+from typing import Optional, Sequence, Union
 
 from gymnasium import Env, error, logger
-
-
-ENV_ID_RE = re.compile(
-    r'^(?:(?P<namespace>[\w:-]+)\/)?(?:(?P<name>[\w:.-]+?))(?:-v(?P<version>\d+))?$'
+from gymnasium.envs.registration import (  # pylint: disable=unused-import
+    EnvSpec,
+    _check_version_exists,
+    find_highest_version,
+    get_env_id,
+    load,
+    namespace,
+    parse_env_id,
 )
+from gymnasium.envs.registration import register as gymnasium_register
+from gymnasium.envs.registration import registry, spec  # pylint: disable=unused-import
+from gymnasium.wrappers import AutoResetWrapper, HumanRendering, OrderEnforcing, RenderCollection
+from gymnasium.wrappers.compatibility import EnvCompatibility
+from safety_gymnasium.wrappers import SafePassiveEnvChecker, SafeTimeLimit
 
 
-def load(name: str) -> callable:
-    """Loads an environment with name and returns an environment creation function
-
-    Args:
-        name: The environment name
-
-    Returns:
-        Calls the environment constructor
-    """
-    mod_name, attr_name = name.split(':')
-    mod = importlib.import_module(mod_name)
-    fn = getattr(mod, attr_name)
-    return fn
-
-
-def parse_env_id(id: str) -> Tuple[Optional[str], str, Optional[int]]:
-    """Parse environment ID string format.
-
-    This format is true today, but it's *not* an official spec.
-    [namespace/](env-name)-v(version)    env-name is group 1, version is group 2
-
-    2016-10-31: We're experimentally expanding the environment ID format
-    to include an optional namespace.
-
-    Args:
-        id: The environment id to parse
-
-    Returns:
-        A tuple of environment namespace, environment name and version number
-
-    Raises:
-        Error: If the environment id does not a valid environment regex
-    """
-    match = ENV_ID_RE.fullmatch(id)
-    if not match:
-        raise error.Error(
-            f'Malformed environment ID: {id}.'
-            f'(Currently all IDs must be of the form [namespace/](env-name)-v(version). (namespace is optional))'
-        )
-    namespace, name, version = match.group('namespace', 'name', 'version')
-    if version is not None:
-        version = int(version)
-
-    return namespace, name, version
-
-
-def get_env_id(ns: Optional[str], name: str, version: Optional[int]) -> str:
-    """Get the full env ID given a name and (optional) version and namespace. Inverse of :meth:`parse_env_id`.
-
-    Args:
-        ns: The environment namespace
-        name: The environment name
-        version: The environment version
-
-    Returns:
-        The environment id
-    """
-
-    full_name = name
-    if version is not None:
-        full_name += f'-v{version}'
-    if ns is not None:
-        full_name = ns + '/' + full_name
-    return full_name
-
-
-@dataclass
-class EnvSpec:
-    """A specification for creating environments with `gymnasium.make`.
-
-    * id: The string used to create the environment with `gymnasium.make`
-    * entry_point: The location of the environment to create from
-    * reward_threshold: The reward threshold for completing the environment.
-    * nondeterministic: If the observation of an environment cannot be repeated with the same initial state, random number generator state and actions.
-    * max_episode_steps: The max number of steps that the environment can take before truncation
-    * order_enforce: If to enforce the order of `reset` before `step` and `render` functions
-    * autoreset: If to automatically reset the environment on episode end
-    * disable_env_checker: If to disable the environment checker wrapper in `gymnasium.make`, by default False (runs the environment checker)
-    * kwargs: Additional keyword arguments passed to the environments through `gymnasium.make`
-    """
-
-    id: str
-    entry_point: Union[Callable, str]
-
-    # Environment attributes
-    reward_threshold: Optional[float] = field(default=None)
-    nondeterministic: bool = field(default=False)
-
-    # Wrappers
-    max_episode_steps: Optional[int] = field(default=None)
-    order_enforce: bool = field(default=True)
-    autoreset: bool = field(default=False)
-    disable_env_checker: bool = field(default=False)
-    apply_api_compatibility: bool = field(default=False)
-
-    # Environment arguments
-    kwargs: dict = field(default_factory=dict)
-
-    # post-init attributes
-    namespace: Optional[str] = field(init=False)
-    name: str = field(init=False)
-    version: Optional[int] = field(init=False)
-
-    def __post_init__(self):
-        # Initialize namespace, name, version
-        self.namespace, self.name, self.version = parse_env_id(self.id)
-
-    def make(self, **kwargs) -> Env:
-        # For compatibility purposes
-        return make(self, **kwargs)
-
-
-def _check_namespace_exists(ns: Optional[str]):
-    """Check if a namespace exists. If it doesn't, print a helpful error message."""
-    if ns is None:
-        return
-    namespaces = {spec_.namespace for spec_ in registry.values() if spec_.namespace is not None}
-    if ns in namespaces:
-        return
-
-    suggestion = difflib.get_close_matches(ns, namespaces, n=1) if len(namespaces) > 0 else None
-    suggestion_msg = (
-        f'Did you mean: `{suggestion[0]}`?'
-        if suggestion
-        else f'Have you installed the proper package for {ns}?'
-    )
-
-    raise error.NamespaceNotFound(f'Namespace {ns} not found. {suggestion_msg}')
-
-
-def _check_name_exists(ns: Optional[str], name: str):
-    """Check if an env exists in a namespace. If it doesn't, print a helpful error message."""
-    _check_namespace_exists(ns)
-    names = {spec_.name for spec_ in registry.values() if spec_.namespace == ns}
-
-    if name in names:
-        return
-
-    suggestion = difflib.get_close_matches(name, names, n=1)
-    namespace_msg = f' in namespace {ns}' if ns else ''
-    suggestion_msg = f'Did you mean: `{suggestion[0]}`?' if suggestion else ''
-
-    raise error.NameNotFound(f"Environment {name} doesn't exist{namespace_msg}. {suggestion_msg}")
-
-
-def _check_version_exists(ns: Optional[str], name: str, version: Optional[int]):
-    """Check if an env version exists in a namespace. If it doesn't, print a helpful error message.
-    This is a complete test whether an environment identifier is valid, and will provide the best available hints.
-
-    Args:
-        ns: The environment namespace
-        name: The environment space
-        version: The environment version
-
-    Raises:
-        DeprecatedEnv: The environment doesn't exist but a default version does
-        VersionNotFound: The ``version`` used doesn't exist
-        DeprecatedEnv: Environment version is deprecated
-    """
-    if get_env_id(ns, name, version) in registry:
-        return
-
-    _check_name_exists(ns, name)
-    if version is None:
-        return
-
-    message = f"Environment version `v{version}` for environment `{get_env_id(ns, name, None)}` doesn't exist."
-
-    env_specs = [
-        spec_ for spec_ in registry.values() if spec_.namespace == ns and spec_.name == name
-    ]
-    env_specs = sorted(env_specs, key=lambda spec_: int(spec_.version or -1))
-
-    default_spec = [spec_ for spec_ in env_specs if spec_.version is None]
-
-    if default_spec:
-        message += f' It provides the default version {default_spec[0].id}`.'
-        if len(env_specs) == 1:
-            raise error.DeprecatedEnv(message)
-
-    # Process possible versioned environments
-
-    versioned_specs = [spec_ for spec_ in env_specs if spec_.version is not None]
-
-    latest_spec = max(versioned_specs, key=lambda spec: spec.version, default=None)  # type: ignore
-    if latest_spec is not None and version > latest_spec.version:
-        version_list_msg = ', '.join(f'`v{spec_.version}`' for spec_ in env_specs)
-        message += f' It provides versioned environments: [ {version_list_msg} ].'
-
-        raise error.VersionNotFound(message)
-
-    if latest_spec is not None and version < latest_spec.version:
-        raise error.DeprecatedEnv(
-            f'Environment version v{version} for `{get_env_id(ns, name, None)}` is deprecated. '
-            f'Please use `{latest_spec.id}` instead.'
-        )
-
-
-def find_highest_version(ns: Optional[str], name: str) -> Optional[int]:
-    version: List[int] = [
-        spec_.version
-        for spec_ in registry.values()
-        if spec_.namespace == ns and spec_.name == name and spec_.version is not None
-    ]
-    return max(version, default=None)
-
-
-def load_env_plugins(entry_point: str = 'gymnasium.envs') -> None:
-    # Load third-party environments
-    for plugin in metadata.entry_points(group=entry_point):
-        # Python 3.8 doesn't support plugin.module, plugin.attr
-        # So we'll have to try and parse this ourselves
-        module, attr = None, None
-        try:
-            module, attr = plugin.module, plugin.attr  # type: ignore  ## error: Cannot access member "attr" for type "EntryPoint"
-        except AttributeError:
-            if ':' in plugin.value:
-                module, attr = plugin.value.split(':', maxsplit=1)
-            else:
-                module, attr = plugin.value, None
-        except Exception as e:
-            warnings.warn(
-                f'While trying to load plugin `{plugin}` from {entry_point}, an exception occurred: {e}'
-            )
-            module, attr = None, None
-        finally:
-            if attr is None:
-                raise error.Error(
-                    f'Gymnasium environment plugin `{module}` must specify a function to execute, not a root module'
-                )
-
-        context = namespace(plugin.name)
-        if plugin.name.startswith('__') and plugin.name.endswith('__'):
-            # `__internal__` is an artifact of the plugin system when
-            # the root namespace had an allow-list. The allow-list is now
-            # removed and plugins can register environments in the root
-            # namespace with the `__root__` magic key.
-            if plugin.name == '__root__' or plugin.name == '__internal__':
-                context = contextlib.nullcontext()
-            else:
-                logger.warn(
-                    f'The environment namespace magic key `{plugin.name}` is unsupported. '
-                    'To register an environment at the root namespace you should specify the `__root__` namespace.'
-                )
-
-        with context:
-            fn = plugin.load()
-            try:
-                fn()
-            except Exception as e:
-                logger.warn(str(e))
-
-
-# fmt: off
-@overload
-def make(id: str, **kwargs) -> Env: ...
-@overload
-def make(id: EnvSpec, **kwargs) -> Env: ...
-
-# Mujoco
-# ----------------------------------------
-@overload
-def make(id: Literal[
-    'SafetyAntVelocity-v0',
-    'Reacher-v2', 'Reacher-v4'
-], **kwargs) -> Env[np.ndarray, np.ndarray]: ...
-# fmt: on
-
-
-# Global registry of environments. Meant to be accessed through `register` and `make`
-registry: Dict[str, EnvSpec] = {}
-current_namespace: Optional[str] = None
-
-
-def _check_spec_register(spec: EnvSpec):
-    """Checks whether the spec is valid to be registered. Helper function for `register`."""
-    global registry
-    latest_versioned_spec = max(
-        (
-            spec_
-            for spec_ in registry.values()
-            if spec_.namespace == spec.namespace
-            and spec_.name == spec.name
-            and spec_.version is not None
-        ),
-        key=lambda spec_: int(spec_.version),  # type: ignore
-        default=None,
-    )
-
-    unversioned_spec = next(
-        (
-            spec_
-            for spec_ in registry.values()
-            if spec_.namespace == spec.namespace
-            and spec_.name == spec.name
-            and spec_.version is None
-        ),
-        None,
-    )
-
-    if unversioned_spec is not None and spec.version is not None:
-        raise error.RegistrationError(
-            "Can't register the versioned environment "
-            f'`{spec.id}` when the unversioned environment '
-            f'`{unversioned_spec.id}` of the same name already exists.'
-        )
-    elif latest_versioned_spec is not None and spec.version is None:
-        raise error.RegistrationError(
-            "Can't register the unversioned environment "
-            f'`{spec.id}` when the versioned environment '
-            f'`{latest_versioned_spec.id}` of the same name '
-            f'already exists. Note: the default behavior is '
-            f'that `gymnasium.make` with the unversioned environment '
-            f'will return the latest versioned environment'
-        )
-
-
-# Public API
-
-
-@contextlib.contextmanager
-def namespace(ns: str):
-    global current_namespace
-    old_namespace = current_namespace
-    current_namespace = ns
-    yield
-    current_namespace = old_namespace
-
-
-def register(
-    id: str,
-    entry_point: Union[Callable, str],
-    reward_threshold: Optional[float] = None,
-    nondeterministic: bool = False,
-    max_episode_steps: Optional[int] = None,
-    order_enforce: bool = True,
-    autoreset: bool = False,
-    disable_env_checker: bool = False,
-    apply_api_compatibility: bool = False,
-    **kwargs,
-):
-    """Register an environment with gymnasium.
-
-    The `id` parameter corresponds to the name of the environment, with the syntax as follows:
-    `(namespace)/(env_name)-v(version)` where `namespace` is optional.
+safe_registry = set()
 
-    It takes arbitrary keyword arguments, which are passed to the `EnvSpec` constructor.
 
-    Args:
-        id: The environment id
-        entry_point: The entry point for creating the environment
-        reward_threshold: The reward threshold considered to have learnt an environment
-        nondeterministic: If the environment is nondeterministic (even with knowledge of the initial seed and all actions)
-        max_episode_steps: The maximum number of episodes steps before truncation. Used by the Time Limit wrapper.
-        order_enforce: If to enable the order enforcer wrapper to ensure users run functions in the correct order
-        autoreset: If to add the autoreset wrapper such that reset does not need to be called.
-        disable_env_checker: If to disable the environment checker for the environment. Recommended to False.
-        apply_api_compatibility: If to apply the `StepAPICompatibility` wrapper.
-        **kwargs: arbitrary keyword arguments which are passed to the environment constructor
-    """
-    global registry, current_namespace
-    ns, name, version = parse_env_id(id)
-
-    if current_namespace is not None:
-        if kwargs.get('namespace') is not None and kwargs.get('namespace') != current_namespace:
-            logger.warn(
-                f"Custom namespace `{kwargs.get('namespace')}` is being overridden by namespace `{current_namespace}`. "
-                f"If you are developing a plugin you shouldn't specify a namespace in `register` calls. "
-                'The namespace is specified through the entry point package metadata.'
-            )
-        ns_id = current_namespace
-    else:
-        ns_id = ns
-
-    full_id = get_env_id(ns_id, name, version)
-
-    new_spec = EnvSpec(
-        id=full_id,
-        entry_point=entry_point,
-        reward_threshold=reward_threshold,
-        nondeterministic=nondeterministic,
-        max_episode_steps=max_episode_steps,
-        order_enforce=order_enforce,
-        autoreset=autoreset,
-        disable_env_checker=disable_env_checker,
-        apply_api_compatibility=apply_api_compatibility,
-        **kwargs,
-    )
-    _check_spec_register(new_spec)
-    if new_spec.id in registry:
-        logger.warn(f'Overriding environment {new_spec.id} already in registry.')
-    registry[new_spec.id] = new_spec
+def register(**kwargs):
+    """Register an environment."""
+    safe_registry.add(kwargs['id'])
+    gymnasium_register(**kwargs)
 
 
+# pylint: disable-next=too-many-arguments,too-many-branches,too-many-statements,too-many-locals
 def make(
-    id: Union[str, EnvSpec],
+    id: Union[str, EnvSpec],  # pylint: disable=redefined-builtin
     max_episode_steps: Optional[int] = None,
     autoreset: bool = False,
     apply_api_compatibility: Optional[bool] = None,
@@ -461,7 +67,8 @@ def make(
             By default, the argument is None to which the environment specification `apply_api_compatibility` is used
             which defaults to False. Otherwise, the value of `apply_api_compatibility` is used.
             If `True`, the wrapper is applied otherwise, the wrapper is not applied.
-        disable_env_checker: If to run the env checker, None will default to the environment specification `disable_env_checker`
+        disable_env_checker: If to run the env checker,
+            None will default to the environment specification `disable_env_checker`
             (which is by default False, running the environment checker),
             otherwise will run according to this parameter (`True` = not run, `False` = run)
         kwargs: Additional arguments to pass to the environment constructor.
@@ -483,11 +90,11 @@ def make(
                 raise ModuleNotFoundError(
                     f'{e}. Environment registration via importing a module failed. '
                     f"Check whether '{module}' contains env registration and can be imported."
-                )
+                ) from e
         spec_ = registry.get(id)
 
-        ns, name, version = parse_env_id(id)
-        latest_version = find_highest_version(ns, name)
+        name_space, name, version = parse_env_id(id)
+        latest_version = find_highest_version(name_space, name)
         if version is not None and latest_version is not None and latest_version > version:
             logger.warn(
                 f'The environment {id} is out of date. You should consider '
@@ -495,7 +102,7 @@ def make(
             )
         if version is None and latest_version is not None:
             version = latest_version
-            new_env_id = get_env_id(ns, name, version)
+            new_env_id = get_env_id(name_space, name, version)
             spec_ = registry.get(new_env_id)
             logger.warn(
                 f'Using the latest versioned environment `{new_env_id}` '
@@ -503,15 +110,17 @@ def make(
             )
 
         if spec_ is None:
-            _check_version_exists(ns, name, version)
+            _check_version_exists(name_space, name, version)
             raise error.Error(f'No registered env with id: {id}')
 
+        assert id in safe_registry, f'Environment {id} is not registered in safety-gymnasium.'
+
     _kwargs = spec_.kwargs.copy()
     _kwargs.update(kwargs)
 
     if spec_.entry_point is None:
         raise error.Error(f'{spec_.id} registered but entry_point is not specified')
-    elif callable(spec_.entry_point):
+    if callable(spec_.entry_point):
         env_creator = spec_.entry_point
     else:
         # Assume it's a string
@@ -521,17 +130,20 @@ def make(
     apply_human_rendering = False
     apply_render_collection = False
 
-    # If we have access to metadata we check that "render_mode" is valid and see if the HumanRendering wrapper needs to be applied
+    # If we have access to metadata we check that "render_mode" is valid and
+    # see if the HumanRendering wrapper needs to be applied
     if mode is not None and hasattr(env_creator, 'metadata'):
         assert isinstance(
             env_creator.metadata, dict
-        ), f'Expect the environment creator ({env_creator}) metadata to be dict, actual type: {type(env_creator.metadata)}'
+        ), f'Expect the environment creator ({env_creator}) metadata to be dict,\
+                actual type: {type(env_creator.metadata)}'
 
         if 'render_modes' in env_creator.metadata:
             render_modes = env_creator.metadata['render_modes']
             if not isinstance(render_modes, Sequence):
                 logger.warn(
-                    f'Expects the environment metadata render_modes to be a Sequence (tuple or list), actual type: {type(render_modes)}'
+                    f'Expects the environment metadata render_modes to be a Sequence (tuple or list),\
+                        actual type: {type(render_modes)}'
                 )
 
             # Apply the `HumanRendering` wrapper, if the mode=="human" but "human" not in render_modes
@@ -558,11 +170,13 @@ def make(
                 apply_render_collection = True
             elif mode not in render_modes:
                 logger.warn(
-                    f'The environment is being initialised with mode ({mode}) that is not in the possible render_modes ({render_modes}).'
+                    f'The environment is being initialised with mode ({mode})\
+                        that is not in the possible render_modes ({render_modes}).'
                 )
         else:
             logger.warn(
-                f"The environment creator metadata doesn't include `render_modes`, contains: {list(env_creator.metadata.keys())}"
+                f"The environment creator metadata doesn't include `render_modes`,\
+                    contains: {list(env_creator.metadata.keys())}"
             )
 
     if apply_api_compatibility is True or (
@@ -585,8 +199,8 @@ def make(
                 'Gym tried to apply the HumanRendering wrapper but it looks like your environment is using the old '
                 'rendering API, which is not supported by the HumanRendering wrapper.'
             )
-        else:
-            raise e
+
+        raise e
 
     # Copies the environment creation specification and kwargs to add to the environment specification details
     spec_ = copy.deepcopy(spec_)
@@ -603,7 +217,7 @@ def make(
     if disable_env_checker is False or (
         disable_env_checker is None and spec_.disable_env_checker is False
     ):
-        env = PassiveEnvChecker(env)
+        env = SafePassiveEnvChecker(env)
 
     # Add the order enforcing wrapper
     if spec_.order_enforce:
@@ -611,9 +225,9 @@ def make(
 
     # Add the time limit wrapper
     if max_episode_steps is not None:
-        env = TimeLimit(env, max_episode_steps)
+        env = SafeTimeLimit(env, max_episode_steps)
     elif spec_.max_episode_steps is not None:
-        env = TimeLimit(env, spec_.max_episode_steps)
+        env = SafeTimeLimit(env, spec_.max_episode_steps)
 
     # Add the autoreset wrapper
     if autoreset:
@@ -626,15 +240,3 @@ def make(
         env = RenderCollection(env)
 
     return env
-
-
-def spec(env_id: str) -> EnvSpec:
-    """Retrieve the spec for the given environment from the global registry."""
-    spec_ = registry.get(env_id)
-    if spec_ is None:
-        ns, name, version = parse_env_id(env_id)
-        _check_version_exists(ns, name, version)
-        raise error.Error(f'No registered env with id: {env_id}')
-    else:
-        assert isinstance(spec_, EnvSpec)
-        return spec_
diff --git a/envs/safety-gymnasium/safety_gymnasium/utils/task_utils.py b/envs/safety-gymnasium/safety_gymnasium/utils/task_utils.py
index e7f73a581..c16a72b3f 100644
--- a/envs/safety-gymnasium/safety_gymnasium/utils/task_utils.py
+++ b/envs/safety-gymnasium/safety_gymnasium/utils/task_utils.py
@@ -56,3 +56,31 @@ def get_body_xvelp(model, data, name):
     jacp = get_body_jacp(model, data, name).reshape((3, model.nv))
     xvelp = np.dot(jacp, data.qvel)
     return xvelp
+
+
+def add_velocity_marker(viewer, pos, vel, cost, velocity_threshold):
+    """Add a marker to the viewer to indicate the velocity of the agent."""
+    pos = pos + np.array([0, 0, 0.6])
+    safe_color = np.array([0.2, 0.8, 0.2, 0.5])
+    unsafe_color = np.array([0.5, 0, 0, 0.5])
+
+    if cost:
+        color = unsafe_color
+    else:
+        vel_ratio = vel / velocity_threshold
+        color = safe_color * (1 - vel_ratio)
+
+    viewer.add_marker(
+        pos=pos,
+        size=0.2 * np.ones(3),
+        type=mujoco.mjtGeom.mjGEOM_SPHERE,  # pylint: disable=no-member
+        rgba=color,
+        label='',
+    )
+
+
+def clear_viewer(viewer):
+    """Clear the viewer's all markers and overlays."""
+    # pylint: disable=protected-access
+    viewer._markers[:] = []
+    viewer._overlays.clear()
diff --git a/envs/safety-gymnasium/safety_gymnasium/world.py b/envs/safety-gymnasium/safety_gymnasium/world.py
index bf8d2fbc9..db4052da4 100644
--- a/envs/safety-gymnasium/safety_gymnasium/world.py
+++ b/envs/safety-gymnasium/safety_gymnasium/world.py
@@ -39,7 +39,7 @@ class World:  # pylint: disable=too-many-instance-attributes
     # Default configuration (this should not be nested since it gets copied)
     # *NOTE:* Changes to this configuration should also be reflected in `Builder` configuration
     DEFAULT = {
-        'robot_base': 'xmls/car.xml',  # Which robot XML to use as the base
+        'robot_base': 'assets/xmls/car.xml',  # Which robot XML to use as the base
         'robot_xy': np.zeros(2),  # Robot XY location
         'robot_rot': 0,  # Robot rotation about Z axis
         'floor_size': [3.5, 3.5, 0.1],  # Used for displaying the floor
@@ -164,8 +164,8 @@ def build(self):  # pylint: disable=too-many-locals, too-many-branches, too-many
                 '@builtin': 'checker',
                 '@height': '100',
                 '@width': '100',
-                '@rgb1': '0.7 0.7 0.7',
-                '@rgb2': '0.8 0.8 0.8',
+                '@rgb1': '0.859 0.843 0.765',
+                '@rgb2': '0.949 0.914 0.855',
                 '@type': '2d',
             }
         )
diff --git a/envs/safety-gymnasium/safety_gymnasium/wrappers/__init__.py b/envs/safety-gymnasium/safety_gymnasium/wrappers/__init__.py
index d6e24d012..531a373f9 100644
--- a/envs/safety-gymnasium/safety_gymnasium/wrappers/__init__.py
+++ b/envs/safety-gymnasium/safety_gymnasium/wrappers/__init__.py
@@ -14,6 +14,5 @@
 # ==============================================================================
 """Env wrappers."""
 
-from safety_gymnasium.wrappers.env_checker import PassiveEnvChecker
-from safety_gymnasium.wrappers.order_enforcing import OrderEnforcing
-from safety_gymnasium.wrappers.time_limit import TimeLimit
+from safety_gymnasium.wrappers.env_checker import SafePassiveEnvChecker
+from safety_gymnasium.wrappers.time_limit import SafeTimeLimit
diff --git a/envs/safety-gymnasium/safety_gymnasium/wrappers/env_checker.py b/envs/safety-gymnasium/safety_gymnasium/wrappers/env_checker.py
index d904a1218..e65e710d6 100644
--- a/envs/safety-gymnasium/safety_gymnasium/wrappers/env_checker.py
+++ b/envs/safety-gymnasium/safety_gymnasium/wrappers/env_checker.py
@@ -12,60 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""A passive environment checker wrapper for an environment's observation and action space along with the reset, step and render functions."""
-# pylint: disable=all
+"""Passive environment checker."""
 
-import gymnasium
 from gymnasium.core import ActType
-from safety_gymnasium.utils.passive_env_checker import (
-    check_action_space,
-    check_observation_space,
-    env_render_passive_checker,
-    env_reset_passive_checker,
-    env_step_passive_checker,
-)
+from gymnasium.wrappers.env_checker import PassiveEnvChecker
+from safety_gymnasium.utils.passive_env_checker import env_step_passive_checker
 
 
-class PassiveEnvChecker(gymnasium.Wrapper):
-    """A passive environment checker wrapper that surrounds the step, reset and render functions to check they follow the gymnasium API."""
+class SafePassiveEnvChecker(PassiveEnvChecker):
+    """Passive environment checker.
 
-    def __init__(self, env):
-        """Initialises the wrapper with the environments, run the observation and action space tests."""
-        super().__init__(env)
-
-        assert hasattr(
-            env, 'action_space'
-        ), 'The environment must specify an action space. https://gymnasium.farama.org/content/environment_creation/'
-        check_action_space(env.action_space)
-        assert hasattr(
-            env, 'observation_space'
-        ), 'The environment must specify an observation space. https://gymnasium.farama.org/content/environment_creation/'
-        check_observation_space(env.observation_space)
-
-        self.checked_reset = False
-        self.checked_step = False
-        self.checked_render = False
+    A passive environment checker wrapper for an environment's observation and action space
+    along with the reset, step and render functions.
+    """
 
     def step(self, action: ActType):
         """Steps through the environment that on the first call will run the `passive_env_step_check`."""
         if self.checked_step is False:
             self.checked_step = True
             return env_step_passive_checker(self.env, action)
-        else:
-            return self.env.step(action)
-
-    def reset(self, **kwargs):
-        """Resets the environment that on the first call will run the `passive_env_reset_check`."""
-        if self.checked_reset is False:
-            self.checked_reset = True
-            return env_reset_passive_checker(self.env, **kwargs)
-        else:
-            return self.env.reset(**kwargs)
 
-    def render(self, *args, **kwargs):
-        """Renders the environment that on the first call will run the `passive_env_render_check`."""
-        if self.checked_render is False:
-            self.checked_render = True
-            return env_render_passive_checker(self.env, *args, **kwargs)
-        else:
-            return self.env.render(*args, **kwargs)
+        return self.env.step(action)
diff --git a/envs/safety-gymnasium/safety_gymnasium/wrappers/order_enforcing.py b/envs/safety-gymnasium/safety_gymnasium/wrappers/order_enforcing.py
deleted file mode 100644
index d1b9f930a..000000000
--- a/envs/safety-gymnasium/safety_gymnasium/wrappers/order_enforcing.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# Copyright 2022 OmniSafe Team. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Wrapper to enforce the proper ordering of environment operations."""
-# pylint: disable=all
-
-import gymnasium
-from gymnasium.error import ResetNeeded
-
-
-class OrderEnforcing(gymnasium.Wrapper):
-    """A wrapper that will produce an error if :meth:`step` is called before an initial :meth:`reset`.
-
-    Example:
-        >>> from gymnasium.envs.classic_control import CartPoleEnv
-        >>> env = CartPoleEnv()
-        >>> env = OrderEnforcing(env)
-        >>> env.step(0)
-        ResetNeeded: Cannot call env.step() before calling env.reset()
-        >>> env.render()
-        ResetNeeded: Cannot call env.render() before calling env.reset()
-        >>> env.reset()
-        >>> env.render()
-        >>> env.step(0)
-    """
-
-    def __init__(self, env: gymnasium.Env, disable_render_order_enforcing: bool = False):
-        """A wrapper that will produce an error if :meth:`step` is called before an initial :meth:`reset`.
-
-        Args:
-            env: The environment to wrap
-            disable_render_order_enforcing: If to disable render order enforcing
-        """
-        super().__init__(env)
-        self._has_reset: bool = False
-        self._disable_render_order_enforcing: bool = disable_render_order_enforcing
-
-    def step(self, action):
-        """Steps through the environment with `kwargs`."""
-        if not self._has_reset:
-            raise ResetNeeded('Cannot call env.step() before calling env.reset()')
-        return self.env.step(action)
-
-    def reset(self, **kwargs):
-        """Resets the environment with `kwargs`."""
-        self._has_reset = True
-        return self.env.reset(**kwargs)
-
-    def render(self, *args, **kwargs):
-        """Renders the environment with `kwargs`."""
-        if not self._disable_render_order_enforcing and not self._has_reset:
-            raise ResetNeeded(
-                'Cannot call `env.render()` before calling `env.reset()`, if this is a intended action, '
-                'set `disable_render_order_enforcing=True` on the OrderEnforcer wrapper.'
-            )
-        return self.env.render(*args, **kwargs)
-
-    @property
-    def has_reset(self):
-        """Returns if the environment has been reset before."""
-        return self._has_reset
diff --git a/envs/safety-gymnasium/safety_gymnasium/wrappers/time_limit.py b/envs/safety-gymnasium/safety_gymnasium/wrappers/time_limit.py
index 45b9c6f60..784bdc361 100644
--- a/envs/safety-gymnasium/safety_gymnasium/wrappers/time_limit.py
+++ b/envs/safety-gymnasium/safety_gymnasium/wrappers/time_limit.py
@@ -13,18 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 """Wrapper for limiting the time steps of an environment."""
-# pylint: disable=all
 
-from typing import Optional
+from gymnasium.wrappers.time_limit import TimeLimit
 
-import gymnasium
 
-
-class TimeLimit(gymnasium.Wrapper):
+class SafeTimeLimit(TimeLimit):
     """This wrapper will issue a `truncated` signal if a maximum number of timesteps is exceeded.
 
-    If a truncation is not defined inside the environment itself, this is the only place that the truncation signal is issued.
-    Critically, this is different from the `terminated` signal that originates from the underlying environment as part of the MDP.
+    If a truncation is not defined inside the environment itself,
+    this is the only place that the truncation signal is issued.
+    Critically, this is different from the `terminated` signal
+    that originates from the underlying environment as part of the MDP.
 
     Example:
        >>> from gymnasium.envs.classic_control import CartPoleEnv
@@ -33,25 +32,6 @@ class TimeLimit(gymnasium.Wrapper):
        >>> env = TimeLimit(env, max_episode_steps=1000)
     """
 
-    def __init__(
-        self,
-        env: gymnasium.Env,
-        max_episode_steps: Optional[int] = None,
-    ):
-        """Initializes the :class:`TimeLimit` wrapper with an environment and the number of steps after which truncation will occur.
-
-        Args:
-            env: The environment to apply the wrapper
-            max_episode_steps: An optional max episode steps (if ``Ǹone``, ``env.spec.max_episode_steps`` is used)
-        """
-        super().__init__(env)
-        if max_episode_steps is None and self.env.spec is not None:
-            max_episode_steps = env.spec.max_episode_steps
-        if self.env.spec is not None:
-            self.env.spec.max_episode_steps = max_episode_steps
-        self._max_episode_steps = max_episode_steps
-        self._elapsed_steps = None
-
     def step(self, action):
         """Steps through the environment and if the number of steps elapsed exceeds ``max_episode_steps`` then truncate.
 
@@ -70,15 +50,3 @@ def step(self, action):
             truncated = True
 
         return observation, reward, cost, terminated, truncated, info
-
-    def reset(self, **kwargs):
-        """Resets the environment with :param:`**kwargs` and sets the number of steps elapsed to zero.
-
-        Args:
-            **kwargs: The kwargs to reset the environment with
-
-        Returns:
-            The reset environment
-        """
-        self._elapsed_steps = 0
-        return self.env.reset(**kwargs)


Test patch:
diff --git a/envs/safety-gymnasium/tests/helpers.py b/envs/safety-gymnasium/tests/helpers.py
index f16bae013..5230ab8b2 100644
--- a/envs/safety-gymnasium/tests/helpers.py
+++ b/envs/safety-gymnasium/tests/helpers.py
@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# pylint: disable=all
+"""A parameterize decorator that allows for multiple dtypes to be tested."""
 
 import itertools
 
-import pytest
+import pytest  # pylint: disable=import-error
 
 
 def parametrize(**argvalues) -> pytest.mark.parametrize:
+    """A parameterize decorator that allows for multiple dtypes to be tested."""
     arguments = list(argvalues)
 
     if 'dtype' in argvalues:
diff --git a/envs/safety-gymnasium/tests/test_env.py b/envs/safety-gymnasium/tests/test_env.py
index db7d3aad1..490a76ecc 100644
--- a/envs/safety-gymnasium/tests/test_env.py
+++ b/envs/safety-gymnasium/tests/test_env.py
@@ -62,7 +62,7 @@ def test_env(agent_id, env_id, level, render_mode):
     render_mode=['rgb_array', 'depth_array'],
 )
 # pylint: disable-next=too-many-locals
-def test_new_env(agent_id, env_id, level, render_mode):
+def test_run_env(agent_id, env_id, level, render_mode):
     """Test env."""
     env_name = 'Safety' + agent_id + env_id + level + '-v0'
     env = safety_gymnasium.make(env_name, render_mode=render_mode)
@@ -89,3 +89,38 @@ def test_new_env(agent_id, env_id, level, render_mode):
         ep_cost += cost
 
         env.render()
+
+
+@helpers.parametrize(
+    agent_id=['Humanoid', 'Ant', 'Hopper', 'HalfCheetah', 'Swimmer', 'Walker2d'],
+    env_id=['Velocity'],
+    render_mode=['rgb_array', 'depth_array'],
+)
+# pylint: disable-next=too-many-locals
+def test_velocity_env(agent_id, env_id, render_mode):
+    """Test env."""
+    env_name = 'Safety' + agent_id + env_id + '-v4'
+    env = safety_gymnasium.make(env_name, render_mode=render_mode)
+    obs, _ = env.reset()
+    # Use below to specify seed.
+    # obs, _ = env.reset(seed=0)
+    terminated, truncated = False, False
+    ep_ret, ep_cost = 0, 0
+    for step in range(1000):  # pylint: disable=unused-variable
+        if terminated or truncated:
+            print(f'Episode Return: {ep_ret} \t Episode Cost: {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'):  # pylint: disable=protected-access
+            max_ep_len = env._max_episode_steps  # pylint: disable=unused-variable,protected-access
+        # pylint: disable-next=unused-variable
+        obs, reward, cost, terminated, truncated, info = env.step(act)
+        ep_ret += reward
+        ep_cost += cost
+
+        env.render()
diff --git a/envs/safety-gymnasium/tests/test_vision_env.py b/envs/safety-gymnasium/tests/test_vision_env.py
index 4c84ffc3c..e6b0f8b71 100644
--- a/envs/safety-gymnasium/tests/test_vision_env.py
+++ b/envs/safety-gymnasium/tests/test_vision_env.py
@@ -56,10 +56,10 @@ def test_vision_env(agent_id, env_id, level):
     level=['0'],
 )
 # pylint: disable-next=too-many-locals
-def test_new_env(agent_id, env_id, level, render_mode):
+def test_new_env(agent_id, env_id, level):
     """Test env."""
     env_name = 'Safety' + agent_id + env_id + level + 'Vision' + '-v0'
-    env = safety_gymnasium.make(env_name, render_mode=render_mode)
+    env = safety_gymnasium.make(env_name)
     obs, _ = env.reset()
     # Use below to specify seed.
     # obs, _ = env.reset(seed=0)
@@ -79,7 +79,6 @@ def test_new_env(agent_id, env_id, level, render_mode):
             max_ep_len = env._max_episode_steps  # pylint: disable=unused-variable,protected-access
         # pylint: disable-next=unused-variable
         obs, reward, cost, terminated, truncated, info = env.step(act)
+
         ep_ret += reward
         ep_cost += cost
-
-        env.render()