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

Issue:
[Feature Request] Style of mypy and docstring of Model-based algorithms
### Required prerequisites

- [X] I have searched the [Issue Tracker](https://github.com/OmniSafeAI/omnisafe/issues) and [Discussions](https://github.com/OmniSafeAI/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/OmniSafeAI/omnisafe/discussions/new).

### Motivation

<!-- Please outline the motivation for the proposal.
Is your feature request related to a problem? E.g., "I'm always frustrated when [...]".
If this is related to another issue, please link here too. -->
I notice the code style and docstring style may not so perfect and it needs some further polish.
e.g. unused argument, puzzling typing and methods docstring.

### Solution

I will open a pull request to fix those problems in PR #212 and I will invite @hdadong to make sure I make no mistake in Model-based code.

### Alternatives

_No response_

### Additional context

_No response_


Discussion / hints:

Fix patch:
diff --git a/docs/source/baserlapi/model_based.rst b/docs/source/baserlapi/model_based.rst
new file mode 100644
index 000000000..e4356f0c0
--- /dev/null
+++ b/docs/source/baserlapi/model_based.rst
@@ -0,0 +1,33 @@
+Base Model-based Algorithms
+===========================
+
+.. currentmodule:: omnisafe.algorithms.model_based.base
+
+
+LOOP
+----
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: LOOP
+        :members:
+        :private-members:
+
+PETS
+----
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: PETS
+        :members:
+        :private-members:
diff --git a/docs/source/envs/adapter.rst b/docs/source/envs/adapter.rst
index 4fd78c637..10187d5c6 100644
--- a/docs/source/envs/adapter.rst
+++ b/docs/source/envs/adapter.rst
@@ -3,6 +3,16 @@ OmniSafe Adapter
 
 .. currentmodule:: omnisafe.adapter
 
+OmniSafe provides a set of adapters to adapt the environment to the framework.
+
+.. autosummary::
+
+    OnlineAdapter
+    OnPolicyAdapter
+    OffPolicyAdapter
+    SauteAdapter
+    SimmerAdapter
+    ModelBasedAdapter
 
 Online Adapter
 --------------
@@ -73,3 +83,17 @@ Simmer Adapter
     .. autoclass:: SimmerAdapter
         :members:
         :private-members:
+
+Model-based Adapter
+-------------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: ModelBasedAdapter
+        :members:
+        :private-members:
diff --git a/docs/source/envs/mujoco_env.rst b/docs/source/envs/mujoco_env.rst
new file mode 100644
index 000000000..7765c28c3
--- /dev/null
+++ b/docs/source/envs/mujoco_env.rst
@@ -0,0 +1,18 @@
+Mujoco Environment
+==================
+
+.. currentmodule:: omnisafe.envs.mujoco_env
+
+MujocoEnv Interface
+-------------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: MujocoEnv
+        :members:
+        :private-members:
diff --git a/docs/source/envs/safety_gymnasium.rst b/docs/source/envs/safety_gymnasium.rst
index f5713c966..ce9c250c4 100644
--- a/docs/source/envs/safety_gymnasium.rst
+++ b/docs/source/envs/safety_gymnasium.rst
@@ -16,3 +16,19 @@ Safety Gymnasium Interface
     .. autoclass:: SafetyGymnasiumEnv
         :members:
         :private-members:
+
+.. currentmodule:: omnisafe.envs.safety_gymnasium_modelbased
+
+Safety Gymnasium World Model
+----------------------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: SafetyGymnasiumModelBased
+        :members:
+        :private-members:
diff --git a/docs/source/envs/wrapper.rst b/docs/source/envs/wrapper.rst
index 456ca5b6e..43d966929 100644
--- a/docs/source/envs/wrapper.rst
+++ b/docs/source/envs/wrapper.rst
@@ -98,6 +98,19 @@ Action Scale
         :members:
         :private-members:
 
+Action Repeat
+-------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: ActionRepeat
+        :members:
+        :private-members:
 
 Unsqueeze Wrapper
 -----------------
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 301bd7acc..72bd93f2d 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -395,6 +395,7 @@ this project, don't hesitate to ask your question on `the GitHub issue page <htt
 
     baserlapi/on_policy
     baserlapi/off_policy
+    baserlapi/model_based
 
 .. toctree::
     :hidden:
@@ -405,6 +406,7 @@ this project, don't hesitate to ask your question on `the GitHub issue page <htt
     saferlapi/second_order
     saferlapi/lagrange
     saferlapi/penalty_function
+    saferlapi/model_based
 
 .. toctree::
     :hidden:
@@ -439,6 +441,8 @@ this project, don't hesitate to ask your question on `the GitHub issue page <htt
     model/actor
     model/critic
     model/actor_critic
+    model/modelbased_model
+    model/modelbased_planner
 
 .. toctree::
     :hidden:
@@ -448,6 +452,7 @@ this project, don't hesitate to ask your question on `the GitHub issue page <htt
     envs/core
     envs/wrapper
     envs/safety_gymnasium
+    envs/mujoco_env
     envs/adapter
 
 
diff --git a/docs/source/model/modelbased_model.rst b/docs/source/model/modelbased_model.rst
new file mode 100644
index 000000000..3aca25c49
--- /dev/null
+++ b/docs/source/model/modelbased_model.rst
@@ -0,0 +1,85 @@
+OmniSafe Model-based Model
+==========================
+
+.. currentmodule:: omnisafe.algorithms.model_based.base.ensemble
+
+Standard Scaler
+---------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: StandardScaler
+        :members:
+        :private-members:
+
+Initialize Weight
+-----------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autofunction:: init_weights
+
+Unbatched Forward
+-----------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autofunction:: unbatched_forward
+
+Ensemble Fully-Connected Layer
+------------------------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: EnsembleFC
+        :members:
+        :private-members:
+
+Ensemble Model
+--------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: EnsembleModel
+        :members:
+        :private-members:
+
+
+Ensemble Dynamics Model
+-----------------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: EnsembleDynamicsModel
+        :members:
+        :private-members:
diff --git a/docs/source/model/modelbased_planner.rst b/docs/source/model/modelbased_planner.rst
new file mode 100644
index 000000000..cf5f6f861
--- /dev/null
+++ b/docs/source/model/modelbased_planner.rst
@@ -0,0 +1,89 @@
+OmniSafe Model-based Planner
+============================
+
+.. currentmodule:: omnisafe.algorithms.model_based.planner
+
+ARC Planner
+-----------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: ARCPlanner
+        :members:
+        :private-members:
+
+CAP Planner
+-----------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: CAPPlanner
+        :members:
+        :private-members:
+
+CCE Planner
+-----------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: CCEPlanner
+        :members:
+        :private-members:
+
+CEM Planner
+-----------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: CEMPlanner
+        :members:
+        :private-members:
+
+RCE Planner
+-----------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: RCEPlanner
+        :members:
+        :private-members:
+
+
+SafeARC Planner
+---------------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: SafeARCPlanner
+        :members:
+        :private-members:
diff --git a/docs/source/saferlapi/first_order.rst b/docs/source/saferlapi/first_order.rst
index 0f6fe93ea..4cae33626 100644
--- a/docs/source/saferlapi/first_order.rst
+++ b/docs/source/saferlapi/first_order.rst
@@ -8,10 +8,8 @@ First Order Algorithms
     FOCOPS
     CUP
 
-
 .. _focopsapi:
 
-
 FOCOPS
 ------
 
diff --git a/docs/source/saferlapi/model_based.rst b/docs/source/saferlapi/model_based.rst
new file mode 100644
index 000000000..7efcd2cc3
--- /dev/null
+++ b/docs/source/saferlapi/model_based.rst
@@ -0,0 +1,60 @@
+Model-based Algorithms
+======================
+
+.. currentmodule:: omnisafe.algorithms.model_based
+
+CAPPETS
+-------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: CAPPETS
+        :members:
+        :private-members:
+
+CCEPETS
+-------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: CCEPETS
+        :members:
+        :private-members:
+
+RCEPETS
+-------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: RCEPETS
+        :members:
+        :private-members:
+
+Safe LOOP
+---------
+
+.. card::
+    :class-header: sd-bg-success sd-text-white
+    :class-card: sd-outline-success  sd-rounded-1
+
+    Documentation
+    ^^^
+
+    .. autoclass:: SafeLOOP
+        :members:
+        :private-members:
diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt
index 72d5bcf6f..6778d1937 100644
--- a/docs/source/spelling_wordlist.txt
+++ b/docs/source/spelling_wordlist.txt
@@ -467,3 +467,19 @@ Jost
 Autoencoder
 Qc
 Qr
+Pinneri
+Shambhuraj
+Sawant
+Blaes
+Achterhold
+Joerg
+Rolinek
+Martius
+Georg
+Michal
+Stueckler
+Unbatched
+DynamicsTrainMseLoss
+DynamicsValMseLoss
+UpdateActorCritic
+UpdateDynamics
diff --git a/docs/source/utils/distributed.rst b/docs/source/utils/distributed.rst
index 77c84128e..127df84e4 100644
--- a/docs/source/utils/distributed.rst
+++ b/docs/source/utils/distributed.rst
@@ -20,7 +20,6 @@ OmniSafe Distributed
     dist_op
     dist_statistics_scalar
 
-
 Set up distributed training
 ---------------------------
 
diff --git a/omnisafe/adapter/modelbased_adapter.py b/omnisafe/adapter/modelbased_adapter.py
index e663a5a6d..8abbd90d7 100644
--- a/omnisafe/adapter/modelbased_adapter.py
+++ b/omnisafe/adapter/modelbased_adapter.py
@@ -13,16 +13,20 @@
 # limitations under the License.
 # ==============================================================================
 """Model-based Adapter for OmniSafe."""
+
+
 from __future__ import annotations
 
 import time
 from typing import Any, Callable
 
+import numpy as np
 import torch
+from gymnasium.spaces import Box
 
 from omnisafe.adapter.online_adapter import OnlineAdapter
 from omnisafe.common.logger import Logger
-from omnisafe.envs.core import make, support_envs
+from omnisafe.envs.core import CMDP, make, support_envs
 from omnisafe.envs.wrapper import (
     ActionRepeat,
     ActionScale,
@@ -34,6 +38,7 @@
     Unsqueeze,
 )
 from omnisafe.utils.config import Config
+from omnisafe.utils.tools import get_device
 
 
 class ModelBasedAdapter(
@@ -42,46 +47,56 @@ class ModelBasedAdapter(
     """Model Based Adapter for OmniSafe.
 
     :class:`ModelBasedAdapter` is used to adapt the environment to the model-based training.
-
+    It trains a world model to provide data for algorithms training.
 
     Args:
         env_id (str): The environment id.
         num_envs (int): The number of environments.
         seed (int): The random seed.
         cfgs (Config): The configuration.
-        kwargs(dict): The other keyword arguments.
+
+    Keyword Args:
+        render_mode (str, optional): The render mode ranges from 'human' to 'rgb_array' and 'rgb_array_list'.
+            Defaults to 'rgb_array'.
+        camera_name (str, optional): The camera name.
+        camera_id (int, optional): The camera id.
+        width (int, optional): The width of the rendered image. Defaults to 256.
+        height (int, optional): The height of the rendered image. Defaults to 256.
 
     Attributes:
-        _env_id (str): The environment id.
-        _device (torch.device): The device.
-        _env (CMDP): The environment.
-        _cfgs (Config): The configuration.
-        _ep_ret (torch.Tensor): The episode return.
-        _ep_cost (torch.Tensor): The episode cost.
-        _ep_len (torch.Tensor): The episode length.
-        _last_dynamics_update (float): The last time of dynamics update.
-        _last_policy_update (float): The last time of policy update.
-        _last_eval (float): The last time of evaluation.
         coordinate_observation_space (OmnisafeSpace): The coordinate observation space.
         lidar_observation_space (OmnisafeSpace): The lidar observation space.
         task (str): The task. eg. The task of SafetyPointGoal-v0 is 'goal'
     """
 
+    coordinate_observation_space: Box | None
+    lidar_observation_space: Box | None
+    task: str | None
+    _ep_ret: torch.Tensor
+    _ep_cost: torch.Tensor
+    _ep_len: torch.Tensor
+    _current_obs: torch.Tensor
+
     def __init__(  # pylint: disable=too-many-arguments
         self,
         env_id: str,
         num_envs: int,
         seed: int,
         cfgs: Config,
-        **kwargs: Any,
+        **env_kwargs: Any,
     ) -> None:
         """Initialize the model-based adapter."""
         assert env_id in support_envs(), f'Env {env_id} is not supported.'
 
-        self._env_id = env_id
-        self._device = cfgs.train_cfgs.device
+        self._env_id: str = env_id
+        self._device: torch.device = get_device(cfgs.train_cfgs.device)
 
-        self._env = make(env_id, num_envs=num_envs, device=cfgs.train_cfgs.device, **kwargs)
+        self._env: CMDP = make(
+            env_id,
+            num_envs=num_envs,
+            device=cfgs.train_cfgs.device,
+            **env_kwargs,
+        )
 
         # wrap the environment, use the action repeat in model-based setting.
         self._wrapper(
@@ -91,7 +106,7 @@ def __init__(  # pylint: disable=too-many-arguments
             action_repeat=cfgs.algo_cfgs.action_repeat,
         )
         self._env.set_seed(seed)
-        self._cfgs = cfgs
+        self._cfgs: Config = cfgs
         if hasattr(self._env, 'coordinate_observation_space') and hasattr(
             self._env,
             'lidar_observation_space',
@@ -106,30 +121,27 @@ def __init__(  # pylint: disable=too-many-arguments
         else:
             self.task = None
 
-        self._ep_ret: torch.Tensor
-        self._ep_cost: torch.Tensor
-        self._ep_len: torch.Tensor
         self._current_obs, _ = self.reset()
-        self._max_ep_len = 1000
+        self._max_ep_len: int = 1000
         self._reset_log()
-        self._last_dynamics_update = 0
-        self._last_policy_update = 0
-        self._last_eval = 0
-        self._first_log = False
+        self._last_dynamics_update: int = 0
+        self._last_policy_update: int = 0
+        self._last_eval: int = 0
+        self._first_log: bool = False
 
-    def get_cost_from_obs_tensor(self, obs: torch.Tensor) -> torch.Tensor | None:
+    def get_cost_from_obs_tensor(self, obs: torch.Tensor) -> torch.Tensor:
         """Get cost from tensor observation.
 
         Args:
-            obs (torch.Tensor): The observation.
+            obs (torch.Tensor): The tensor version of observation.
         """
         return (
             self._env.get_cost_from_obs_tensor(obs)
             if hasattr(self._env, 'get_cost_from_obs_tensor')
-            else None
+            else torch.zeros(1)
         )
 
-    def get_lidar_from_coordinate(self, obs: torch.Tensor) -> torch.Tensor | None:
+    def get_lidar_from_coordinate(self, obs: np.ndarray) -> torch.Tensor | None:
         """Get lidar from numpy coordinate.
 
         Args:
@@ -141,12 +153,19 @@ def get_lidar_from_coordinate(self, obs: torch.Tensor) -> torch.Tensor | None:
             else None
         )
 
-    def render(self, *args: str, **kwargs: int) -> Any:
+    def render(self, *args: str, **kwargs: Any) -> Any:
         """Render the environment.
 
         Args:
             args (str): The arguments.
-            kwargs (int): The keyword arguments.
+
+        Keyword Args:
+            render_mode (str, optional): The render mode, ranging from ``human``, ``rgb_array``, ``rgb_array_list``.
+                Defaults to ``rgb_array``.
+            camera_name (str, optional): The camera name.
+            camera_id (int, optional): The camera id.
+            width (int, optional): The width of the rendered image. Defaults to 256.
+            height (int, optional): The height of the rendered image. Defaults to 256.
         """
         return self._env.render(*args, **kwargs)
 
@@ -163,26 +182,25 @@ def _wrapper(
 
             OmniSafe supports the following wrappers:
 
-            .. list-table::
-
-                *   -   Wrapper
-                    -   Description
-                *   -   TimeLimit
-                    -   Limit the time steps of the environment.
-                *   -   AutoReset
-                    -   Reset the environment when the episode is done.
-                *   -   ObsNormalize
-                    -   Normalize the observation.
-                *   -   RewardNormalize
-                    -   Normalize the reward.
-                *   -   CostNormalize
-                    -   Normalize the cost.
-                *   -   ActionScale
-                    -   Scale the action.
-                *   -   ActionRepeat
-                    -   Repeat the action.
-                *   -   Unsqueeze
-                    -   Unsqueeze the step result for single environment case.
+            +-----------------+--------------------------------------------------------+
+            | Wrapper         | Description                                            |
+            +=================+========================================================+
+            | TimeLimit       | Limit the time steps of the environment.               |
+            +-----------------+--------------------------------------------------------+
+            | AutoReset       | Reset the environment when the episode is done.        |
+            +-----------------+--------------------------------------------------------+
+            | ObsNormalize    | Normalize the observation.                             |
+            +-----------------+--------------------------------------------------------+
+            | RewardNormalize | Normalize the reward.                                  |
+            +-----------------+--------------------------------------------------------+
+            | CostNormalize   | Normalize the cost.                                    |
+            +-----------------+--------------------------------------------------------+
+            | ActionScale     | Scale the action.                                      |
+            +-----------------+--------------------------------------------------------+
+            | ActionRepeat    | Repeat the action.                                     |
+            +-----------------+--------------------------------------------------------+
+            | Unsqueeze       | Unsqueeze the step result for single environment case. |
+            +-----------------+--------------------------------------------------------+
 
         Args:
             obs_normalize (bool): Whether to normalize the observation.
@@ -206,34 +224,46 @@ def _wrapper(
         if self._env.num_envs == 1:
             self._env = Unsqueeze(self._env, device=self._device)
 
-    def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
+    def rollout(  # pylint: disable=too-many-arguments,too-many-locals
         self,
         current_step: int,
-        roll_out_step: int,
+        rollout_step: int,
         use_actor_critic: bool,
-        act_func: Callable,
-        store_data_func: Callable,
-        update_dynamics_func: Callable,
+        act_func: Callable[[int, torch.Tensor], torch.Tensor],
+        store_data_func: Callable[
+            [
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                dict[str, Any],
+            ],
+            None,
+        ],
+        update_dynamics_func: Callable[[], None],
         logger: Logger,
         use_eval: bool,
-        eval_func: Callable,
-        algo_reset_func: Callable,
-        update_actor_func: Callable,
+        eval_func: Callable[[int, bool], None],
+        algo_reset_func: Callable[[], None],
+        update_actor_func: Callable[[int], None],
     ) -> int:
         """Roll out the environment and store the data in the buffer.
 
         Args:
             current_step (int): Current training step.
-            roll_out_step (int): Number of steps to roll out.
+            rollout_step (int): Number of steps to roll out.
             use_actor_critic (bool): Whether to use actor-critic.
-            act_func (Callable): Function to get action.
-            store_data_func (Callable): Function to store data.
-            update_dynamics_func (Callable): Function to update dynamics.
-            logger (Logger): Logger.
+            act_func (Callable[[int, torch.Tensor], torch.Tensor]): Function to get action.
+            store_data_func (Callable[[torch.Tensor, ..., dict[str, Any], ], None,]): Function to store data.
+            update_dynamics_func (Callable[[], None]): Function to update dynamics.
+            logger (Logger): Logger, to log ``EpRet``, ``EpCost``, ``EpLen``.
             use_eval (bool): Whether to use evaluation.
-            eval_func (Callable): Function to evaluate the agent.
-            algo_reset_func (Callable): Function to reset the algorithm.
-            update_actor_func (Callable): Function to update the actor.
+            eval_func (Callable[[int, bool], None]): Function to evaluate the agent.
+            algo_reset_func (Callable[[], None]): Function to reset the algorithm.
+            update_actor_func (Callable[[int], None]): Function to update the actor.
         """
         epoch_start_time = time.time()
 
@@ -244,16 +274,14 @@ def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
 
         epoch_steps = 0
 
-        while epoch_steps < roll_out_step and current_step < self._cfgs.train_cfgs.total_steps:
-            action, action_info = act_func(current_step, self._current_obs, self._env)
+        while epoch_steps < rollout_step and current_step < self._cfgs.train_cfgs.total_steps:
+            action = act_func(current_step, self._current_obs)
             next_state, reward, cost, terminated, truncated, info = self.step(action)
             epoch_steps += info['num_step']
             current_step += info['num_step']
             self._log_value(reward=reward, cost=cost, info=info)
 
             store_data_func(
-                current_step,
-                self._ep_len,
                 self._current_obs,
                 action,
                 reward,
@@ -262,7 +290,6 @@ def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
                 truncated,
                 next_state,
                 info,
-                action_info,
             )
             self._current_obs = next_state
             if terminated or truncated:
@@ -270,7 +297,7 @@ def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
                 self._reset_log()
                 self._current_obs, _ = self.reset()
                 if algo_reset_func is not None:
-                    algo_reset_func(current_step)
+                    algo_reset_func()
             if (
                 current_step % self._cfgs.algo_cfgs.update_dynamics_cycle
                 < self._cfgs.algo_cfgs.action_repeat
@@ -278,7 +305,7 @@ def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
                 >= self._cfgs.algo_cfgs.update_dynamics_cycle
             ):
                 update_dynamics_start = time.time()
-                update_dynamics_func(current_step)
+                update_dynamics_func()
                 self._last_dynamics_update = current_step
                 update_dynamics_time += time.time() - update_dynamics_start
 
@@ -301,7 +328,7 @@ def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
                 and current_step - self._last_eval >= self._cfgs.evaluation_cfgs.eval_cycle
             ):
                 eval_start = time.time()
-                eval_func(current_step)
+                eval_func(current_step, True)
                 self._last_eval = current_step
                 eval_time += time.time() - eval_start
 
@@ -311,23 +338,23 @@ def roll_out(  # pylint: disable=too-many-arguments,too-many-locals
         epoch_time = time.time() - epoch_start_time
         logger.store(**{'Time/Epoch': epoch_time})
         logger.store(**{'Time/UpdateDynamics': update_dynamics_time})
-        roll_out_time = epoch_time - update_dynamics_time
+        rollout_time = epoch_time - update_dynamics_time
 
         if use_eval:
             logger.store(**{'Time/Eval': eval_time})
-            roll_out_time -= eval_time
+            rollout_time -= eval_time
 
         if use_actor_critic:
             logger.store(**{'Time/UpdateActorCritic': update_actor_critic_time})
-            roll_out_time -= update_actor_critic_time
-        logger.store(**{'Time/Rollout': roll_out_time})
+            rollout_time -= update_actor_critic_time
+        logger.store(**{'Time/Rollout': rollout_time})
         return current_step
 
     def _log_value(
         self,
         reward: torch.Tensor,
         cost: torch.Tensor,
-        info: dict,
+        info: dict[str, Any],
     ) -> None:
         """Log value.
 
@@ -336,9 +363,9 @@ def _log_value(
             be stored in ``info['original_reward']`` and ``info['original_cost']``.
 
         Args:
-            reward (torch.Tensor): The reward.
-            cost (torch.Tensor): The cost.
-            info (dict): Information.
+            reward (torch.Tensor): The immediate step reward.
+            cost (torch.Tensor): The immediate step cost.
+            info (dict[str, Any]): Some information logged by the environment.
         """
         self._ep_ret += info.get('original_reward', reward).cpu()
         self._ep_cost += info.get('original_cost', cost).cpu()
@@ -348,23 +375,19 @@ def _log_metrics(self, logger: Logger) -> None:
         """Log metrics.
 
         Args:
-            logger (Logger): Logger.
+            logger (Logger): Logger, to log ``EpRet``, ``EpCost``, ``EpLen``.
         """
         self._first_log = True
         logger.store(
-            **{
+            {
                 'Metrics/EpRet': self._ep_ret,
                 'Metrics/EpCost': self._ep_cost,
                 'Metrics/EpLen': self._ep_len,
             },
         )
 
-    def _reset_log(self, idx: int | None = None) -> None:  # pylint: disable=unused-argument
-        """Reset log.
-
-        Args:
-            idx (int | None): The index of the environment.
-        """
+    def _reset_log(self) -> None:
+        """Reset log."""
         self._ep_ret = torch.zeros(1)
         self._ep_cost = torch.zeros(1)
         self._ep_len = torch.zeros(1)
diff --git a/omnisafe/adapter/online_adapter.py b/omnisafe/adapter/online_adapter.py
index 9b42e4d29..f61444b26 100644
--- a/omnisafe/adapter/online_adapter.py
+++ b/omnisafe/adapter/online_adapter.py
@@ -42,14 +42,6 @@ class OnlineAdapter:
     any existing RL algorithms. The online adapter is used to adapt the environment to the
     framework.
 
-    OmniSafe provides a set of adapters to adapt the environment to the framework.
-
-    - OnPolicyAdapter: Adapt the environment to the on-policy framework.
-    - OffPolicyAdapter: Adapt the environment to the off-policy framework.
-    - OfflineAdapter: Adapt the environment to the offline framework.
-    - SauteAdapter: Adapt the environment to the SAUTE framework.
-    - SimmerAdapter: Adapt the environment to the SIMMER framework.
-
     Args:
         env_id (str): The environment id.
         num_envs (int): The number of parallel environments.
diff --git a/omnisafe/algorithms/model_based/base/ensemble.py b/omnisafe/algorithms/model_based/base/ensemble.py
index 9e7cc3d20..6611517d2 100644
--- a/omnisafe/algorithms/model_based/base/ensemble.py
+++ b/omnisafe/algorithms/model_based/base/ensemble.py
@@ -16,6 +16,8 @@
 # original version doesn't validate model error batch-wise and is highly memory intensive.
 # ==============================================================================
 """The Dynamics Model of MBPO and PETS."""
+
+
 from __future__ import annotations
 
 import itertools
@@ -38,21 +40,28 @@ def swish(data: torch.Tensor) -> torch.Tensor:
 
 
 class StandardScaler:
-    """Normalize data."""
+    """Normalizes data using standardization.
+
+    This class provides methods to fit the scaler to the input data and transform
+    the input data using the parameters learned during the fitting process.
+
+    Args:
+        device (torch.device): The device to use.
+    """
 
     def __init__(self, device: torch.device) -> None:
         """Initialize an instance of :class:`StandardScaler`."""
-        self._mean = 0.0
-        self._std = 1.0
-        self._mean_t = torch.tensor(self._mean).to(device)
-        self._std_t = torch.tensor(self._std).to(device)
-        self._device = device
+        self._mean: float = 0.0
+        self._std: float = 1.0
+        self._mean_t: torch.Tensor = torch.tensor(self._mean).to(device)
+        self._std_t: torch.Tensor = torch.tensor(self._std).to(device)
+        self._device: torch.device = device
 
-    def fit(self, data: torch.Tensor | np.ndarray) -> None:
+    def fit(self, data: np.ndarray) -> None:
         """Fits the scaler to the input data.
 
         Args:
-            data (np.ndarray): A numpy array containing the input
+            data (np.ndarray): A numpy array containing the input.
         """
         self._mean = np.mean(data, axis=0, keepdims=True)
         self._std = np.std(data, axis=0, keepdims=True)
@@ -60,29 +69,40 @@ def fit(self, data: torch.Tensor | np.ndarray) -> None:
         self._mean_t = torch.FloatTensor(self._mean).to(self._device)
         self._std_t = torch.FloatTensor(self._std).to(self._device)
 
-    def transform(self, data: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
         """Transforms the input matrix data using the parameters of this scaler.
 
-        Arguments:
-            data (torch.Tensor|np.ndarray): A numpy array containing the input
+        Args:
+            data (torch.Tensor): The input data to transform.
 
         Returns:
-            transformed_data (torch.Tensor|np.ndarray): The transformed data.
+            transformed_data: The transformed data.
         """
-        if torch.is_tensor(data):
-            return (data - self._mean_t) / self._std_t
-        return (data - self._mean) / self._std
+        return (data - self._mean_t) / self._std_t
 
 
 def init_weights(layer: nn.Module) -> None:
-    """Initialize network weight."""
+    """Initialize network weight.
+
+    Args:
+        layer (nn.Module): The layer to initialize.
+    """
 
     def truncated_normal_init(
         weight: torch.Tensor,
         mean: float = 0.0,
         std: float = 0.01,
     ) -> torch.Tensor:
-        """Initialize network weight."""
+        """Initialize network weight.
+
+        Args:
+            weight (torch.Tensor): The weight to be initialized.
+            mean (float): The mean of the normal distribution.
+            std (float): The standard deviation of the normal distribution.
+
+        Returns:
+            weight: The initialized weight.
+        """
         torch.nn.init.normal_(weight, mean=mean, std=std)
         while True:
             cond = torch.logical_or(weight < mean - 2 * std, weight > mean + 2 * std)
@@ -109,12 +129,12 @@ def unbatched_forward(
     """Special forward for nn.Sequential modules which contain BatchedLinear layers we want to use.
 
     Args:
-        layer (nn.Module|EnsembleFC): The layer to forward through.
+        layer (nn.Module | EnsembleFC): The layer to forward through.
         input_data (torch.Tensor): The input data.
         index (int): The index of the model to use.
 
     Returns:
-        output (torch.Tensor): The output of the layer.
+        output: The output of the layer.
     """
     if isinstance(layer, EnsembleFC):
         output = F.linear(
@@ -129,13 +149,30 @@ def unbatched_forward(
 
 
 class EnsembleFC(nn.Module):
-    """Ensemble fully connected network."""
+    """Ensemble fully connected network.
+
+    A fully connected network with ensemble_size models.
 
-    _constants_ = ['in_features', 'out_features']
+    Args:
+        in_features (int): The number of input features.
+        out_features (int): The number of output features.
+        ensemble_size (int): The number of models in the ensemble.
+        weight_decay (float): The decaying factor.
+        bias (bool): Whether to use bias.
+
+    Attributes:
+        in_features (int): The number of input features.
+        out_features (int): The number of output features.
+        ensemble_size (int): The number of models in the ensemble.
+        weight (nn.Parameter): The weight of the network.
+        bias (nn.Parameter): The bias of the network.
+    """
+
+    _constants_: list[str]
     in_features: int
     out_features: int
     ensemble_size: int
-    weight: torch.Tensor
+    weight: nn.Parameter
 
     def __init__(
         self,
@@ -145,8 +182,9 @@ def __init__(
         weight_decay: float = 0.0,
         bias: bool = True,
     ) -> None:
-        """Initialize network weight."""
+        """Initialize an instance of fully connected network."""
         super().__init__()
+        self._constants_ = ['in_features', 'out_features']
         self.in_features = in_features
         self.out_features = out_features
         self.ensemble_size = ensemble_size
@@ -158,7 +196,14 @@ def __init__(
             self.register_parameter('bias', None)
 
     def forward(self, input_data: torch.Tensor) -> torch.Tensor:
-        """Forward pass."""
+        """Forward pass.
+
+        Args:
+            input_data (torch.Tensor): The input data.
+
+        Returns:
+            The forward output of the network.
+        """
         w_times_x = torch.bmm(input_data, self.weight)
         if self.bias is not None:
             return torch.add(w_times_x, self.bias[:, None, :])  # w times x + b
@@ -167,7 +212,31 @@ def forward(self, input_data: torch.Tensor) -> torch.Tensor:
 
 # pylint: disable-next=too-many-instance-attributes
 class EnsembleModel(nn.Module):
-    """Ensemble dynamics model."""
+    """Ensemble dynamics model.
+
+    A dynamics model with ensemble_size models.
+
+    Args:
+        device (torch.device): The device to use.
+        state_size (int): The size of the state.
+        action_size (int): The size of the action.
+        reward_size (int): The size of the reward.
+        cost_size (int): The size of the cost.
+        ensemble_size (int): The number of models in the ensemble.
+        predict_reward (bool): Whether to predict reward.
+        predict_cost (bool): Whether to predict cost.
+        hidden_size (int): The size of the hidden layer.
+        learning_rate (float): The learning rate.
+        use_decay (bool): Whether to use weight decay.
+
+    Attributes:
+        max_logvar (torch.Tensor): The maximum log variance.
+        min_logvar (torch.Tensor): The minimum log variance.
+        scaler (StandardScaler): The scaler.
+    """
+
+    max_logvar: torch.Tensor
+    min_logvar: torch.Tensor
 
     # pylint: disable-next=too-many-arguments
     def __init__(
@@ -179,7 +248,7 @@ def __init__(
         cost_size: int,
         ensemble_size: int,
         predict_reward: bool,
-        predict_cost: bool | None = None,
+        predict_cost: bool = False,
         hidden_size: int = 200,
         learning_rate: float = 1e-3,
         use_decay: bool = False,
@@ -187,31 +256,46 @@ def __init__(
         """Initialize network weight."""
         super().__init__()
 
-        self._state_size = state_size
-        self._reward_size = reward_size
-        self._cost_size = cost_size
-        self._predict_reward = predict_reward
-        self._predict_cost = predict_cost
+        self._state_size: int = state_size
+        self._reward_size: int = reward_size
+        self._cost_size: int = cost_size
+        self._predict_reward: bool = predict_reward
+        self._predict_cost: bool = predict_cost
 
-        self._output_dim = state_size
+        self._output_dim: int = state_size
         if predict_reward:
             self._output_dim += reward_size
         if predict_cost:
             self._output_dim += cost_size
 
-        self._hidden_size = hidden_size
-        self._use_decay = use_decay
+        self._hidden_size: int = hidden_size
+        self._use_decay: bool = use_decay
 
-        self._nn1 = EnsembleFC(
+        self._nn1: EnsembleFC = EnsembleFC(
             state_size + action_size,
             hidden_size,
             ensemble_size,
             weight_decay=0.000025,
         )
-        self._nn2 = EnsembleFC(hidden_size, hidden_size, ensemble_size, weight_decay=0.00005)
-        self._nn3 = EnsembleFC(hidden_size, hidden_size, ensemble_size, weight_decay=0.000075)
-        self._nn4 = EnsembleFC(hidden_size, hidden_size, ensemble_size, weight_decay=0.000075)
-        self._nn5 = EnsembleFC(
+        self._nn2: EnsembleFC = EnsembleFC(
+            hidden_size,
+            hidden_size,
+            ensemble_size,
+            weight_decay=0.00005,
+        )
+        self._nn3: EnsembleFC = EnsembleFC(
+            hidden_size,
+            hidden_size,
+            ensemble_size,
+            weight_decay=0.000075,
+        )
+        self._nn4: EnsembleFC = EnsembleFC(
+            hidden_size,
+            hidden_size,
+            ensemble_size,
+            weight_decay=0.000075,
+        )
+        self._nn5: EnsembleFC = EnsembleFC(
             hidden_size,
             self._output_dim * 2,
             ensemble_size,
@@ -220,10 +304,10 @@ def __init__(
 
         self.register_buffer('max_logvar', (torch.ones((1, self._output_dim)).float() / 2))
         self.register_buffer('min_logvar', (-torch.ones((1, self._output_dim)).float() * 10))
-        self._optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
+        self._optimizer: torch.optim.Adam = torch.optim.Adam(self.parameters(), lr=learning_rate)
         self.apply(init_weights)
-        self._device = device
-        self.scaler = StandardScaler(self._device)
+        self._device: torch.device = device
+        self.scaler: StandardScaler = StandardScaler(self._device)
 
     def forward(
         self,
@@ -234,16 +318,18 @@ def forward(
 
         Args:
             data (torch.Tensor): Input data.
-            ret_log_var (bool): Whether to return the log variance.
+            ret_log_var (bool, optional): Whether to return the log variance, defaults to False.
 
         Returns:
-            mean (torch.Tensor): Mean of the next state, reward, cost.
-            logvar or var (torch.Tensor): Log variance of the next state, reward, cost.
+            mean: Mean of the next state, reward, cost.
+            logvar or var: Log variance of the next state, reward, cost.
         """
-        if not torch.is_tensor(data):
-            data = torch.tensor(data, dtype=torch.float32, device=self._device)
-        data = self.scaler.transform(data)
-        nn1_output = swish(self._nn1(data))
+        if isinstance(data, torch.Tensor):
+            data_t = data
+        else:
+            data_t = torch.tensor(data, dtype=torch.float32, device=self._device)
+        data_t = self.scaler.transform(data_t)
+        nn1_output = swish(self._nn1(data_t))
         nn2_output = swish(self._nn2(nn1_output))
         nn3_output = swish(self._nn3(nn2_output))
         nn4_output = swish(self._nn4(nn3_output))
@@ -267,20 +353,22 @@ def forward_idx(
         """Compute next state, reward, cost from an certain model.
 
         Args:
-            data (torch.Tensor): Input data.
+            data (torch.Tensor | np.ndarray): Input data.
             idx_model (int): Index of the model.
             ret_log_var (bool): Whether to return the log variance.
 
         Returns:
-            mean (torch.Tensor): Mean of the next state, reward, cost.
-            logvar or var (torch.Tensor): Log variance of the next state, reward, cost.
+            mean: Mean of the next state, reward, cost.
+            logvar or var: Log variance of the next state, reward, cost.
         """
         assert data.shape[0] == 1
-        if not torch.is_tensor(data):
-            data = torch.tensor(data, dtype=torch.float32, device=self._device)
-        data = self.scaler.transform(data[0])
+        if isinstance(data, torch.Tensor):
+            data_t = data
+        else:
+            data_t = torch.tensor(data, dtype=torch.float32, device=self._device)
+        data_t = self.scaler.transform(data_t[0])
         unbatched_forward_fn = partial(unbatched_forward, index=idx_model)
-        nn1_output = swish(unbatched_forward_fn(self._nn1, data))
+        nn1_output = swish(unbatched_forward_fn(self._nn1, data_t))
         nn2_output = swish(unbatched_forward_fn(self._nn2, nn1_output))
         nn3_output = swish(unbatched_forward_fn(self._nn3, nn2_output))
         nn4_output = swish(unbatched_forward_fn(self._nn4, nn3_output))
@@ -315,7 +403,7 @@ def loss(
             mean (torch.Tensor): Mean of the next state, reward, cost.
             logvar (torch.Tensor): Log variance of the next state, reward, cost.
             labels (torch.Tensor): Ground truth of the next state, reward, cost.
-            inc_var_loss (bool): Whether to include the variance loss.
+            inc_var_loss (bool, optional): Whether to include the variance loss. Defaults to True.
 
         Returns:
             total_loss (torch.Tensor): Total loss.
@@ -334,7 +422,11 @@ def loss(
         return total_loss, mse_loss
 
     def train_ensemble(self, loss: torch.Tensor) -> None:
-        """Train the dynamics model."""
+        """Train the dynamics model.
+
+        Args:
+            loss (torch.Tensor): The loss of the dynamics model.
+        """
         self._optimizer.zero_grad()
         loss += 0.01 * torch.sum(torch.Tensor(self.max_logvar)) - 0.01 * torch.sum(
             torch.Tensor(self.min_logvar),
@@ -347,7 +439,24 @@ def train_ensemble(self, loss: torch.Tensor) -> None:
 
 # pylint: disable-next=too-many-instance-attributes
 class EnsembleDynamicsModel:
-    """Dynamics model for predict next state, reward and cost."""
+    """Dynamics model for predict next state, reward and cost.
+
+    Args:
+        model_cfgs (Config): The configuration of the dynamics model.
+        device (torch.device): The device to use.
+        state_shape (tuple[int, ...]): The shape of the state.
+        action_shape (tuple[int, ...]): The shape of the action.
+        actor_critic (ConstraintActorCritic | ConstraintActorQCritic | None, optional): The actor critic model.
+            Defaults to None.
+        rew_func (Callable[[torch.Tensor], torch.Tensor] | None, optional): The reward function. Defaults to None.
+        cost_func (Callable[[torch.Tensor], torch.Tensor] | None, optional): The cost function.
+            Defaults to None.
+        terminal_func (Callable[[torch.Tensor], torch.Tensor] | None, optional): The terminal function.
+            Defaults to None.
+
+    Attributes:
+        elite_model_idxes (list[int]): The index of the elite models.
+    """
 
     # pylint: disable-next=too-many-arguments
     def __init__(
@@ -357,41 +466,41 @@ def __init__(
         state_shape: tuple[int, ...],
         action_shape: tuple[int, ...],
         actor_critic: ConstraintActorCritic | ConstraintActorQCritic | None = None,
-        rew_func: Callable | None = None,
-        cost_func: Callable | None = None,
-        terminal_func: Callable | None = None,
+        rew_func: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        cost_func: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        terminal_func: Callable[[torch.Tensor], torch.Tensor] | None = None,
     ) -> None:
         """Initialize the dynamics model."""
-        self._num_ensemble = model_cfgs.num_ensemble
-        self._elite_size = model_cfgs.elite_size
-        self._predict_reward = model_cfgs.predict_reward
-        self._predict_cost = model_cfgs.predict_cost
-        self._batch_size = model_cfgs.batch_size
-        self._max_epoch_since_update = model_cfgs.max_epoch
+        self._num_ensemble: int = model_cfgs.num_ensemble
+        self._elite_size: int = model_cfgs.elite_size
+        self._predict_reward: bool = model_cfgs.predict_reward
+        self._predict_cost: bool = model_cfgs.predict_cost
+        self._batch_size: int = model_cfgs.batch_size
+        self._max_epoch_since_update: int = model_cfgs.max_epoch
 
-        self._reward_size = model_cfgs.reward_size
-        self._cost_size = model_cfgs.cost_size
-        self._use_cost = model_cfgs.use_cost
-        self._use_terminal = model_cfgs.use_terminal
-        self._use_var = model_cfgs.use_var
-        self._use_reward_critic = model_cfgs.use_reward_critic
-        self._use_cost_critic = model_cfgs.use_cost_critic
+        self._reward_size: int = model_cfgs.reward_size
+        self._cost_size: int = model_cfgs.cost_size
+        self._use_cost: bool = model_cfgs.use_cost
+        self._use_terminal: bool = model_cfgs.use_terminal
+        self._use_var: bool = model_cfgs.use_var
+        self._use_reward_critic: bool = model_cfgs.use_reward_critic
+        self._use_cost_critic: bool = model_cfgs.use_cost_critic
 
-        self._device = device
+        self._device: torch.device = device
 
-        self._state_size = state_shape[0]
-        self._action_size = action_shape[0]
+        self._state_size: int = state_shape[0]
+        self._action_size: int = action_shape[0]
 
-        self._rew_func = rew_func
-        self._cost_func = cost_func
-        self._terminal_func = terminal_func
+        self._rew_func: Callable[[torch.Tensor], torch.Tensor] | None = rew_func
+        self._cost_func: Callable[[torch.Tensor], torch.Tensor] | None = cost_func
+        self._terminal_func: Callable[[torch.Tensor], torch.Tensor] | None = terminal_func
 
-        self._actor_critic = actor_critic
+        self._actor_critic: ConstraintActorCritic | ConstraintActorQCritic | None = actor_critic
 
         self._model_list: list[int] = []
 
-        self.elite_model_idxes = list(range(self._elite_size))
-        self._ensemble_model = EnsembleModel(
+        self.elite_model_idxes: list[int] = list(range(self._elite_size))
+        self._ensemble_model: EnsembleModel = EnsembleModel(
             device=self._device,
             state_size=self._state_size,
             action_size=self._action_size,
@@ -405,10 +514,11 @@ def __init__(
             use_decay=model_cfgs.use_decay,
         )
         self._ensemble_model.to(self._device)
-
         self._max_epoch_since_update = 5
-        self._epochs_since_update = 0
-        self._snapshots = {i: (0, 1e10) for i in range(self._num_ensemble)}
+        self._epochs_since_update: int = 0
+        self._snapshots: dict[int, tuple[int, float]] = {
+            i: (0, 1e10) for i in range(self._num_ensemble)
+        }
 
         if self._predict_reward is False:
             assert rew_func is not None, 'rew_func should not be None'
@@ -427,17 +537,17 @@ def __init__(
 
     @property
     def ensemble_model(self) -> EnsembleModel:
-        """Return the ensemble model."""
+        """The ensemble model."""
         return self._ensemble_model
 
     @property
     def num_models(self) -> int:
-        """Return the number of ensemble."""
+        """The number of ensemble."""
         return self._num_ensemble
 
     @property
     def state_size(self) -> int:
-        """Return the state size."""
+        """The state size."""
         return self._state_size
 
     # pylint: disable-next=too-many-locals, too-many-arguments
@@ -453,6 +563,10 @@ def train(
             inputs (np.ndarray): Input data.
             labels (np.ndarray): Ground truth of the next state, reward, cost.
             holdout_ratio (float): The ratio of the data hold out for validation.
+
+        Returns:
+            train_mse_losses: The training loss.
+            val_mse_losses: The validation loss.
         """
         self._epochs_since_update = 0
         self._snapshots = {i: (0, 1e10) for i in range(self._num_ensemble)}
@@ -466,8 +580,9 @@ def train(
         holdout_inputs, holdout_labels = inputs[:num_holdout], labels[:num_holdout]
         self._ensemble_model.scaler.fit(train_inputs)
 
+        train_mse_losses = []
+        val_losses = []
         for epoch in itertools.count():
-            train_mse_losses = []
             # training
             train_idx = np.vstack(
                 [np.random.permutation(train_inputs.shape[0]) for _ in range(self._num_ensemble)],
@@ -489,9 +604,7 @@ def train(
             )
             val_batch_size = 512
             val_losses_list = []
-            for _len_valid, start_pos in enumerate(
-                range(0, holdout_inputs.shape[0], val_batch_size),
-            ):
+            for start_pos in range(0, holdout_inputs.shape[0], val_batch_size):
                 with torch.no_grad():
                     idx = val_idx[:, start_pos : start_pos + val_batch_size]
                     val_input = torch.from_numpy(holdout_inputs[idx]).float().to(self._device)
@@ -508,11 +621,13 @@ def train(
                     )
                     holdout_mse_losses = holdout_mse_losses.detach().cpu().numpy()
                     val_losses_list.append(holdout_mse_losses)
-            val_losses = np.array(val_losses_list)
-            val_losses = np.sum(val_losses, axis=0) / (_len_valid + 1)
-            sorted_loss_idx = np.argsort(val_losses)
+            current_loss = np.sum(np.array(val_losses_list), axis=0) / (
+                int(holdout_inputs.shape[0] / val_batch_size) + 1
+            )
+            val_losses.append(current_loss)
+            sorted_loss_idx = np.argsort(current_loss)
             self.elite_model_idxes = sorted_loss_idx[: self._elite_size].tolist()
-            break_train = self._save_best(epoch, val_losses)
+            break_train = self._save_best(epoch, current_loss)
             if break_train:
                 break
 
@@ -521,7 +636,15 @@ def train(
         return np.array(train_mse_losses), np.array(val_mse_losses)
 
     def _save_best(self, epoch: int, holdout_losses: list) -> bool:
-        """Save the best model."""
+        """Save the best model.
+
+        Args:
+            epoch (int): The current epoch.
+            holdout_losses (list): The holdout loss.
+
+        Returns:
+            Whether to break the training.
+        """
         updated = False
         for i, current_loss in enumerate(holdout_losses):
             _, best = self._snapshots[i]
@@ -541,7 +664,17 @@ def _compute_reward(
         self,
         network_output: torch.Tensor,
     ) -> torch.Tensor:
-        """Compute the reward from the network output."""
+        """Compute the reward from the network output.
+
+        Args:
+            network_output (torch.Tensor): The output of the network.
+
+        Returns:
+            reward: The reward, from the network output or the reward function.
+
+        Raises:
+            ValueError: If the reward function is not defined.
+        """
         if self._predict_reward:
             reward_start_dim = int(self._predict_cost) * self._cost_size
             reward_end_dim = reward_start_dim + self._reward_size
@@ -555,7 +688,17 @@ def _compute_cost(
         self,
         network_output: torch.Tensor,
     ) -> torch.Tensor:
-        """Compute the cost from the network output."""
+        """Compute the cost from the network output.
+
+        Args:
+            network_output (torch.Tensor): The output of the network.
+
+        Returns:
+            cost: The cost, from the network output or the cost function.
+
+        Raises:
+            ValueError: If the cost function is not defined.
+        """
         if self._predict_cost:
             return network_output[:, :, : self._cost_size]
         if self._cost_func is not None:
@@ -567,7 +710,17 @@ def _compute_terminal(
         self,
         network_output: torch.Tensor,
     ) -> torch.Tensor:
-        """Compute the terminal from the network output."""
+        """Compute the terminal from the network output.
+
+        Args:
+            network_output (torch.Tensor): The output of the network.
+
+        Returns:
+            terminal: The terminal signal, from the network output or the terminal function.
+
+        Raises:
+            ValueError: If the terminal function is not defined.
+        """
         if self._terminal_func is not None:
             return self._terminal_func(network_output[:, :, self._state_start_dim :])
         raise ValueError('Terminal function is not defined.')
@@ -586,6 +739,10 @@ def _predict(
             batch_size (int, optional): the batch size for prediction.
             idx (Union[int, None], optional): the index of the model to use.
             ret_log_var (bool, optional): whether to return the log variance.
+
+        Returns:
+            ensemble_mean_tensor: The mean of the ensemble.
+            ensemble_var_tensor: The variance of the ensemble.
         """
         if idx is not None:
             assert inputs.shape[0] == 1
@@ -639,7 +796,7 @@ def sample(
         Returns:
             sample_states (torch.Tensor): the sampled states.
             rewards (torch.Tensor): the rewards.
-            info (Dict[str, List[torch.Tensor]]): the info dict, contains the costs if use_cost is True
+            info: the info dict, contains the costs if `use_cost` is True.
         """
         assert (
             states.shape[:-1] == actions.shape[:-1]
@@ -709,7 +866,7 @@ def imagine(
             idx (int, optional): the index of the model to use.
 
         Returns:
-            traj (Dict[str, List[torch.Tensor]]): the trajectory dict, contains the states, rewards, etc.
+            traj: the trajectory dict, contains the states, rewards, etc.
         """
         assert (
             states.shape[1] == self._state_size
diff --git a/omnisafe/algorithms/model_based/base/loop.py b/omnisafe/algorithms/model_based/base/loop.py
index 4b78509d6..23f5321b3 100644
--- a/omnisafe/algorithms/model_based/base/loop.py
+++ b/omnisafe/algorithms/model_based/base/loop.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 """Implementation of the Learning Off-Policy with Online Planning algorithm."""
 
+
 from __future__ import annotations
 
 from typing import Any
@@ -23,13 +24,13 @@
 from torch import nn, optim
 from torch.nn.utils.clip_grad import clip_grad_norm_
 
-from omnisafe.adapter import ModelBasedAdapter
 from omnisafe.algorithms import registry
 from omnisafe.algorithms.model_based.base.ensemble import EnsembleDynamicsModel
 from omnisafe.algorithms.model_based.base.pets import PETS
 from omnisafe.algorithms.model_based.planner.arc import ARCPlanner
 from omnisafe.common.buffer import OffPolicyBuffer
 from omnisafe.models.actor_critic.constraint_actor_q_critic import ConstraintActorQCritic
+from omnisafe.typing import OmnisafeSpace
 
 
 @registry.register
@@ -43,13 +44,28 @@ class LOOP(PETS):
         - URL: `LOOP <https://arxiv.org/abs/2008.10066>`_
     """
 
+    _log_alpha: torch.Tensor
+    _alpha_optimizer: optim.Optimizer
+    _target_entropy: float
+
     def _init_model(self) -> None:
-        """Initialize the dynamics model and the planner."""
-        self._dynamics_state_space = (
+        """Initialize the dynamics model and the planner.
+
+        LOOP uses following models:
+
+        - dynamics model: to predict the next state and the cost.
+        - actor_critic: to predict the action and the value.
+        - planner: to generate the action.
+        """
+        self._dynamics_state_space: OmnisafeSpace = (
             self._env.coordinate_observation_space
             if self._env.coordinate_observation_space is not None
             else self._env.observation_space
         )
+        assert self._dynamics_state_space is not None and isinstance(
+            self._dynamics_state_space.shape,
+            tuple,
+        )
         assert self._env.action_space is not None and isinstance(
             self._env.action_space.shape,
             tuple,
@@ -58,15 +74,15 @@ def _init_model(self) -> None:
             self._action_space = self._env.action_space
         else:
             raise NotImplementedError
-        self._actor_critic = ConstraintActorQCritic(
+        self._actor_critic: ConstraintActorQCritic = ConstraintActorQCritic(
             obs_space=self._dynamics_state_space,
             act_space=self._env.action_space,
             model_cfgs=self._cfgs.model_cfgs,
             epochs=self._epochs,
         ).to(self._device)
-        self._use_actor_critic = True
-        self._update_count = 0
-        self._dynamics = EnsembleDynamicsModel(
+        self._use_actor_critic: bool = True
+        self._update_count: int = 0
+        self._dynamics: EnsembleDynamicsModel = EnsembleDynamicsModel(
             model_cfgs=self._cfgs.dynamics_cfgs,
             device=self._device,
             state_shape=self._dynamics_state_space.shape,
@@ -77,7 +93,7 @@ def _init_model(self) -> None:
             terminal_func=None,
         )
         self._update_dynamics_cycle = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
-        self._planner = ARCPlanner(
+        self._planner: ARCPlanner = ARCPlanner(
             dynamics=self._dynamics,
             planner_cfgs=self._cfgs.planner_cfgs,
             gamma=float(self._cfgs.algo_cfgs.gamma),
@@ -91,10 +107,17 @@ def _init_model(self) -> None:
         )
 
     def _init(self) -> None:
+        """The initialization of the algorithm.
+
+        User can define the initialization of the algorithm by inheriting this method.
+
+        Examples:
+            >>> def _init(self) -> None:
+            ...     super()._init()
+            ...     self._buffer = CustomBuffer()
+            ...     self._model = CustomModel()
+        """
         super()._init()
-        self._log_alpha: torch.Tensor
-        self._alpha_optimizer: optim.Optimizer
-        self._target_entropy: float
 
         self._alpha = self._cfgs.algo_cfgs.alpha
         self._alpha_gamma = self._cfgs.algo_cfgs.alpha_gamma
@@ -111,7 +134,24 @@ def _alpha_discount(self) -> None:
         self._alpha *= self._alpha_gamma
 
     def _init_log(self) -> None:
-        """Initialize logger."""
+        """Initialize logger.
+
+        +-------------------------+----------------------------------------------------------------------+
+        | Things to log           | Description                                                          |
+        +=========================+======================================================================+
+        | Value/alpha             | The value of alpha.                                                  |
+        +-------------------------+----------------------------------------------------------------------+
+        | Values/reward_critic    | Average value in :meth:`rollout` (from critic network) of the epoch. |
+        +-------------------------+----------------------------------------------------------------------+
+        | Values/cost_critic      | Average cost in :meth:`rollout` (from critic network) of the epoch.  |
+        +-------------------------+----------------------------------------------------------------------+
+        | Loss/Loss_cost_critic   | Loss of the cost critic network.                                     |
+        +-------------------------+----------------------------------------------------------------------+
+        | Loss/Loss_reward_critic | Loss of the cost critic network.                                     |
+        +-------------------------+----------------------------------------------------------------------+
+        | Loss/Loss_pi            | Loss of the policy network.                                          |
+        +-------------------------+----------------------------------------------------------------------+
+        """
         super()._init_log()
         self._logger.register_key('Value/alpha')
         # log information about actor
@@ -144,18 +184,15 @@ def _select_action(  # pylint: disable=unused-argument
         self,
         current_step: int,
         state: torch.Tensor,
-        env: ModelBasedAdapter,
-    ) -> tuple[torch.Tensor, dict]:
+    ) -> torch.Tensor:
         """Select action.
 
         Args:
-            current_step (int): current step
-            state (torch.Tensor): current state
-            env (ModelBasedAdapter): environment
+            current_step (int): The current step.
+            state (torch.Tensor): The current state.
 
         Returns:
-            action (torch.Tensor): action
-            action_info (dict): action information
+            The selected action.
         """
         if current_step < self._cfgs.algo_cfgs.start_learning_steps:
             action = torch.tensor(self._env.action_space.sample()).to(self._device).unsqueeze(0)
@@ -166,14 +203,42 @@ def _select_action(  # pylint: disable=unused-argument
         assert action.shape == torch.Size(
             [1, *self._action_space.shape],
         ), 'action shape should be [batch_size, action_dim]'
-        info = {}
-        return action, info
+        return action
 
     def _update_policy(self, current_step: int) -> None:
         """Update policy.
 
+        -  Get the ``data`` from buffer
+
+        .. note::
+
+            +----------+---------------------------------------+
+            | obs      | ``observaion`` stored in buffer.      |
+            +==========+=======================================+
+            | act      | ``action`` stored in buffer.          |
+            +----------+---------------------------------------+
+            | reward   | ``reward`` stored in buffer.          |
+            +----------+---------------------------------------+
+            | cost     | ``cost`` stored in buffer.            |
+            +----------+---------------------------------------+
+            | next_obs | ``next observaion`` stored in buffer. |
+            +----------+---------------------------------------+
+            | done     | ``terminated`` stored in buffer.      |
+            +----------+---------------------------------------+
+
+        -  Update value net by :meth:`_update_reward_critic`.
+        -  Update cost net by :meth:`_update_cost_critic`.
+        -  Update policy net by :meth:`_update_actor`.
+
+        The basic process of each update is as follows:
+
+        #. Get the mini-batch data from buffer.
+        #. Get the loss of network.
+        #. Update the network by loss.
+        #. Repeat steps 2, 3 until the ``update_policy_iters`` times.
+
         Args:
-            current_step (int): current step
+            current_step (int): The current step.
         """
         if current_step >= self._cfgs.algo_cfgs.start_learning_steps:
             for _step in range(self._cfgs.algo_cfgs.update_policy_iters):
@@ -218,8 +283,6 @@ def _update_policy(self, current_step: int) -> None:
 
     def _store_real_data(  # pylint: disable=too-many-arguments,unused-argument
         self,
-        current_step: int,
-        ep_len: int,
         state: torch.Tensor,
         action: torch.Tensor,
         reward: torch.Tensor,
@@ -227,23 +290,19 @@ def _store_real_data(  # pylint: disable=too-many-arguments,unused-argument
         terminated: torch.Tensor,
         truncated: torch.Tensor,
         next_state: torch.Tensor,
-        info: dict,
-        action_info: dict,
+        info: dict[str, Any],
     ) -> None:  # pylint: disable=too-many-arguments
         """Store real data in buffer.
 
         Args:
-            current_step (int): current step
-            ep_len (int): episode length
-            state (torch.Tensor): current state
-            action (torch.Tensor): action
-            reward (torch.Tensor): reward
-            cost (torch.Tensor): cost
-            terminated (torch.Tensor): terminated
-            truncated (torch.Tensor): truncated
-            next_state (torch.Tensor): next state
-            info (dict): information
-            action_info (dict): action information
+            state (torch.Tensor): The state from the environment.
+            action (torch.Tensor): The action from the agent.
+            reward (torch.Tensor): The reward signal from the environment.
+            cost (torch.Tensor): The cost signal from the environment.
+            terminated (torch.Tensor): The terminated signal from the environment.
+            truncated (torch.Tensor): The truncated signal from the environment.
+            next_state (torch.Tensor): The next state from the environment.
+            info (dict[str, Any]): The information from the environment.
         """
         done = terminated or truncated
         goal_met = False if 'goal_met' not in info.keys() else info['goal_met']
@@ -279,12 +338,16 @@ def _update_reward_critic(
     ) -> None:
         """Update reward critic using Soft Actor-Critic.
 
+        - Get the TD loss of reward critic.
+        - Update critic network by loss.
+        - Log useful information.
+
         Args:
-            obs (torch.Tensor): observation
-            action (torch.Tensor): action
-            reward (torch.Tensor): reward
-            done (torch.Tensor): done
-            next_obs (torch.Tensor): next observation
+            obs (torch.Tensor): The ``observation`` sampled from buffer.
+            action (torch.Tensor): The ``action`` sampled from buffer.
+            reward (torch.Tensor): The ``reward`` sampled from buffer.
+            done (torch.Tensor): The ``terminated`` sampled from buffer.
+            next_obs (torch.Tensor): The ``next observation`` sampled from buffer.
         """
         self._actor_critic.reward_critic_optimizer.zero_grad()
 
@@ -331,12 +394,16 @@ def _update_cost_critic(
     ) -> None:
         """Update cost critic using TD3 algorithm.
 
+        - Get the TD loss of cost critic.
+        - Update critic network by loss.
+        - Log useful information.
+
         Args:
-            obs (torch.Tensor): current observation
-            action (torch.Tensor): current action
-            cost (torch.Tensor): current cost
-            done (torch.Tensor): current done signal
-            next_obs (torch.Tensor): next observation
+            obs (torch.Tensor): The ``observation`` sampled from buffer.
+            action (torch.Tensor): The ``action`` sampled from buffer.
+            cost (torch.Tensor): The ``cost`` sampled from buffer.
+            done (torch.Tensor): The ``terminated`` sampled from buffer.
+            next_obs (torch.Tensor): The ``next observation`` sampled from buffer.
         """
         with torch.no_grad():
             next_action = self._actor_critic.actor.predict(next_obs, deterministic=True)
@@ -372,8 +439,12 @@ def _update_actor(
     ) -> None:
         """Update actor using Soft Actor-Critic algorithm.
 
+        - Get the loss of actor.
+        - Update actor by loss.
+        - Log useful information.
+
         Args:
-            obs (torch.Tensor): observation
+            obs (torch.Tensor): The ``observation`` sampled from buffer.
         """
         self._actor_critic.actor_optimizer.zero_grad()
         loss = self._loss_pi(obs)
@@ -400,10 +471,22 @@ def _loss_pi(
         self,
         obs: torch.Tensor,
     ) -> torch.Tensor:
-        """Compute loss for actor using Soft Actor-Critic algorithm.
+        r"""Computing ``pi/actor`` loss.
+
+        The loss function in SAC is defined as:
+
+        .. math::
+
+            L = -Q^V (s, \pi (s)) + \alpha \log \pi (s)
+
+        where :math:`Q^V` is the min value of two reward critic networks, and :math:`\pi` is the
+        policy network, and :math:`\alpha` is the temperature parameter.
 
         Args:
-            obs (torch.Tensor): observation
+            obs (torch.Tensor): The ``observation`` sampled from buffer.
+
+        Returns:
+            The loss of pi/actor.
         """
         action = self._actor_critic.actor.predict(
             obs,
diff --git a/omnisafe/algorithms/model_based/base/pets.py b/omnisafe/algorithms/model_based/base/pets.py
index a0c80b39c..de3994aae 100644
--- a/omnisafe/algorithms/model_based/base/pets.py
+++ b/omnisafe/algorithms/model_based/base/pets.py
@@ -13,11 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 """Implementation of the Probabilistic Ensembles with Trajectory Sampling algorithm."""
+
+
 from __future__ import annotations
 
 import os
 import time
-from typing import Any
+from typing import Any, Callable
 
 import numpy as np
 import torch
@@ -32,6 +34,7 @@
 from omnisafe.algorithms.model_based.planner.cem import CEMPlanner
 from omnisafe.common.buffer import OffPolicyBuffer
 from omnisafe.common.logger import Logger
+from omnisafe.typing import OmnisafeSpace
 
 
 @registry.register
@@ -46,29 +49,36 @@ class PETS(BaseAlgo):
     """
 
     def _init_env(self) -> None:
-        self._env = ModelBasedAdapter(
+        self._env: ModelBasedAdapter = ModelBasedAdapter(
             self._env_id,
             1,
             self._seed,
             self._cfgs,
         )
-        self._total_steps = int(self._cfgs.train_cfgs.total_steps)
-        self._steps_per_epoch = int(self._cfgs.algo_cfgs.steps_per_epoch)
-        self._epochs = self._total_steps // self._cfgs.algo_cfgs.steps_per_epoch
-        print(f'Total steps: {self._total_steps}, epochs: {self._epochs}')
+        self._total_steps: int = int(self._cfgs.train_cfgs.total_steps)
+        self._steps_per_epoch: int = int(self._cfgs.algo_cfgs.steps_per_epoch)
+        self._epochs: int = self._total_steps // self._cfgs.algo_cfgs.steps_per_epoch
 
     def _init_model(self) -> None:
         """Initialize dynamics model and planner."""
-        self._dynamics_state_space = (
+        self._dynamics_state_space: OmnisafeSpace = (
             self._env.coordinate_observation_space
             if self._env.coordinate_observation_space is not None
             else self._env.observation_space
         )
+        assert self._dynamics_state_space is not None and isinstance(
+            self._dynamics_state_space.shape,
+            tuple,
+        )
+        assert self._env.action_space is not None and isinstance(
+            self._env.action_space.shape,
+            tuple,
+        )
         if isinstance(self._env.action_space, Box):
             self._action_space = self._env.action_space
         else:
             raise NotImplementedError
-        self._dynamics = EnsembleDynamicsModel(
+        self._dynamics: EnsembleDynamicsModel = EnsembleDynamicsModel(
             model_cfgs=self._cfgs.dynamics_cfgs,
             device=self._device,
             state_shape=self._dynamics_state_space.shape,
@@ -79,7 +89,7 @@ def _init_model(self) -> None:
             terminal_func=None,
         )
 
-        self._planner = CEMPlanner(
+        self._planner: CEMPlanner = CEMPlanner(
             dynamics=self._dynamics,
             planner_cfgs=self._cfgs.planner_cfgs,
             gamma=float(self._cfgs.algo_cfgs.gamma),
@@ -90,12 +100,12 @@ def _init_model(self) -> None:
             action_min=-1.0,
             device=self._device,
         )
-        self._use_actor_critic = False
-        self._update_dynamics_cycle = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
+        self._use_actor_critic: bool = False
+        self._update_dynamics_cycle: int = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
 
     def _init(self) -> None:
         """Initialize the algorithm."""
-        self._dynamics_buf = OffPolicyBuffer(
+        self._dynamics_buf: OffPolicyBuffer = OffPolicyBuffer(
             obs_space=self._dynamics_state_space,
             act_space=self._env.action_space,
             size=self._cfgs.train_cfgs.total_steps,
@@ -106,17 +116,71 @@ def _init(self) -> None:
             'render_mode': 'rgb_array',
             'camera_name': 'track',
         }
-        self._eval_env = ModelBasedAdapter(
+        self._eval_env: ModelBasedAdapter = ModelBasedAdapter(
             self._env_id,
             1,
             self._seed,
             self._cfgs,
             **env_kwargs,
         )
-        self._eval_fn = self._evaluation_single_step
+        self._eval_fn: Callable[[int, bool], None] = self._evaluation_single_step
 
     def _init_log(self) -> None:
-        """Initialize logger."""
+        """Initialize logger.
+
+        +---------------------------+-------------------------------------------------+
+        | Things to log             | Description                                     |
+        +===========================+=================================================+
+        | Train/Epoch               | Current epoch.                                  |
+        +---------------------------+-------------------------------------------------+
+        | TotalEnvSteps             | Total steps of the experiment.                  |
+        +---------------------------+-------------------------------------------------+
+        | Metrics/EpRet             | Average return of the epoch.                    |
+        +---------------------------+-------------------------------------------------+
+        | Metrics/EpCost            | Average cost of the epoch.                      |
+        +---------------------------+-------------------------------------------------+
+        | Metrics/EpLen             | Average length of the epoch.                    |
+        +---------------------------+-------------------------------------------------+
+        | EvalMetrics/EpRet         | Average episode return in evaluation.           |
+        +---------------------------+-------------------------------------------------+
+        | EvalMetrics/EpCost        | Average episode cost in evaluation.             |
+        +---------------------------+-------------------------------------------------+
+        | EvalMetrics/EpLen         | Average episode length in evaluation.           |
+        +---------------------------+-------------------------------------------------+
+        | Loss/DynamicsTrainMseLoss | The training loss of dynamics model.            |
+        +---------------------------+-------------------------------------------------+
+        | Loss/DynamicsValMseLoss   | The validation loss of dynamics model.          |
+        +---------------------------+-------------------------------------------------+
+        | Plan/iter                 | The number of iterations in the planner.        |
+        +---------------------------+-------------------------------------------------+
+        | Plan/last_var_mean        | The mean of the last variance in the planner.   |
+        +---------------------------+-------------------------------------------------+
+        | Plan/last_var_max         | The max of the last variance in the planner.    |
+        +---------------------------+-------------------------------------------------+
+        | Plan/last_var_min         | The min of the last variance in the planner.    |
+        +---------------------------+-------------------------------------------------+
+        | Plan/episode_returns_max  | The max of the episode returns in the planner.  |
+        +---------------------------+-------------------------------------------------+
+        | Plan/episode_returns_mean | The mean of the episode returns in the planner. |
+        +---------------------------+-------------------------------------------------+
+        | Plan/episode_returns_min  | The min of the episode returns in the planner.  |
+        +---------------------------+-------------------------------------------------+
+        | Time/Total                | The total time of the algorithm.                |
+        +---------------------------+-------------------------------------------------+
+        | Time/Rollout              | The time of the rollout.                        |
+        +---------------------------+-------------------------------------------------+
+        | Time/UpdateActorCritic    | The time of the actor-critic update.            |
+        +---------------------------+-------------------------------------------------+
+        | Time/Eval                 | The time of the evaluation.                     |
+        +---------------------------+-------------------------------------------------+
+        | Time/Epoch                | The time of the epoch.                          |
+        +---------------------------+-------------------------------------------------+
+        | Time/FPS                  | The FPS of the algorithm.                       |
+        +---------------------------+-------------------------------------------------+
+        | Time/UpdateDynamics       | The time of the dynamics update.                |
+        +---------------------------+-------------------------------------------------+
+
+        """
         self._logger = Logger(
             output_dir=self._cfgs.logger_cfgs.log_dir,
             exp_name=self._cfgs.exp_name,
@@ -159,9 +223,8 @@ def _init_log(self) -> None:
 
     def _save_model(self) -> None:
         """Save the model."""
-        what_to_save: dict[str, Any] = {}
         # set up model saving
-        what_to_save = {
+        what_to_save: dict[str, Any] = {
             'dynamics': self._dynamics.ensemble_model,
         }
         if self._cfgs.algo_cfgs.obs_normalize:
@@ -174,22 +237,23 @@ def learn(self) -> tuple[float, float, float]:
         """This is main function for algorithm update.
 
         It is divided into the following steps:
+
         - :meth:`rollout`: collect interactive data from environment.
         - :meth:`update`: perform actor/critic updates.
         - :meth:`log`: epoch/update information for visualization and terminal log print.
 
         Returns:
-            ep_ret: average episode return in final epoch.
-            ep_cost: average episode cost in final epoch.
-            ep_len: average episode length in final epoch.
+            ep_ret: Average episode return in final epoch.
+            ep_cost: Average episode cost in final epoch.
+            ep_len: Average episode length in final epoch.
         """
         self._logger.log('INFO: Start training')
         start_time = time.time()
         current_step = 0
         for epoch in range(self._epochs):
-            current_step = self._env.roll_out(
+            current_step = self._env.rollout(
                 current_step=current_step,
-                roll_out_step=self._steps_per_epoch,
+                rollout_step=self._steps_per_epoch,
                 use_actor_critic=self._use_actor_critic,
                 act_func=self._select_action,
                 store_data_func=self._store_real_data,
@@ -206,7 +270,7 @@ def learn(self) -> tuple[float, float, float]:
                 self._update_epoch()
             # evaluate episode
             self._logger.store(
-                **{
+                {
                     'Train/Epoch': epoch,
                     'TotalEnvSteps': current_step,
                     'Time/Total': time.time() - start_time,
@@ -226,45 +290,36 @@ def learn(self) -> tuple[float, float, float]:
 
     def _algo_reset(
         self,
-        current_step: int,  # pylint: disable=unused-argument
     ) -> None:
-        ...
+        """Reset the algorithm."""
 
     def _update_policy(
         self,
         current_step: int,  # pylint: disable=unused-argument
     ) -> None:
-        ...
+        """Update policy."""
 
     def _update_dynamics_model(
         self,
-        current_step: int,  # pylint: disable=unused-argument
     ) -> None:
-        """Update dynamics.
-
-        Args:
-            current_step (int): current step.
-        """
+        """Update dynamics model."""
         state = self._dynamics_buf.data['obs'][: self._dynamics_buf.size, :]
         action = self._dynamics_buf.data['act'][: self._dynamics_buf.size, :]
         reward = self._dynamics_buf.data['reward'][: self._dynamics_buf.size]
         cost = self._dynamics_buf.data['cost'][: self._dynamics_buf.size]
         next_state = self._dynamics_buf.data['next_obs'][: self._dynamics_buf.size, :]
         delta_state = next_state - state
-        if torch.is_tensor(delta_state):
-            inputs = torch.cat((state, action), -1)
-            inputs = torch.reshape(inputs, (inputs.shape[0], -1))
-
-            labels = torch.reshape(delta_state, (delta_state.shape[0], -1))
-            if self._cfgs.dynamics_cfgs.predict_reward:
-                labels = torch.cat(((torch.reshape(reward, (reward.shape[0], -1))), labels), -1)
-            if self._cfgs.dynamics_cfgs.predict_cost:
-                labels = torch.cat(((torch.reshape(cost, (cost.shape[0], -1))), labels), -1)
-            inputs = inputs.cpu().detach().numpy()
-            labels = labels.cpu().detach().numpy()
-        assert not torch.is_tensor(inputs) and not torch.is_tensor(
-            labels,
-        ), 'inputs and labels should be numpy array'
+        assert isinstance(delta_state, torch.Tensor), 'delta_state should be torch.Tensor'
+        inputs = torch.cat((state, action), -1)
+        inputs = torch.reshape(inputs, (inputs.shape[0], -1))
+
+        labels = torch.reshape(delta_state, (delta_state.shape[0], -1))
+        if self._cfgs.dynamics_cfgs.predict_reward:
+            labels = torch.cat(((torch.reshape(reward, (reward.shape[0], -1))), labels), -1)
+        if self._cfgs.dynamics_cfgs.predict_cost:
+            labels = torch.cat(((torch.reshape(cost, (cost.shape[0], -1))), labels), -1)
+        inputs = inputs.cpu().detach().numpy()
+        labels = labels.cpu().detach().numpy()
         train_mse_losses, val_mse_losses = self._dynamics.train(
             inputs,
             labels,
@@ -284,18 +339,15 @@ def _select_action(  # pylint: disable=unused-argument
         self,
         current_step: int,
         state: torch.Tensor,
-        env: ModelBasedAdapter,
-    ) -> tuple[torch.Tensor, dict]:
-        """Action selection.
+    ) -> torch.Tensor:
+        """Select action.
 
         Args:
-            current_step (int): current step.
-            state (torch.Tensor): current state.
-            env (ModelBasedAdapter): environment.
+            current_step (int): The current step.
+            state (torch.Tensor): The current state.
 
         Returns:
-            action (torch.Tensor): action.
-            info (dict): information.
+            The selected action.
         """
         assert state.shape[0] == 1, 'state shape should be [1, state_dim]'
         if current_step < self._cfgs.algo_cfgs.start_learning_steps:
@@ -306,13 +358,10 @@ def _select_action(  # pylint: disable=unused-argument
         assert action.shape == torch.Size(
             [1, *self._action_space.shape],
         ), 'action shape should be [batch_size, action_dim]'
-        info = {}
-        return action, info
+        return action
 
     def _store_real_data(  # pylint: disable=too-many-arguments,unused-argument
         self,
-        current_step: int,
-        ep_len: int,
         state: torch.Tensor,
         action: torch.Tensor,
         reward: torch.Tensor,
@@ -320,23 +369,19 @@ def _store_real_data(  # pylint: disable=too-many-arguments,unused-argument
         terminated: torch.Tensor,
         truncated: torch.Tensor,
         next_state: torch.Tensor,
-        info: dict,
-        action_info: dict,
+        info: dict[str, Any],
     ) -> None:  # pylint: disable=too-many-arguments
         """Store real data in buffer.
 
         Args:
-            current_step (int): current step.
-            ep_len (int): episode length.
-            state (torch.Tensor): current state.
-            action (torch.Tensor): action.
-            reward (torch.Tensor): reward.
-            cost (torch.Tensor): cost.
-            terminated (torch.Tensor): terminated.
-            truncated (torch.Tensor): truncated.
-            next_state (torch.Tensor): next state.
-            info (dict): information.
-            action_info (dict): action information.
+            state (torch.Tensor): The state from the environment.
+            action (torch.Tensor): The action from the agent.
+            reward (torch.Tensor): The reward signal from the environment.
+            cost (torch.Tensor): The cost signal from the environment.
+            terminated (torch.Tensor): The terminated signal from the environment.
+            truncated (torch.Tensor): The truncated signal from the environment.
+            next_state (torch.Tensor): The next state from the environment.
+            info (dict[str, Any]): The information from the environment.
         """
         done = terminated or truncated
         goal_met = False if 'goal_met' not in info.keys() else info['goal_met']
@@ -360,9 +405,8 @@ def _evaluation_single_step(  # pylint: disable=too-many-locals
         """Evaluation dynamics model single step.
 
         Args:
-            current_step (int): current step.
-            use_real_input (bool): use real input or not.
-
+            current_step (int): The current step.
+            use_real_input (bool): Whether to use real input or not.
         """
         obs, _ = self._eval_env.reset()
         obs_dynamics = obs
@@ -379,7 +423,7 @@ def _evaluation_single_step(  # pylint: disable=too-many-locals
                 print(f'Eval Episode Return: {ep_ret} \t Cost: {ep_cost}')
                 save_replay_path = os.path.join(self._logger.log_dir, 'video-pic')
                 self._logger.store(
-                    **{
+                    {
                         'EvalMetrics/EpRet': ep_ret,
                         'EvalMetrics/EpCost': ep_cost,
                         'EvalMetrics/EpLen': ep_len,
@@ -419,7 +463,7 @@ def _evaluation_single_step(  # pylint: disable=too-many-locals
                 num_episode += 1
                 if num_episode == self._cfgs.evaluation_cfgs.num_episode:
                     break
-            action, _ = self._select_action(current_step, obs, self._eval_env)
+            action = self._select_action(current_step, obs)
 
             idx = np.random.choice(self._dynamics.elite_model_idxes, size=1)[0]
             traj = self._dynamics.imagine(
@@ -453,20 +497,20 @@ def draw_picture(
         self,
         timestep: int,
         num_episode: int,
-        pred_state: list,
-        true_state: list,
+        pred_state: list[float],
+        true_state: list[float],
         save_replay_path: str = './',
         name: str = 'reward',
     ) -> None:
         """Draw a curve of the predicted value and the ground true value.
 
         Args:
-            timestep (int): current step.
-            num_episode (int): number of episodes.
-            pred_state (list): predicted state.
-            true_state (list): true state.
-            save_replay_path (str): save replay path.
-            name (str): name of the curve.
+            timestep (int): The current step.
+            num_episode (int): The number of episodes.
+            pred_state (list[float]): The predicted state.
+            true_state (list[float]): The true state.
+            save_replay_path (str): The path for saving replay.
+            name (str): The name of the curve.
         """
         target1 = list(pred_state)
         target2 = list(true_state)
diff --git a/omnisafe/algorithms/model_based/cap_pets.py b/omnisafe/algorithms/model_based/cap_pets.py
index dd174a223..ee1d96135 100644
--- a/omnisafe/algorithms/model_based/cap_pets.py
+++ b/omnisafe/algorithms/model_based/cap_pets.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Implementation of the Conservative and Adaptive Penalty algorithm."""
+
+
 from __future__ import annotations
 
 from typing import Any
@@ -39,12 +41,23 @@ class CAPPETS(PETS):
     """
 
     def _init_model(self) -> None:
-        """Initialize the dynamics model and the planner."""
+        """Initialize the dynamics model and the planner.
+
+        CAP uses following models:
+
+        - dynamics model: to predict the next state and the cost.
+        - lagrange multiplier: to trade off between the cost and the reward.
+        - planner: to generate the action.
+        """
         self._dynamics_state_space = (
             self._env.coordinate_observation_space
             if self._env.coordinate_observation_space is not None
             else self._env.observation_space
         )
+        assert self._dynamics_state_space is not None and isinstance(
+            self._dynamics_state_space.shape,
+            tuple,
+        )
         assert self._env.action_space is not None and isinstance(
             self._env.action_space.shape,
             tuple,
@@ -54,7 +67,7 @@ def _init_model(self) -> None:
         else:
             raise NotImplementedError
 
-        self._dynamics = EnsembleDynamicsModel(
+        self._dynamics: EnsembleDynamicsModel = EnsembleDynamicsModel(
             model_cfgs=self._cfgs.dynamics_cfgs,
             device=self._device,
             state_shape=self._dynamics_state_space.shape,
@@ -65,9 +78,9 @@ def _init_model(self) -> None:
             terminal_func=None,
         )
 
-        self._lagrange = Lagrange(**self._cfgs.lagrange_cfgs)
+        self._lagrange: Lagrange = Lagrange(**self._cfgs.lagrange_cfgs)
 
-        self._planner = CAPPlanner(
+        self._planner: CAPPlanner = CAPPlanner(
             dynamics=self._dynamics,
             planner_cfgs=self._cfgs.planner_cfgs,
             gamma=float(self._cfgs.algo_cfgs.gamma),
@@ -81,11 +94,33 @@ def _init_model(self) -> None:
             lagrange=self._lagrange.lagrangian_multiplier,
         )
 
-        self._use_actor_critic = False
-        self._update_dynamics_cycle = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
+        self._use_actor_critic: bool = False
+        self._update_dynamics_cycle: int = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
 
     def _init_log(self) -> None:
-        """Initialize the logger."""
+        """Initialize the logger.
+
+        +----------------------------+-------------------------------+
+        | Things to log              | Description                   |
+        +============================+===============================+
+        | Plan/feasible_num          | The number of feasible plans. |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_max     | The maximum planning cost.    |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_mean    | The mean planning cost.       |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_min     | The minimum planning cost.    |
+        +----------------------------+-------------------------------+
+        | Metrics/LagrangeMultiplier | The lagrange multiplier.      |
+        +----------------------------+-------------------------------+
+        | Plan/var_penalty_max       | The maximum planning penalty. |
+        +----------------------------+-------------------------------+
+        | Plan/var_penalty_mean      | The mean planning penalty.    |
+        +----------------------------+-------------------------------+
+        | Plan/var_penalty_min       | The minimum planning penalty. |
+        +----------------------------+-------------------------------+
+
+        """
         super()._init_log()
         self._logger.register_key('Plan/feasible_num')
         self._logger.register_key('Plan/episode_costs_max')
diff --git a/omnisafe/algorithms/model_based/cce_pets.py b/omnisafe/algorithms/model_based/cce_pets.py
index b4751500e..c138a32c3 100644
--- a/omnisafe/algorithms/model_based/cce_pets.py
+++ b/omnisafe/algorithms/model_based/cce_pets.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Implementation of the Constrained Cross-Entropy algorithm."""
+
+
 from __future__ import annotations
 
 from gymnasium.spaces import Box
@@ -21,6 +23,7 @@
 from omnisafe.algorithms.model_based.base import PETS
 from omnisafe.algorithms.model_based.base.ensemble import EnsembleDynamicsModel
 from omnisafe.algorithms.model_based.planner.cce import CCEPlanner
+from omnisafe.typing import OmnisafeSpace
 
 
 @registry.register
@@ -31,17 +34,27 @@ class CCEPETS(PETS):
     References:
         - Title: Constrained Cross-Entropy Method for Safe Reinforcement Learning
         - Authors: Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess,
-        Tom Erez, Yuval Tassa, David Silver, Daan Wierstra.
+            Tom Erez, Yuval Tassa, David Silver, Daan Wierstra.
         - URL: `CCE <https://proceedings.neurips.cc/paper/2018/hash/34ffeb359a192eb8174b6854643cc046-Abstract.html>`_
     """
 
     def _init_model(self) -> None:
-        """Initialize the dynamics model and the planner."""
-        self._dynamics_state_space = (
+        """Initialize the dynamics model and the planner.
+
+        CCEPETS uses following models:
+
+        - dynamics model: to predict the next state and the cost.
+        - planner: to generate the action.
+        """
+        self._dynamics_state_space: OmnisafeSpace = (
             self._env.coordinate_observation_space
             if self._env.coordinate_observation_space is not None
             else self._env.observation_space
         )
+        assert self._dynamics_state_space is not None and isinstance(
+            self._dynamics_state_space.shape,
+            tuple,
+        )
         assert self._env.action_space is not None and isinstance(
             self._env.action_space.shape,
             tuple,
@@ -50,7 +63,7 @@ def _init_model(self) -> None:
             self._action_space = self._env.action_space
         else:
             raise NotImplementedError
-        self._dynamics = EnsembleDynamicsModel(
+        self._dynamics: EnsembleDynamicsModel = EnsembleDynamicsModel(
             model_cfgs=self._cfgs.dynamics_cfgs,
             device=self._device,
             state_shape=self._dynamics_state_space.shape,
@@ -61,7 +74,7 @@ def _init_model(self) -> None:
             terminal_func=None,
         )
 
-        self._planner = CCEPlanner(
+        self._planner: CCEPlanner = CCEPlanner(
             dynamics=self._dynamics,
             planner_cfgs=self._cfgs.planner_cfgs,
             gamma=float(self._cfgs.algo_cfgs.gamma),
@@ -74,11 +87,24 @@ def _init_model(self) -> None:
             cost_limit=self._cfgs.algo_cfgs.cost_limit,
         )
 
-        self._use_actor_critic = False
-        self._update_dynamics_cycle = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
+        self._use_actor_critic: bool = False
+        self._update_dynamics_cycle: int = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
 
     def _init_log(self) -> None:
-        """Initialize the logger keys for the CCE algorithm."""
+        """Initialize the logger keys for the CCE algorithm.
+
+        +----------------------------+-------------------------------+
+        | Things to log              | Description                   |
+        +============================+===============================+
+        | Plan/feasible_num          | The number of feasible plans. |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_max     | The maximum planning cost.    |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_mean    | The mean planning cost.       |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_min     | The minimum planning cost.    |
+        +----------------------------+-------------------------------+
+        """
         super()._init_log()
         self._logger.register_key('Plan/feasible_num')
         self._logger.register_key('Plan/episode_costs_max')
diff --git a/omnisafe/algorithms/model_based/planner/arc.py b/omnisafe/algorithms/model_based/planner/arc.py
index 68d3c83b7..5817f8569 100644
--- a/omnisafe/algorithms/model_based/planner/arc.py
+++ b/omnisafe/algorithms/model_based/planner/arc.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Model Predictive Control Planner of the Actor Regularized Control (ARC) algorithm."""
+
+
 from __future__ import annotations
 
 from typing import Any
@@ -63,8 +65,8 @@ def __init__(  # pylint: disable=too-many-locals, too-many-arguments
         self._actor_critic: ConstraintActorQCritic = kwargs['actor_critic']
         self._mixture_coefficient: float = planner_cfgs.mixture_coefficient
         self._temperature: float = planner_cfgs.temperature
-        self._actor_traj = int(self._mixture_coefficient * self._num_samples)
-        self._num_action = self._actor_traj + self._num_samples
+        self._actor_traj: int = int(self._mixture_coefficient * self._num_samples)
+        self._num_action: int = self._actor_traj + self._num_samples
         assert (
             self._num_samples + self._mixture_coefficient * self._num_samples
         ) > self._num_elites, 'The number of samples should be larger than the number of elites.'
@@ -78,7 +80,7 @@ def _act_from_last_gaus(self, last_mean: torch.Tensor, last_var: torch.Tensor) -
             last_var (torch.Tensor): Last variance of the gaussian distribution.
 
         Returns:
-            Sample actions (torch.Tensor): Sampled actions from the last gaussian distribution.
+            sampled actions: Sampled actions from the last gaussian distribution.
         """
         constrained_std = torch.sqrt(last_var)
         actions = torch.clamp(
@@ -101,10 +103,10 @@ def _act_from_actor(self, state: torch.Tensor) -> torch.Tensor:
         """Sample actions from the actor.
 
         Args:
-            state (torch.Tensor): Current state.
+            state (torch.Tensor): The current state.
 
         Returns:
-            Sample actions (torch.Tensor): Sampled actions from the actor.
+            sampled actions: Sampled actions from the actor.
         """
         assert state.shape == torch.Size(
             [1, *self._dynamics_state_shape],
@@ -135,12 +137,12 @@ def _state_action_repeat(
         """Repeat the state for num_repeat * action.shape[0] times and action for num_repeat times.
 
         Args:
-            state (torch.Tensor): Current state.
-            action (torch.Tensor): Sampled actions.
+            state (torch.Tensor): The current state.
+            action (torch.Tensor): The sampled actions.
 
         Returns:
-            states (torch.Tensor): Repeated states.
-            actions (torch.Tensor): Repeated actions.
+            states: The repeated states.
+            actions: The repeated actions.
         """
         assert action.shape == torch.Size(
             [self._horizon, self._num_action, *self._action_shape],
@@ -161,18 +163,18 @@ def _state_action_repeat(
     def _select_elites(
         self,
         actions: torch.Tensor,
-        traj: dict,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        traj: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
         """Select elites from the sampled actions.
 
         Args:
             actions (torch.Tensor): Sampled actions.
-            traj (dict): Trajectory dictionary.
+            traj (dict[str, torch.Tensor]): Trajectory dictionary.
 
         Returns:
-            elites_value (torch.Tensor): Value of the elites.
-            elites_action (torch.Tensor): Action of the elites.
-            info (dict): Dictionary containing the information of elites value and action.
+            elites_value: The value of the elites.
+            elites_action: The action of the elites.
+            info: The dictionary containing the information of elites value and action.
         """
         rewards = traj['rewards']
         values = traj['values']
@@ -222,18 +224,18 @@ def _update_mean_var(
         self,
         elite_actions: torch.Tensor,
         elite_values: torch.Tensor,
-        info: dict,
+        info: dict[str, int | float],
     ) -> tuple[torch.Tensor, torch.Tensor]:  # pylint: disable-next=unused-argument
         """Update the mean and variance of the elite actions.
 
         Args:
-            elite_actions (torch.Tensor): Elite actions.
-            elite_values (torch.Tensor): Elite values.
-            kwargs (dict): Keyword arguments.
+            elite_actions (torch.Tensor): The elite actions.
+            elite_values (torch.Tensor): The elite values.
+            info (dict[str, int | float]): The dictionary containing the information of the elite values and actions.
 
         Returns:
-            new_mean (torch.Tensor): New mean of the elite actions.
-            new_var (torch.Tensor): New variance of the elite actions.
+            new_mean: The new mean of the elite actions.
+            new_var: The new variance of the elite actions.
         """
         assert (
             elite_actions.shape[0] == self._horizon
@@ -259,15 +261,15 @@ def _update_mean_var(
         return new_mean, new_var
 
     @torch.no_grad()
-    def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict]:
+    def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict[str, int | float]]:
         """Output the action given the state.
 
         Args:
             state (torch.Tensor): State of the environment.
 
         Returns:
-            action (torch.Tensor): Action of the environment.
-            logger_info (dict): Dictionary containing the information of the action.
+            action: The action of the agent.
+            info: The dictionary containing the information of the action.
         """
         assert state.shape == torch.Size(
             [1, *self._dynamics_state_shape],
@@ -279,6 +281,7 @@ def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict]:
 
         current_iter = 0
         actions_actor = self._act_from_actor(state)
+        info: dict[str, float | int] = {}
         while current_iter < self._num_iterations and last_var.max() > self._epsilon:
             actions_gauss = self._act_from_last_gaus(last_mean=last_mean, last_var=last_var)
             actions = torch.cat([actions_gauss, actions_actor], dim=1)
diff --git a/omnisafe/algorithms/model_based/planner/cap.py b/omnisafe/algorithms/model_based/planner/cap.py
index a5332f711..c31c5bfce 100644
--- a/omnisafe/algorithms/model_based/planner/cap.py
+++ b/omnisafe/algorithms/model_based/planner/cap.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Model Predictive Control Planner of the Conservative and Adaptive Penalty algorithm."""
+
+
 from __future__ import annotations
 
 from typing import Any
@@ -65,18 +67,18 @@ def __init__(  # pylint: disable=too-many-locals, too-many-arguments
     def _select_elites(  # pylint: disable=too-many-locals
         self,
         actions: torch.Tensor,
-        traj: dict,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        traj: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
         """Select elites from the sampled actions.
 
         Args:
             actions (torch.Tensor): Sampled actions.
-            traj (dict): Trajectory dictionary.
+            traj (dict[str, torch.Tensor]): Trajectory dictionary.
 
         Returns:
-            elites_value (torch.Tensor): Value of the elites.
-            elites_action (torch.Tensor): Action of the elites.
-            info (dict): Dictionary containing the information of elites value and action.
+            elites_value: The value of the elites.
+            elites_action: The action of the elites.
+            info: The dictionary containing the information of elites value and action.
         """
         rewards = traj['rewards']
         costs = traj['costs']
diff --git a/omnisafe/algorithms/model_based/planner/cce.py b/omnisafe/algorithms/model_based/planner/cce.py
index 4d87b0082..2383dad1a 100644
--- a/omnisafe/algorithms/model_based/planner/cce.py
+++ b/omnisafe/algorithms/model_based/planner/cce.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Model Predictive Control Planner of the Constrained Cross-Entropy algorithm."""
+
+
 from __future__ import annotations
 
 from typing import Any
@@ -30,7 +32,7 @@ class CCEPlanner(CEMPlanner):
     References:
         - Title: Constrained Cross-Entropy Method for Safe Reinforcement Learning
         - Authors: Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess,
-        Tom Erez, Yuval Tassa, David Silver, Daan Wierstra.
+            Tom Erez, Yuval Tassa, David Silver, Daan Wierstra.
         - URL: `CCE <https://proceedings.neurips.cc/paper/2018/hash/34ffeb359a192eb8174b6854643cc046-Abstract.html>`_
     """
 
@@ -66,18 +68,18 @@ def __init__(  # pylint: disable=too-many-locals, too-many-arguments
     def _select_elites(
         self,
         actions: torch.Tensor,
-        traj: dict,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        traj: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
         """Select elites from the sampled actions.
 
         Args:
             actions (torch.Tensor): Sampled actions.
-            traj (dict): Trajectory dictionary.
+            traj (dict[str, torch.Tensor]): Trajectory dictionary.
 
         Returns:
-            elites_value (torch.Tensor): Value of the elites.
-            elites_action (torch.Tensor): Action of the elites.
-            info (dict): Dictionary containing the information of elites value and action.
+            elites_value: The value of the elites.
+            elites_action: The action of the elites.
+            info: The dictionary containing the information of elites value and action.
         """
         rewards = traj['rewards']
         costs = traj['costs']
diff --git a/omnisafe/algorithms/model_based/planner/cem.py b/omnisafe/algorithms/model_based/planner/cem.py
index 99efdbce0..401b5db73 100644
--- a/omnisafe/algorithms/model_based/planner/cem.py
+++ b/omnisafe/algorithms/model_based/planner/cem.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Model Predictive Control Planner of Cross-Entropy Method optimization algorithm."""
+
+
 from __future__ import annotations
 
 from typing import Any
@@ -27,7 +29,10 @@ class CEMPlanner:  # pylint: disable=too-many-instance-attributes
     """The planner of  Cross-Entropy Method optimization (CEM) algorithm.
 
     References:
-        - URL: `A good description of CEM <https://arxiv.org/pdf/2008.06389.pdf>`_
+        - Title: Sample-efficient Cross-Entropy Method for Real-time Planning
+        - Authors: Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, Jan Achterhold,
+            Joerg Stueckler, Michal Rolinek, Georg Martius
+        - URL: `CEM <https://arxiv.org/pdf/2008.06389.pdf>`_
     """
 
     def __init__(  # pylint: disable=too-many-locals, too-many-arguments
@@ -88,7 +93,7 @@ def _act_from_last_gaus(self, last_mean: torch.Tensor, last_var: torch.Tensor) -
             last_var (torch.Tensor): Last variance of the gaussian distribution.
 
         Returns:
-            Sample actions (torch.Tensor): Sampled actions from the last gaussian distribution.
+            sampled actions: Sampled actions from the last gaussian distribution.
         """
         constrained_std = torch.sqrt(last_var)
         actions = torch.clamp(
@@ -115,12 +120,12 @@ def _state_action_repeat(
         """Repeat the state for num_repeat * action.shape[0] times and action for num_repeat times.
 
         Args:
-            state (torch.Tensor): Current state.
-            action (torch.Tensor): Sampled actions.
+            state (torch.Tensor): The current state.
+            action (torch.Tensor): The sampled actions.
 
         Returns:
-            states (torch.Tensor): Repeated states.
-            actions (torch.Tensor): Repeated actions.
+            states: The repeated states.
+            actions: The repeated actions.
         """
         assert (
             self._num_particles % self._num_models == 0
@@ -144,18 +149,18 @@ def _state_action_repeat(
     def _select_elites(
         self,
         actions: torch.Tensor,
-        traj: dict,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        traj: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
         """Select elites from the sampled actions.
 
         Args:
             actions (torch.Tensor): Sampled actions.
-            traj (dict): Trajectory dictionary.
+            traj (dict[str, torch.Tensor]): Trajectory dictionary.
 
         Returns:
-            elites_value (torch.Tensor): Value of the elites.
-            elites_action (torch.Tensor): Action of the elites.
-            info (dict): Dictionary containing the information of elites value and action.
+            elites_value: The value of the elites.
+            elites_action: The action of the elites.
+            info: The dictionary containing the information of elites value and action.
         """
         rewards = traj['rewards']
         assert actions.shape == torch.Size(
@@ -193,17 +198,18 @@ def _update_mean_var(  # pylint: disable=unused-argument
         self,
         elite_actions: torch.Tensor,
         elite_values: torch.Tensor,
-        info: dict,
+        info: dict[str, int | float],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Update the mean and variance of the elite actions.
 
         Args:
-            elite_actions (torch.Tensor): Elite actions.
-            elite_values (torch.Tensor): Elite values.
+            elite_actions (torch.Tensor): The elite actions.
+            elite_values (torch.Tensor): The elite values.
+            info (dict[str, int | float]): The dictionary containing the information of the elite values and actions.
 
         Returns:
-            new_mean (torch.Tensor): New mean of the elite actions.
-            new_var (torch.Tensor): New variance of the elite actions.
+            new_mean: The new mean of the elite actions.
+            new_var: The new variance of the elite actions.
         """
         assert elite_actions.shape == torch.Size(
             [self._horizon, self._num_elites, *self._action_shape],
@@ -225,8 +231,8 @@ def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict]:
             state (torch.Tensor): State of the environment.
 
         Returns:
-            action (torch.Tensor): Action of the environment.
-            logger_info (dict): Dictionary containing the information of the action.
+            action: The action of the agent.
+            info: The dictionary containing the information of the action.
         """
         assert state.shape == torch.Size(
             [1, *self._dynamics_state_shape],
@@ -237,6 +243,7 @@ def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict]:
         last_mean[-1] = self._action_sequence_mean[-1].clone()
 
         current_iter = 0
+        info: dict[str, float | int] = {}
         while current_iter < self._num_iterations:
             actions = self._act_from_last_gaus(last_mean=last_mean, last_var=last_var)
             # [horizon, num_sample, action_shape]
diff --git a/omnisafe/algorithms/model_based/planner/rce.py b/omnisafe/algorithms/model_based/planner/rce.py
index 1560497e6..f35045008 100644
--- a/omnisafe/algorithms/model_based/planner/rce.py
+++ b/omnisafe/algorithms/model_based/planner/rce.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Model Predictive Control Planner of Robust Cross Entropy algorithm."""
+
+
 from __future__ import annotations
 
 import torch
@@ -33,18 +35,18 @@ class RCEPlanner(CCEPlanner):
     def _select_elites(
         self,
         actions: torch.Tensor,
-        traj: dict,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        traj: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
         """Select elites from the sampled actions.
 
         Args:
             actions (torch.Tensor): Sampled actions.
-            traj (dict): Trajectory dictionary.
+            traj (dict[str, torch.Tensor]): Trajectory dictionary.
 
         Returns:
-            elites_value (torch.Tensor): Value of the elites.
-            elites_action (torch.Tensor): Action of the elites.
-            info (dict): Dictionary containing the information of elites value and action.
+            elites_value: The value of the elites.
+            elites_action: The action of the elites.
+            info: The dictionary containing the information of elites value and action.
         """
         rewards = traj['rewards']
         costs = traj['costs']
diff --git a/omnisafe/algorithms/model_based/planner/safe_arc.py b/omnisafe/algorithms/model_based/planner/safe_arc.py
index 643e0ad18..3d7aa1e0a 100644
--- a/omnisafe/algorithms/model_based/planner/safe_arc.py
+++ b/omnisafe/algorithms/model_based/planner/safe_arc.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Model Predictive Control Planner of the Safe Actor Regularized Control algorithm."""
+
+
 from __future__ import annotations
 
 from typing import Any
@@ -67,17 +69,18 @@ def _update_mean_var(
         self,
         elite_actions: torch.Tensor,
         elite_values: torch.Tensor,
-        info: dict,
+        info: dict[str, int | float],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Update the mean and variance of the elite actions.
 
         Args:
-            elite_actions (torch.Tensor): Elite actions.
-            elite_values (torch.Tensor): Elite values.
+            elite_actions (torch.Tensor): The elite actions.
+            elite_values (torch.Tensor): The elite values.
+            info (dict[str, int | float]): The dictionary containing the information of the elite values and actions.
 
         Returns:
-            new_mean (torch.Tensor): New mean of the elite actions.
-            new_var (torch.Tensor): New variance of the elite actions.
+            new_mean: The new mean of the elite actions.
+            new_var: The new variance of the elite actions.
         """
         assert (
             elite_actions.shape[0] == self._horizon
@@ -110,18 +113,18 @@ def _update_mean_var(
     def _select_elites(
         self,
         actions: torch.Tensor,
-        traj: dict,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        traj: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
         """Select elites from the sampled actions.
 
         Args:
             actions (torch.Tensor): Sampled actions.
-            traj (dict): Trajectory dictionary.
+            traj (dict[str, torch.Tensor]): Trajectory dictionary.
 
         Returns:
-            elites_value (torch.Tensor): Value of the elites.
-            elites_action (torch.Tensor): Action of the elites.
-            info (dict): Dictionary containing the information of elites value and action.
+            elites_value: The value of the elites.
+            elites_action: The action of the elites.
+            info: The dictionary containing the information of elites value and action.
         """
         rewards = traj['rewards']
         values = traj['values']
@@ -194,15 +197,15 @@ def _select_elites(
         return elite_values, elite_actions, info
 
     @torch.no_grad()
-    def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict]:
+    def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict[str, int | float]]:
         """Output the action given the state.
 
         Args:
             state (torch.Tensor): State of the environment.
 
         Returns:
-            action (torch.Tensor): Action of the environment.
-            logger_info (dict): Dictionary containing the information of the action.
+            action: The action of the agent.
+            info: The dictionary containing the information of the action.
         """
         assert state.shape == torch.Size(
             [1, *self._dynamics_state_shape],
@@ -214,6 +217,7 @@ def output_action(self, state: torch.Tensor) -> tuple[torch.Tensor, dict]:
 
         current_iter = 0
         actions_actor = self._act_from_actor(state)
+        info: dict[str, int | float] = {}
         while current_iter < self._num_iterations and last_var.max() > self._epsilon:
             actions_gauss = self._act_from_last_gaus(last_mean=last_mean, last_var=last_var)
             actions = torch.cat([actions_gauss, actions_actor], dim=1)
diff --git a/omnisafe/algorithms/model_based/rce_pets.py b/omnisafe/algorithms/model_based/rce_pets.py
index 5e2b9d4a0..05e5c15f5 100644
--- a/omnisafe/algorithms/model_based/rce_pets.py
+++ b/omnisafe/algorithms/model_based/rce_pets.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Implementation of the Robust Cross Entropy algorithm."""
+
+
 from __future__ import annotations
 
 from gymnasium.spaces import Box
@@ -35,12 +37,22 @@ class RCEPETS(PETS):
     """
 
     def _init_model(self) -> None:
-        """Initialize the dynamics model and the planner."""
+        """Initialize the dynamics model and the planner.
+
+        RCEPETS uses following models:
+
+        - dynamics model: to predict the next state and the cost.
+        - planner: to generate the action.
+        """
         self._dynamics_state_space = (
             self._env.coordinate_observation_space
             if self._env.coordinate_observation_space is not None
             else self._env.observation_space
         )
+        assert self._dynamics_state_space is not None and isinstance(
+            self._dynamics_state_space.shape,
+            tuple,
+        )
         assert self._env.action_space is not None and isinstance(
             self._env.action_space.shape,
             tuple,
@@ -49,7 +61,7 @@ def _init_model(self) -> None:
             self._action_space = self._env.action_space
         else:
             raise NotImplementedError
-        self._dynamics = EnsembleDynamicsModel(
+        self._dynamics: EnsembleDynamicsModel = EnsembleDynamicsModel(
             model_cfgs=self._cfgs.dynamics_cfgs,
             device=self._device,
             state_shape=self._dynamics_state_space.shape,
@@ -60,7 +72,7 @@ def _init_model(self) -> None:
             terminal_func=None,
         )
 
-        self._planner = RCEPlanner(
+        self._planner: RCEPlanner = RCEPlanner(
             dynamics=self._dynamics,
             planner_cfgs=self._cfgs.planner_cfgs,
             gamma=float(self._cfgs.algo_cfgs.gamma),
@@ -73,11 +85,26 @@ def _init_model(self) -> None:
             cost_limit=self._cfgs.algo_cfgs.cost_limit,
         )
 
-        self._use_actor_critic = False
-        self._update_dynamics_cycle = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
+        self._use_actor_critic: bool = False
+        self._update_dynamics_cycle: int = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
 
     def _init_log(self) -> None:
-        """Initialize the logger."""
+        """Initialize the logger.
+
+        +----------------------------+-------------------------------+
+        | Things to log              | Description                   |
+        +============================+===============================+
+        | Plan/feasible_num          | The number of feasible plans. |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_max     | The maximum planning cost.    |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_mean    | The mean planning cost.       |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_min     | The minimum planning cost.    |
+        +----------------------------+-------------------------------+
+        | Metrics/LagrangeMultiplier | The lagrange multiplier.      |
+        +----------------------------+-------------------------------+
+        """
         super()._init_log()
         self._logger.register_key('Plan/feasible_num')
         self._logger.register_key('Plan/episode_costs_max')
diff --git a/omnisafe/algorithms/model_based/safeloop.py b/omnisafe/algorithms/model_based/safeloop.py
index ec0914c9e..6d38ca499 100644
--- a/omnisafe/algorithms/model_based/safeloop.py
+++ b/omnisafe/algorithms/model_based/safeloop.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Implementation of the Safe Learning Off-Policy with Online Planning algorithm."""
+
+
 from __future__ import annotations
 
 from gymnasium.spaces import Box
@@ -22,6 +24,7 @@
 from omnisafe.algorithms.model_based.base.loop import LOOP
 from omnisafe.algorithms.model_based.planner.safe_arc import SafeARCPlanner
 from omnisafe.models.actor_critic.constraint_actor_q_critic import ConstraintActorQCritic
+from omnisafe.typing import OmnisafeSpace
 from omnisafe.utils import distributed
 
 
@@ -37,12 +40,22 @@ class SafeLOOP(LOOP):
     """
 
     def _init_model(self) -> None:
-        """Initialize the dynamics model and the planner."""
-        self._dynamics_state_space = (
+        """Initialize the dynamics model and the planner.
+
+        SafeLOOP uses following models:
+
+        - dynamics model: to predict the next state and the cost.
+        - planner: to generate the action.
+        """
+        self._dynamics_state_space: OmnisafeSpace = (
             self._env.coordinate_observation_space
             if self._env.coordinate_observation_space is not None
             else self._env.observation_space
         )
+        assert self._dynamics_state_space is not None and isinstance(
+            self._dynamics_state_space.shape,
+            tuple,
+        )
         assert self._env.action_space is not None and isinstance(
             self._env.action_space.shape,
             tuple,
@@ -51,7 +64,7 @@ def _init_model(self) -> None:
             self._action_space = self._env.action_space
         else:
             raise NotImplementedError
-        self._actor_critic = ConstraintActorQCritic(
+        self._actor_critic: ConstraintActorQCritic = ConstraintActorQCritic(
             obs_space=self._dynamics_state_space,
             act_space=self._action_space,
             model_cfgs=self._cfgs.model_cfgs,
@@ -59,9 +72,9 @@ def _init_model(self) -> None:
         ).to(self._device)
         if distributed.world_size() > 1:
             distributed.sync_params(self._actor_critic)
-        self._use_actor_critic = True
-        self._update_count = 0
-        self._dynamics = EnsembleDynamicsModel(
+        self._use_actor_critic: bool = True
+        self._update_count: int = 0
+        self._dynamics: EnsembleDynamicsModel = EnsembleDynamicsModel(
             model_cfgs=self._cfgs.dynamics_cfgs,
             device=self._device,
             state_shape=self._dynamics_state_space.shape,
@@ -71,9 +84,9 @@ def _init_model(self) -> None:
             cost_func=self._env.get_cost_from_obs_tensor,
             terminal_func=None,
         )
-        self._update_dynamics_cycle = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
+        self._update_dynamics_cycle: int = int(self._cfgs.algo_cfgs.update_dynamics_cycle)
 
-        self._planner = SafeARCPlanner(
+        self._planner: SafeARCPlanner = SafeARCPlanner(
             dynamics=self._dynamics,
             planner_cfgs=self._cfgs.planner_cfgs,
             gamma=float(self._cfgs.algo_cfgs.gamma),
@@ -88,7 +101,20 @@ def _init_model(self) -> None:
         )
 
     def _init_log(self) -> None:
-        """Initialize the logger keys for the algorithm."""
+        """Initialize the logger keys for the algorithm.
+
+        +----------------------------+-------------------------------+
+        | Things to log              | Description                   |
+        +============================+===============================+
+        | Plan/feasible_num          | The number of feasible plans. |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_max     | The maximum planning cost.    |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_mean    | The mean planning cost.       |
+        +----------------------------+-------------------------------+
+        | Plan/episode_costs_min     | The minimum planning cost.    |
+        +----------------------------+-------------------------------+
+        """
         super()._init_log()
         self._logger.register_key('Plan/feasible_num')
         self._logger.register_key('Plan/episode_costs_max')
diff --git a/omnisafe/algorithms/on_policy/base/policy_gradient.py b/omnisafe/algorithms/on_policy/base/policy_gradient.py
index 07ece2b80..1ad7eb772 100644
--- a/omnisafe/algorithms/on_policy/base/policy_gradient.py
+++ b/omnisafe/algorithms/on_policy/base/policy_gradient.py
@@ -232,9 +232,9 @@ def learn(self) -> tuple[float, float, float]:
         - :meth:`log`: epoch/update information for visualization and terminal log print.
 
         Returns:
-            ep_ret: average episode return in final epoch.
-            ep_cost: average episode cost in final epoch.
-            ep_len: average episode length in final epoch.
+            ep_ret: Average episode return in final epoch.
+            ep_cost: Average episode cost in final epoch.
+            ep_len: Average episode length in final epoch.
         """
         start_time = time.time()
         self._logger.log('INFO: Start training')
diff --git a/omnisafe/common/logger.py b/omnisafe/common/logger.py
index fc07b26a4..3ec0bd045 100644
--- a/omnisafe/common/logger.py
+++ b/omnisafe/common/logger.py
@@ -249,7 +249,7 @@ def store(
         self,
         data: dict[str, Any] | None = None,
         /,
-        **kwargs: Any,
+        **kwargs: Any | float | np.ndarray | torch.Tensor,
     ) -> None:
         """Store the data to the logger.
 
@@ -259,7 +259,9 @@ def store(
         Args:
             data (dict[str, int | float | np.ndarray | torch.Tensor] or None, optional): The data to
                 be stored. Defaults to None.
-            **kwargs (int, float, np.ndarray, or torch.Tensor): The data to be stored.
+
+        Keyword Args:
+            kwargs (int, float, np.ndarray, or torch.Tensor): The data to be stored.
         """
         if data is not None:
             kwargs.update(data)
diff --git a/omnisafe/envs/core.py b/omnisafe/envs/core.py
index 7709e08df..bca9bfca5 100644
--- a/omnisafe/envs/core.py
+++ b/omnisafe/envs/core.py
@@ -155,7 +155,7 @@ def render(self) -> Any:
         """Compute the render frames as specified by :attr:`render_mode` during the initialization of the environment.
 
         Returns:
-            The render frames, we recommend to use `np.ndarray` which could construct video by
+            The rendered frames, we recommend using `np.ndarray` which could construct video by
             moviepy.
         """
 
@@ -268,7 +268,7 @@ def render(self) -> Any:
         """Compute the render frames as specified by :attr:`render_mode` during the initialization of the environment.
 
         Returns:
-            The render frames, we recommend to use `np.ndarray` which could construct video by
+            The rendered frames, we recommend using `np.ndarray` which could construct video by
             moviepy.
         """
         return self._env.render()
@@ -384,7 +384,14 @@ def make(env_id: str, class_name: str | None = None, **kwargs: Any) -> CMDP:
     Args:
         env_id (str): The environment id.
         class_name (str or None): The environment class name.
-        **kwargs: the keyword arguments for the environment initialization.
+
+    Keyword Args:
+        render_mode (str, optional): The render mode ranges from 'human' to 'rgb_array' and 'rgb_array_list'.
+            Defaults to 'rgb_array'.
+        camera_name (str, optional): The camera name.
+        camera_id (int, optional): The camera id.
+        width (int, optional): The width of the rendered image. Defaults to 256.
+        height (int, optional): The height of the rendered image. Defaults to 256.
 
     Returns:
         The environment class.
diff --git a/omnisafe/envs/mujoco_env.py b/omnisafe/envs/mujoco_env.py
index 128a8450c..eccad10a3 100644
--- a/omnisafe/envs/mujoco_env.py
+++ b/omnisafe/envs/mujoco_env.py
@@ -30,7 +30,6 @@ class MujocoEnv(CMDP):
     """Gymnasium Mujoco environment.
 
     Attributes:
-        _support_envs (list[str]): List of supported environments.
         need_auto_reset_wrapper (bool): Whether to use auto reset wrapper.
         need_time_limit_wrapper (bool): Whether to use time limit wrapper.
     """
@@ -61,7 +60,14 @@ def __init__(
             env_id (str): Environment id.
             num_envs (int, optional): Number of environments. Defaults to 1.
             device (torch.device, optional): Device to store the data. Defaults to 'cpu'.
-            **kwargs: Other arguments.
+
+        Keyword Args:
+            render_mode (str, optional): The render mode, ranging from ``human``, ``rgb_array``, ``rgb_array_list``.
+                Defaults to ``rgb_array``.
+            camera_name (str, optional): The camera name.
+            camera_id (int, optional): The camera id.
+            width (int, optional): The width of the rendered image. Defaults to 256.
+            height (int, optional): The height of the rendered image. Defaults to 256.
         """
         super().__init__(env_id)
         self._env_id = env_id
@@ -85,7 +91,14 @@ def __init__(
     def step(
         self,
         action: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict]:
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        dict[str, Any],
+    ]:
         """Step the environment.
 
         .. note::
@@ -98,12 +111,12 @@ def step(
             action (torch.Tensor): Action to take.
 
         Returns:
-            observation (torch.Tensor): agent's observation of the current environment.
-            reward (torch.Tensor): amount of reward returned after previous action.
-            cost (torch.Tensor): amount of cost returned after previous action.
-            terminated (torch.Tensor): whether the episode has ended.
-            truncated (torch.Tensor): whether the episode has been truncated due to a time limit.
-            info (Dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning).
+            observation: Agent's observation of the current environment.
+            reward: Amount of reward returned after previous action.
+            cost: Amount of cost returned after previous action.
+            terminated: Whether the episode has ended.
+            truncated: Whether the episode has been truncated due to a time limit.
+            info: Auxiliary diagnostic information (helpful for debugging, and sometimes learning).
         """
         obs, reward, terminated, truncated, info = self._env.step(
             action.detach().cpu().numpy(),
@@ -135,8 +148,8 @@ def reset(self, seed: int | None = None) -> tuple[torch.Tensor, dict]:
             seed (int, optional): Seed to reset the environment. Defaults to None.
 
         Returns:
-            observation (torch.Tensor): agent's observation of the current environment.
-            info (Dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning).
+            observation: Agent's observation of the current environment.
+            info: Auxiliary diagnostic information (helpful for debugging, and sometimes learning).
         """
         obs, info = self._env.reset(seed=seed)
         return torch.as_tensor(obs, dtype=torch.float32, device=self._device), info
@@ -153,7 +166,7 @@ def sample_action(self) -> torch.Tensor:
         """Sample a random action.
 
         Returns:
-            torch.Tensor: A random action.
+            A random action.
         """
         return torch.as_tensor(
             self._env.action_space.sample(),
@@ -165,7 +178,7 @@ def render(self) -> Any:
         """Render the environment.
 
         Returns:
-            Any: Rendered environment.
+            Rendered environment.
         """
         return self._env.render()
 
diff --git a/omnisafe/envs/safety_gymnasium_env.py b/omnisafe/envs/safety_gymnasium_env.py
index 7f151eab5..f87871807 100644
--- a/omnisafe/envs/safety_gymnasium_env.py
+++ b/omnisafe/envs/safety_gymnasium_env.py
@@ -35,7 +35,14 @@ class SafetyGymnasiumEnv(CMDP):
         num_envs (int, optional): Number of environments. Defaults to 1.
         device (torch.device, optional): Device to store the data. Defaults to
             ``torch.device('cpu')``.
-        **kwargs (Any): Other arguments.
+
+    Keyword Args:
+        render_mode (str, optional): The render mode ranges from 'human' to 'rgb_array' and 'rgb_array_list'.
+            Defaults to 'rgb_array'.
+        camera_name (str, optional): The camera name.
+        camera_id (int, optional): The camera id.
+        width (int, optional): The width of the rendered image. Defaults to 256.
+        height (int, optional): The height of the rendered image. Defaults to 256.
 
     Attributes:
         need_auto_reset_wrapper (bool): Whether to use auto reset wrapper.
diff --git a/omnisafe/envs/safety_gymnasium_modelbased.py b/omnisafe/envs/safety_gymnasium_modelbased.py
index 34aebd4b1..2c1f86043 100644
--- a/omnisafe/envs/safety_gymnasium_modelbased.py
+++ b/omnisafe/envs/safety_gymnasium_modelbased.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Environments in the Safety Gymnasium."""
+"""World model of the Safety Gymnasium."""
+
 
 from __future__ import annotations
 
@@ -24,6 +25,7 @@
 import torch
 
 from omnisafe.envs.core import CMDP, env_register
+from omnisafe.typing import Box, OmnisafeSpace
 
 
 @env_register
@@ -62,7 +64,14 @@ def __init__(
             num_envs (int, optional): Number of environments. Defaults to 1.
             device (torch.device, optional): Device to store the data. Defaults to 'cpu'.
             use_lidar (bool, optional): Whether to use lidar observation. Defaults to False.
-            **kwargs: Other arguments.
+
+        Keyword Args:
+            render_mode (str, optional): The render mode, ranging from ``human``, ``rgb_array``, ``rgb_array_list``.
+                Defaults to ``rgb_array``.
+            camera_name (str, optional): The camera name.
+            camera_id (int, optional): The camera id.
+            width (int, optional): The width of the rendered image. Defaults to 256.
+            height (int, optional): The height of the rendered image. Defaults to 256.
         """
         super().__init__(env_id)
         self._use_lidar = use_lidar
@@ -72,6 +81,11 @@ def __init__(
                 autoreset=False,
                 **kwargs,
             )
+            assert isinstance(self._env.action_space, Box), 'Only support Box action space.'
+            assert isinstance(
+                self._env.observation_space,
+                Box,
+            ), 'Only support Box observation space.'
             self._action_space = self._env.action_space
         else:
             raise NotImplementedError
@@ -80,10 +94,10 @@ def __init__(
 
         self._num_envs = num_envs
         self._metadata = self._env.metadata
-        self._constraints = ['hazards']  # gremlins, vase, buttons
-        self._xyz_sensors = ['velocimeter', 'accelerometer']
-        self._angle_sensors = ['gyro', 'magnetometer']
-        self._flatten_order = (
+        self._constraints: list[str] = ['hazards']  # gremlins, vase, buttons
+        self._xyz_sensors: list[str] = ['velocimeter', 'accelerometer']
+        self._angle_sensors: list[str] = ['gyro', 'magnetometer']
+        self._flatten_order: list[str] = (
             self._xyz_sensors
             + self._angle_sensors
             + ['goal']
@@ -91,21 +105,21 @@ def __init__(
             + ['robot_m']
             + ['robot']
         )
-        self._base_state = self._xyz_sensors + self._angle_sensors
-        self._task = 'Goal'
+        self._base_state: list[str] = self._xyz_sensors + self._angle_sensors
+        self._task: str = 'Goal'
         self._env.reset()
-        self.goal_position = self._env.task.goal.pos
-        self.robot_position = self._env.task.agent.pos
-        self.hazards_position = self._env.task.hazards.pos
-        self.goal_distance = self._dist_xy(self.robot_position, self.goal_position)
+        self.goal_position: np.ndarray = self._env.task.goal.pos
+        self.robot_position: np.ndarray = self._env.task.agent.pos
+        self.hazards_position: list[np.ndarray] = self._env.task.hazards.pos
+        self.goal_distance: float = self._dist_xy(self.robot_position, self.goal_position)
 
-        coordinate_sensor_obs = self._get_coordinate_sensor()
-        self._coordinate_obs_size = sum(
+        coordinate_sensor_obs: dict[str, Any] = self._get_coordinate_sensor()
+        self._coordinate_obs_size: int = sum(
             np.prod(i.shape) for i in list(coordinate_sensor_obs.values())
         )
-        offset = 0
-        self.key_to_slice = {}
-        self.key_to_slice_tensor = {}
+        offset: int = 0
+        self.key_to_slice: dict[str, slice] = {}
+        self.key_to_slice_tensor: dict[str, torch.Tensor] = {}
 
         for k in self._flatten_order:
             k_size = np.prod(coordinate_sensor_obs[k].shape)
@@ -113,25 +127,25 @@ def __init__(
             self.key_to_slice_tensor[k] = torch.arange(offset, offset + k_size)
 
             offset += k_size
-        self._base_state_size = sum(
+        self._base_state_size: int = sum(
             np.prod(coordinate_sensor_obs[k].shape) for k in list(self._base_state)
         )
         self.key_to_slice['base_state'] = slice(0, self._base_state_size)
         self.key_to_slice_tensor['base_state'] = torch.arange(0, self._base_state_size)
 
-        self._num_lidar_bin = 16
-        self._max_lidar_dist = 3
-        self.hazards_size = 0.2
-        self.goal_size = 0.3
-        self.original_observation_space = self._env.observation_space
-        self.coordinate_observation_space = gymnasium.spaces.Box(
+        self._num_lidar_bin: int = 16
+        self._max_lidar_dist: int = 3
+        self.hazards_size: float = 0.2
+        self.goal_size: float = 0.3
+        self.original_observation_space: OmnisafeSpace = self._env.observation_space
+        self.coordinate_observation_space: OmnisafeSpace = gymnasium.spaces.Box(
             -np.inf,
             np.inf,
             (self._coordinate_obs_size,),
             dtype=np.float32,
         )
         flat_coordinate_obs = self._get_flat_coordinate(coordinate_sensor_obs)
-        self.lidar_observation_space = gymnasium.spaces.Box(
+        self.lidar_observation_space: OmnisafeSpace = gymnasium.spaces.Box(
             -np.inf,
             np.inf,
             (self.get_lidar_from_coordinate(flat_coordinate_obs).shape[0],),
@@ -144,7 +158,7 @@ def __init__(
 
     @property
     def task(self) -> str:
-        """Get the name of the task."""
+        """The name of the task."""
         return self._task
 
     def get_cost_from_obs_tensor(self, obs: torch.Tensor, is_binary: bool = True) -> torch.Tensor:
@@ -155,7 +169,7 @@ def get_cost_from_obs_tensor(self, obs: torch.Tensor, is_binary: bool = True) ->
             is_binary (bool, optional): Whether to use binary cost. Defaults to True.
 
         Returns:
-            cost (torch.Tensor): Batch cost.
+            cost: Batch cost.
         """
         assert torch.is_tensor(obs), 'obs must be tensor'
         hazards_key = self.key_to_slice_tensor['hazards']
@@ -187,17 +201,17 @@ def get_lidar_from_coordinate(self, obs: np.ndarray) -> torch.Tensor:
         """Get lidar observation.
 
         Args:
-            obs (np.ndarray): observation.
+            obs (np.ndarray): The observation.
 
         Returns:
-            lidar_obs (np.ndarray): lidar observation.
+            lidar_obs: The lidar observation.
         """
         robot_matrix_x_y = obs[self.key_to_slice['robot_m']]
-        robot_matrix_x = robot_matrix_x_y[0]
-        robot_matrix_y = robot_matrix_x_y[1]
-        first_row = [robot_matrix_x, robot_matrix_y, 0]
-        second_row = [-robot_matrix_y, robot_matrix_x, 0]
-        third_row = [0, 0, 1]
+        robot_matrix_x = float(robot_matrix_x_y[0])
+        robot_matrix_y = float(robot_matrix_x_y[1])
+        first_row = [robot_matrix_x, robot_matrix_y, 0.0]
+        second_row = [-robot_matrix_y, robot_matrix_x, 0.0]
+        third_row = [0.0, 0.0, 1.0]
         robot_matrix = [first_row, second_row, third_row]
         robot_pos = obs[self.key_to_slice['robot']]
         hazards_lidar_vec = self._obs_lidar_pseudo(robot_matrix, robot_pos, self.hazards_position)
@@ -214,19 +228,19 @@ def get_lidar_from_coordinate(self, obs: np.ndarray) -> torch.Tensor:
 
     def _ego_xy(
         self,
-        robot_matrix: list,
+        robot_matrix: list[list[float]],
         robot_pos: np.ndarray,
         pos: np.ndarray,
     ) -> np.ndarray:
         """Return the egocentric XY vector to a position from the robot.
 
         Args:
-            robot_matrix (np.ndarray): 3x3 rotation matrix.
+            robot_matrix (list[list[float]]): 3x3 rotation matrix.
             robot_pos (np.ndarray): 2D robot position.
             pos (np.ndarray): 2D position.
 
         Returns:
-            2D_egocentric_vector (np.ndarray): 2D egocentric vector.
+            2D_egocentric_vector: The 2D egocentric vector.
         """
         assert pos.shape == (2,), f'Bad pos {pos}'
         assert robot_pos.shape == (2,), f'Bad robot_pos {robot_pos}'
@@ -240,9 +254,9 @@ def _ego_xy(
 
     def _obs_lidar_pseudo(
         self,
-        robot_matrix: list,
+        robot_matrix: list[list[float]],
         robot_pos: np.ndarray,
-        positions: list,
+        positions: list[np.ndarray],
     ) -> np.ndarray:  # pylint: disable=too-many-locals
         """Return a robot-centric lidar observation of a list of positions.
 
@@ -264,12 +278,12 @@ def _obs_lidar_pseudo(
             - constant size observation with variable numbers of objects
 
         Args:
-            robot_matrix (np.ndarray): 3x3 rotation matrix.
+            robot_matrix (list[list[float]]): 3x3 rotation matrix.
             robot_pos (np.ndarray): 2D robot position.
-            positions (np.ndarray): 2D positions.
+            positions (list[np.ndarray]): 2D positions.
 
         Returns:
-            lidar_observation (np.ndarray): lidar observation.
+            lidar_observation: The lidar observation.
         """
         obs = np.zeros(self._num_lidar_bin)
         for pos in positions:
@@ -297,29 +311,32 @@ def _obs_lidar_pseudo(
             obs[bin_minus] = max(obs[bin_minus], (1 - alias) * sensor)
         return obs
 
-    def _get_flat_coordinate(self, coordinate_obs: dict) -> np.ndarray:
+    def _get_flat_coordinate(self, coordinate_obs: dict[str, Any]) -> np.ndarray:
         """Get the flattened obs.
 
         Args:
-            coordinate_obs: dict of coordinate and sensor observations.
+            coordinate_obs (dict[str, Any]): The dict of coordinate and sensor observations.
 
         Returns:
-            flat_obs (np.ndarray): flattened observation.
+            flat_obs: The flattened observation.
         """
+        assert (
+            self.coordinate_observation_space.shape is not None
+        ), 'Bad coordinate_observation_space'
         flat_obs = np.zeros(self.coordinate_observation_space.shape[0])
         for k in self._flatten_order:
             idx = self.key_to_slice[k]
             flat_obs[idx] = coordinate_obs[k].flat
         return flat_obs
 
-    def _get_coordinate_sensor(self) -> dict:
+    def _get_coordinate_sensor(self) -> dict[str, Any]:
         """Return the coordinate observation and sensor observation.
 
         We will ignore the z-axis coordinates in every poses.
         The returned obs coordinates are all in the robot coordinates.
 
         Returns:
-            coordinate_obs (dict): coordinate observation.
+            coordinate_obs: The coordinate observation.
         """
         obs = {}
         robot_matrix = self._env.task.agent.mat
@@ -360,17 +377,17 @@ def _get_coordinate_sensor(self) -> dict:
 
     def _dist_xy(
         self,
-        pos1: np.ndarray | list,
-        pos2: np.ndarray | list,
+        pos1: np.ndarray | list[np.ndarray],
+        pos2: np.ndarray | list[np.ndarray],
     ) -> float:
         """Return the distance from the robot to an XY position.
 
         Args:
-            pos1 (np.ndarray | list): The first position.
-            pos2 (np.ndarray | list): The second position.
+            pos1 (np.ndarray | list[np.ndarray]): The first position.
+            pos2 (np.ndarray | list[np.ndarray]): The second position.
 
         Returns:
-            distance (float): The distance between the two positions.
+            distance: The distance between the two positions.
         """
         pos1 = np.asarray(pos1)
         pos2 = np.asarray(pos2)
@@ -383,7 +400,14 @@ def _dist_xy(
     def step(
         self,
         action: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict]:
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        dict[str, Any],
+    ]:
         """Step the environment.
 
         .. note::
@@ -396,12 +420,12 @@ def step(
             action (torch.Tensor): Action to take.
 
         Returns:
-            observation (torch.Tensor): agent's observation of the current environment.
-            reward (torch.Tensor): amount of reward returned after previous action.
-            cost (torch.Tensor): amount of cost returned after previous action.
-            terminated (torch.Tensor): whether the episode has ended.
-            truncated (torch.Tensor): whether the episode has been truncated due to a time limit.
-            info (Dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning).
+            observation: The agent's observation of the current environment.
+            reward: The amount of reward returned after previous action.
+            cost: The amount of cost returned after previous action.
+            terminated: Whether the episode has ended.
+            truncated: Whether the episode has been truncated due to a time limit.
+            info: Some information logged by the environment.
         """
         obs_original, reward, cost, terminated, truncated, info = self._env.step(
             action.detach().cpu().numpy(),
@@ -440,15 +464,15 @@ def step(
 
         return obs, reward, cost, terminated, truncated, info
 
-    def reset(self, seed: int | None = None) -> tuple[torch.Tensor, dict]:
+    def reset(self, seed: int | None = None) -> tuple[torch.Tensor, dict[str, Any]]:
         """Reset the environment.
 
         Args:
             seed (int, optional): Seed to reset the environment. Defaults to None.
 
         Returns:
-            observation (torch.Tensor): agent's observation of the current environment.
-            info (Dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning).
+            observation: The initial observation of the space.
+            info: Some information logged by the environment.
         """
         obs_original, info = self._env.reset(seed=seed)
         if self._task == 'Goal':
@@ -477,7 +501,7 @@ def sample_action(self) -> torch.Tensor:
         """Sample a random action.
 
         Returns:
-            torch.Tensor: A random action.
+            The sampled action.
         """
         return torch.as_tensor(
             self._env.action_space.sample(),
@@ -489,7 +513,8 @@ def render(self) -> Any:
         """Render the environment.
 
         Returns:
-            Any: Rendered environment.
+            The rendered frames, we recommend using `np.ndarray` which could construct video by
+            moviepy.
         """
         return self._env.render()
 
diff --git a/omnisafe/envs/wrapper.py b/omnisafe/envs/wrapper.py
index d2ae23d8c..558aecad2 100644
--- a/omnisafe/envs/wrapper.py
+++ b/omnisafe/envs/wrapper.py
@@ -504,7 +504,7 @@ def step(
 
 
 class ActionRepeat(Wrapper):
-    """Repeat ab action given times.
+    """Repeat action given times.
 
     Example:
         >>> env = ActionRepeat(env, times=3)
@@ -530,8 +530,27 @@ def __init__(
     def step(
         self,
         action: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict]:
-        """Run self._times timesteps of the environment's dynamics using the agent actions."""
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        dict[str, Any],
+    ]:
+        """Run self._times timesteps of the environment's dynamics using the agent actions.
+
+        Args:
+            action (torch.Tensor): The action from the agent or random.
+
+        Returns:
+            observation: The agent's observation of the current environment.
+            reward: The amount of reward returned after previous action.
+            cost: The amount of cost returned after previous action.
+            terminated: Whether the episode has ended.
+            truncated: Whether the episode has been truncated due to a time limit.
+            info: Some information logged by the environment.
+        """
         rewards, costs = torch.tensor(0.0).to(self._device), torch.tensor(0.0).to(self._device)
         for _step, _ in enumerate(range(self._times)):
             obs, reward, cost, terminated, truncated, info = super().step(action)
diff --git a/omnisafe/utils/config.py b/omnisafe/utils/config.py
index 892e23df5..a2d4d4df4 100644
--- a/omnisafe/utils/config.py
+++ b/omnisafe/utils/config.py
@@ -79,8 +79,8 @@ class Config(dict):
         max_len (int): Maximum length.
         num_threads (int): Number of threads.
 
-    Args:
-        **kwargs (Any): Keyword arguments to set config.
+    Keyword Args:
+        kwargs (Any): keyword arguments to set the attributes.
     """
 
     seed: int
diff --git a/omnisafe/utils/plotter.py b/omnisafe/utils/plotter.py
index e6df09137..5bdbb7ec2 100644
--- a/omnisafe/utils/plotter.py
+++ b/omnisafe/utils/plotter.py
@@ -102,7 +102,9 @@ def plot_data(
             value (str, optional): The y-axis label. Defaults to 'Rewards'.
             condition (str, optional): The condition label. Defaults to 'Condition1'.
             smooth (int, optional): The smoothing window size. Defaults to 1.
-            **kwargs (Any): Additional arguments.
+
+        Keyword Args:
+            kwargs: Other keyword arguments for ``sns.lineplot``.
         """
         if smooth > 1:
             y = np.ones(smooth)


Test patch:
diff --git a/tests/test_ensemble.py b/tests/test_ensemble.py
index e9378331e..8f6df0630 100644
--- a/tests/test_ensemble.py
+++ b/tests/test_ensemble.py
@@ -30,9 +30,7 @@
 
 def test_standard_scaler():
     standard_scaler = StandardScaler(device='cpu')
-    numpy_input = np.random.rand(10, 10)
-    torch_input = torch.from_numpy(numpy_input)
-    assert isinstance(standard_scaler.transform(numpy_input), np.ndarray)
+    torch_input = torch.rand(10, 10)
     assert isinstance(standard_scaler.transform(torch_input), torch.Tensor)